Clean up streams in SetNumInstancesCommand

Also using Service instead of serviceId String - because it helps readability.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=228952033
This commit is contained in:
guyben 2019-01-11 15:03:15 -08:00 committed by Ben McIlwain
parent 773af1da35
commit fd8a18b72e
2 changed files with 68 additions and 73 deletions

View file

@ -15,28 +15,27 @@
package google.registry.tools; package google.registry.tools;
import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.collect.ImmutableSetMultimap.flatteningToImmutableSetMultimap;
import static google.registry.util.CollectionUtils.nullToEmpty;
import static google.registry.util.PreconditionsUtils.checkArgumentPresent;
import com.beust.jcommander.Parameter; import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters; import com.beust.jcommander.Parameters;
import com.google.api.services.appengine.v1.Appengine; import com.google.api.services.appengine.v1.Appengine;
import com.google.common.base.Enums;
import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap; import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.MultimapBuilder;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Sets;
import com.google.common.flogger.FluentLogger; import com.google.common.flogger.FluentLogger;
import google.registry.config.RegistryConfig.Config; import google.registry.config.RegistryConfig.Config;
import google.registry.tools.AppEngineConnection.Service; import google.registry.tools.AppEngineConnection.Service;
import google.registry.util.AppEngineServiceUtils; import google.registry.util.AppEngineServiceUtils;
import java.io.IOException; import java.io.IOException;
import java.util.AbstractMap.SimpleEntry;
import java.util.Arrays;
import java.util.Collection;
import java.util.List; import java.util.List;
import java.util.Set; import java.util.Set;
import java.util.stream.Stream;
import javax.inject.Inject; import javax.inject.Inject;
/** A command to set the number of instances for an App Engine service. */ /** A command to set the number of instances for an App Engine service. */
@ -49,11 +48,12 @@ final class SetNumInstancesCommand implements CommandWithRemoteApi {
private static final FluentLogger logger = FluentLogger.forEnclosingClass(); private static final FluentLogger logger = FluentLogger.forEnclosingClass();
private static final ImmutableSet<String> ALL_VALID_SERVICES = private static final ImmutableSet<Service> ALL_DEPLOYED_SERVICES =
Arrays.stream(Service.values()).map(Service::name).collect(toImmutableSet()); ImmutableSet.copyOf(Service.values());
private static final ImmutableSet<String> ALL_DEPLOYED_SERVICE_IDS = private static final ImmutableMap<String, Service> SERVICE_ID_TO_SERVICE =
Arrays.stream(Service.values()).map(Service::getServiceId).collect(toImmutableSet()); ALL_DEPLOYED_SERVICES.stream()
.collect(toImmutableMap(service -> service.getServiceId(), service -> service));
// TODO(b/119629679): Use List<Service> after upgrading jcommander to latest version. // TODO(b/119629679): Use List<Service> after upgrading jcommander to latest version.
@Parameter( @Parameter(
@ -61,7 +61,7 @@ final class SetNumInstancesCommand implements CommandWithRemoteApi {
description = description =
"Comma-delimited list of App Engine services to set. " "Comma-delimited list of App Engine services to set. "
+ "Allowed values: [DEFAULT, TOOLS, BACKEND, PUBAPI]") + "Allowed values: [DEFAULT, TOOLS, BACKEND, PUBAPI]")
private List<String> services = ImmutableList.of(); private List<String> serviceNames = ImmutableList.of();
@Parameter( @Parameter(
names = "--versions", names = "--versions",
@ -93,103 +93,98 @@ final class SetNumInstancesCommand implements CommandWithRemoteApi {
@Override @Override
public void run() throws Exception { public void run() throws Exception {
Set<String> invalidServiceIds =
Sets.difference(ImmutableSet.copyOf(services), ALL_VALID_SERVICES);
checkArgument(invalidServiceIds.isEmpty(), "Invalid service(s): %s", invalidServiceIds);
Set<String> serviceIds = ImmutableSet<Service> services =
services.stream() serviceNames.stream()
.map(service -> Service.valueOf(service).getServiceId()) .map(
name ->
checkArgumentPresent(
Enums.getIfPresent(Service.class, name).toJavaUtil(),
"Invalid service '%s'. Allowed values are %s",
name,
ALL_DEPLOYED_SERVICES))
.collect(toImmutableSet()); .collect(toImmutableSet());
if (nonLiveVersions) { if (nonLiveVersions) {
checkArgument(versions.isEmpty(), "--versions cannot be set if --non_live_versions is set"); checkArgument(versions.isEmpty(), "--versions cannot be set if --non_live_versions is set");
serviceIds = serviceIds.isEmpty() ? ALL_DEPLOYED_SERVICE_IDS : serviceIds; services = services.isEmpty() ? ALL_DEPLOYED_SERVICES : services;
Multimap<String, String> allLiveVersionsMap = getAllLiveVersionsMap(serviceIds); ImmutableSetMultimap<Service, String> allLiveVersionsMap = getAllLiveVersionsMap(services);
Multimap<String, String> manualScalingVersionsMap = getManualScalingVersionsMap(serviceIds); ImmutableSetMultimap<Service, String> manualScalingVersionsMap =
getManualScalingVersionsMap(services);
// Set number of instances for versions which are manual scaling and non-live // Set number of instances for versions which are manual scaling and non-live
manualScalingVersionsMap.forEach( manualScalingVersionsMap.forEach(
(serviceId, versionId) -> { (service, versionId) -> {
if (!allLiveVersionsMap.containsEntry(serviceId, versionId)) { if (!allLiveVersionsMap.containsEntry(service, versionId)) {
setNumInstances(serviceId, versionId, numInstances); setNumInstances(service, versionId, numInstances);
} }
}); });
} else { } else {
checkArgument(!serviceIds.isEmpty(), "Service must be specified"); checkArgument(!services.isEmpty(), "Service must be specified");
checkArgument(!versions.isEmpty(), "Version must be specified"); checkArgument(!versions.isEmpty(), "Version must be specified");
checkArgument(numInstances > 0, "Number of instances must be greater than zero"); checkArgument(numInstances > 0, "Number of instances must be greater than zero");
Multimap<String, String> manualScalingVersionsMap = getManualScalingVersionsMap(serviceIds); ImmutableSetMultimap<Service, String> manualScalingVersionsMap =
getManualScalingVersionsMap(services);
for (String serviceId : serviceIds) { for (Service service : services) {
for (String versionId : versions) { for (String versionId : versions) {
checkArgument( checkArgument(
manualScalingVersionsMap.containsEntry(serviceId, versionId), manualScalingVersionsMap.containsEntry(service, versionId),
"Version %s of service %s is not managed through manual scaling", "Version %s of service %s is not managed through manual scaling",
versionId, versionId,
serviceId); service);
setNumInstances(serviceId, versionId, numInstances); setNumInstances(service, versionId, numInstances);
} }
} }
} }
} }
private void setNumInstances(String service, String version, long numInstances) { private void setNumInstances(Service service, String version, long numInstances) {
appEngineServiceUtils.setNumInstances(service, version, numInstances); appEngineServiceUtils.setNumInstances(service.getServiceId(), version, numInstances);
logger.atInfo().log( logger.atInfo().log(
"Successfully set version %s of service %s to %d instances.", "Successfully set version %s of service %s to %d instances.",
version, service, numInstances); version, service, numInstances);
} }
private Multimap<String, String> getAllLiveVersionsMap(Set<String> services) { private ImmutableSetMultimap<Service, String> getAllLiveVersionsMap(Set<Service> services) {
try { try {
return Stream.of(appengine.apps().services().list(projectId).execute().getServices()) return nullToEmpty(appengine.apps().services().list(projectId).execute().getServices())
.flatMap(Collection::stream) .stream()
.filter(service -> services.contains(service.getId())) .filter(
.flatMap(
service -> service ->
// getAllocations returns only live versions or null services.contains(SERVICE_ID_TO_SERVICE.getOrDefault(service.getId(), null)))
Stream.of(service.getSplit().getAllocations())
.flatMap(
allocations ->
allocations.keySet().stream()
.map(versionId -> new SimpleEntry<>(service.getId(), versionId))))
.collect( .collect(
Multimaps.toMultimap( flatteningToImmutableSetMultimap(
SimpleEntry::getKey, service -> SERVICE_ID_TO_SERVICE.get(service.getId()),
SimpleEntry::getValue, service -> nullToEmpty(service.getSplit().getAllocations()).keySet().stream()));
MultimapBuilder.treeKeys().arrayListValues()::build));
} catch (IOException e) { } catch (IOException e) {
throw new RuntimeException(e); throw new RuntimeException(e);
} }
} }
private Multimap<String, String> getManualScalingVersionsMap(Set<String> services) { private ImmutableSetMultimap<Service, String> getManualScalingVersionsMap(Set<Service> services) {
return services.stream() return services.stream()
.flatMap( .collect(
serviceId -> { flatteningToImmutableSetMultimap(
service -> service,
service -> {
try { try {
return Stream.of( return nullToEmpty(
appengine appengine
.apps() .apps()
.services() .services()
.versions() .versions()
.list(projectId, serviceId) .list(projectId, service.getServiceId())
.execute() .execute()
.getVersions()) .getVersions())
.flatMap(Collection::stream) .stream()
.filter(version -> version.getManualScaling() != null) .filter(version -> version.getManualScaling() != null)
.map(version -> new SimpleEntry<>(serviceId, version.getId())); .map(version -> version.getId());
} catch (IOException e) { } catch (IOException e) {
throw new RuntimeException(e); throw new RuntimeException(e);
} }
}) }));
.collect(
Multimaps.toMultimap(
SimpleEntry::getKey,
SimpleEntry::getValue,
MultimapBuilder.treeKeys().arrayListValues()::build));
} }
} }

View file

@ -60,7 +60,7 @@ public class SetNumInstancesCommandTest extends CommandTestCase<SetNumInstancesC
assertThrows( assertThrows(
IllegalArgumentException.class, IllegalArgumentException.class,
() -> runCommand("--services=", "--versions=version", "--num_instances=5")); () -> runCommand("--services=", "--versions=version", "--num_instances=5"));
assertThat(thrown).hasMessageThat().contains("Invalid service(s): []"); assertThat(thrown).hasMessageThat().contains("Invalid service ''");
} }
@Test @Test
@ -71,7 +71,7 @@ public class SetNumInstancesCommandTest extends CommandTestCase<SetNumInstancesC
() -> () ->
runCommand( runCommand(
"--services=INVALID,DEFAULT", "--versions=version", "--num_instances=5")); "--services=INVALID,DEFAULT", "--versions=version", "--num_instances=5"));
assertThat(thrown).hasMessageThat().contains("Invalid service(s): [INVALID]"); assertThat(thrown).hasMessageThat().contains("Invalid service 'INVALID'");
} }
@Test @Test