Expose host status cache, and use it for all bulk operations

9 files changed, 53 insertions, 26 deletions

diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/Orchestrator.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/Orchestrator.java
index a639d07e504..59b320cf501 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/Orchestrator.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/Orchestrator.java
@@ -10,7 +10,9 @@ import com.yahoo.vespa.orchestrator.status.ApplicationInstanceStatus;
 import com.yahoo.vespa.orchestrator.status.HostStatus;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 
 /**
  * The orchestrator is used to coordinate the need of vespa services to restart or
@@ -46,6 +48,14 @@ public interface Orchestrator {
      */
     HostStatus getNodeStatus(HostName hostName) throws HostNameNotFoundException;
 
+    /**
+     * Returns a not necessarily consistent mapping from host names to their statuses, for hosts known by this.
+     *
+     * Prefer this to {@link #getNodeStatus(HostName)} when consistency is not required, and when doing bulk reads.
+     * @return a mapping from host names to their statuses. Unknown hosts map to {@code Optional.empty()}.
+     */
+    Function<HostName, Optional<HostStatus>> getNodeStatuses();
+
     void setNodeStatus(HostName hostName, HostStatus state) throws OrchestrationException;
 
     /**
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
index 33cfa310a68..5d5d6df0fb3 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
@@ -35,8 +35,10 @@ import java.time.Clock;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 
@@ -104,6 +106,14 @@ public class OrchestratorImpl implements Orchestrator {
     }
 
     @Override
+    public Function<HostName, Optional<HostStatus>> getNodeStatuses() {
+        Function<ApplicationInstanceReference, Set<HostName>> suspendedHosts = statusService.getSuspendedHostsByApplication();
+        return hostName -> instanceLookupService.findInstanceByHost(hostName)
+                                                .map(application -> suspendedHosts.apply(application.reference()).contains(hostName)
+                                                                    ? HostStatus.ALLOWED_TO_BE_DOWN : HostStatus.NO_REMARKS);
+    }
+
+    @Override
     public void setNodeStatus(HostName hostName, HostStatus status) throws OrchestrationException {
         ApplicationInstanceReference reference = getApplicationInstance(hostName).reference();
         OrchestratorContext context = OrchestratorContext.createContextForSingleAppOp(clock);
@@ -141,7 +151,7 @@ public class OrchestratorImpl implements Orchestrator {
 
             ApplicationInstanceStatus appStatus = statusService.forApplicationInstance(appInstance.reference()).getApplicationInstanceStatus();
             if (appStatus == ApplicationInstanceStatus.NO_REMARKS) {
-                policy.releaseSuspensionGrant(context.createSubcontextWithinLock(), appInstance, hostName, statusRegistry);
+                policy.releaseSuspensionGrant(context.createSubcontextWithinLock(), appInstance, hostName, statusRegistry, statusService);
             }
         }
     }
@@ -164,6 +174,7 @@ public class OrchestratorImpl implements Orchestrator {
             ApplicationApi applicationApi = new ApplicationApiImpl(
                     nodeGroup,
                     statusRegistry,
+                    statusService,
                     clusterControllerClientFactory);
 
             policy.acquirePermissionToRemove(context.createSubcontextWithinLock(), applicationApi);
@@ -188,6 +199,7 @@ public class OrchestratorImpl implements Orchestrator {
 
             ApplicationApi applicationApi = new ApplicationApiImpl(nodeGroup,
                                                                    hostStatusRegistry,
+                                                                   statusService,
                                                                    clusterControllerClientFactory);
             policy.grantSuspensionRequest(context.createSubcontextWithinLock(), applicationApi);
         }
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorUtil.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorUtil.java
index 79506d042e2..de8f7e9f2bc 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorUtil.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorUtil.java
@@ -57,14 +57,6 @@ public class OrchestratorUtil {
                 .collect(toSet());
     }
 
-    public static Map<HostName, HostStatus> getHostStatusMap(Collection<HostName> hosts,
-                                                             ReadOnlyStatusRegistry hostStatusService) {
-        return hosts.stream()
-                .collect(Collectors.toMap(
-                        hostName -> hostName,
-                        hostName -> hostStatusService.getHostStatus(hostName)));
-    }
-
     private static boolean hasServiceInstanceOnHost(ServiceCluster serviceCluster, HostName hostName) {
         return serviceInstancesOnHost(serviceCluster, hostName).count() > 0;
     }
@@ -75,11 +67,6 @@ public class OrchestratorUtil {
                 .filter(instance -> instance.hostName().equals(hostName));
     }
 
-    public static <K, V1, V2> Map<K, V2> mapValues(Map<K, V1> map, Function<V1, V2> valueConverter) {
-        return map.entrySet().stream()
-                .collect(toMap(Map.Entry::getKey, entry -> valueConverter.apply(entry.getValue())));
-    }
-
     private static final Pattern APPLICATION_INSTANCE_REFERENCE_REST_FORMAT_PATTERN = Pattern.compile("^([^:]+):(.+)$");
 
     /** Returns an ApplicationInstanceReference constructed from the serialized format used in the REST API. */
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/ApplicationApiImpl.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/ApplicationApiImpl.java
index 5800b48da75..1d74fb337d9 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/ApplicationApiImpl.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/ApplicationApiImpl.java
@@ -12,6 +12,7 @@ import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactory;
 import com.yahoo.vespa.orchestrator.status.ApplicationInstanceStatus;
 import com.yahoo.vespa.orchestrator.status.HostStatus;
 import com.yahoo.vespa.orchestrator.status.MutableStatusRegistry;
+import com.yahoo.vespa.orchestrator.status.StatusService;
 
 import java.util.Collection;
 import java.util.Comparator;
@@ -35,12 +36,16 @@ public class ApplicationApiImpl implements ApplicationApi {
 
     public ApplicationApiImpl(NodeGroup nodeGroup,
                               MutableStatusRegistry hostStatusService,
+                              StatusService statusService,
                               ClusterControllerClientFactory clusterControllerClientFactory) {
         this.applicationInstance = nodeGroup.getApplication();
         this.nodeGroup = nodeGroup;
         this.hostStatusService = hostStatusService;
         Collection<HostName> hosts = getHostsUsedByApplicationInstance(applicationInstance);
-        this.hostStatusMap = hosts.stream().collect(Collectors.toMap(Function.identity(), hostStatusService::getHostStatus));
+        Collection<HostName> suspendedHosts = statusService.getSuspendedHostsByApplication().apply(nodeGroup.getApplicationReference());
+        this.hostStatusMap = hosts.stream().collect(Collectors.toMap(Function.identity(),
+                                                                     hostName -> suspendedHosts.contains(hostName)
+                                                                                 ? HostStatus.ALLOWED_TO_BE_DOWN : HostStatus.NO_REMARKS));
         this.clusterInOrder = makeClustersInOrder(nodeGroup, hostStatusMap, clusterControllerClientFactory);
     }
 
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java
index b6e7014cac0..e761991ac1a 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java
@@ -14,6 +14,7 @@ import com.yahoo.vespa.orchestrator.model.StorageNode;
 import com.yahoo.vespa.orchestrator.status.ApplicationInstanceStatus;
 import com.yahoo.vespa.orchestrator.status.HostStatus;
 import com.yahoo.vespa.orchestrator.status.MutableStatusRegistry;
+import com.yahoo.vespa.orchestrator.status.StatusService;
 
 import java.util.logging.Logger;
 
@@ -104,9 +105,10 @@ public class HostedVespaPolicy implements Policy {
             OrchestratorContext context,
             ApplicationInstance applicationInstance,
             HostName hostName,
-            MutableStatusRegistry hostStatusService) throws HostStateChangeDeniedException {
+            MutableStatusRegistry hostStatusService,
+            StatusService statusService) throws HostStateChangeDeniedException {
         NodeGroup nodeGroup = new NodeGroup(applicationInstance, hostName);
-        ApplicationApi applicationApi = new ApplicationApiImpl(nodeGroup, hostStatusService, clusterControllerClientFactory);
+        ApplicationApi applicationApi = new ApplicationApiImpl(nodeGroup, hostStatusService, statusService, clusterControllerClientFactory);
         releaseSuspensionGrant(context, applicationApi);
     }
 
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/Policy.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/Policy.java
index e2487301326..7744d75681b 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/Policy.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/Policy.java
@@ -6,6 +6,7 @@ import com.yahoo.vespa.applicationmodel.HostName;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
 import com.yahoo.vespa.orchestrator.model.ApplicationApi;
 import com.yahoo.vespa.orchestrator.status.MutableStatusRegistry;
+import com.yahoo.vespa.orchestrator.status.StatusService;
 
 /**
  * @author oyving
@@ -32,6 +33,7 @@ public interface Policy {
     void releaseSuspensionGrant(
             OrchestratorContext context, ApplicationInstance applicationInstance,
             HostName hostName,
-            MutableStatusRegistry hostStatusService) throws HostStateChangeDeniedException;
+            MutableStatusRegistry hostStatusService,
+            StatusService statusService) throws HostStateChangeDeniedException;
 
 }
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/resources/InstanceResource.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/resources/InstanceResource.java
index cd20a01f6af..65ef1c9eaf0 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/resources/InstanceResource.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/resources/InstanceResource.java
@@ -34,7 +34,6 @@ import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
-import static com.yahoo.vespa.orchestrator.OrchestratorUtil.getHostStatusMap;
 import static com.yahoo.vespa.orchestrator.OrchestratorUtil.getHostsUsedByApplicationInstance;
 import static com.yahoo.vespa.orchestrator.OrchestratorUtil.parseAppInstanceReference;
 
@@ -82,11 +81,13 @@ public class InstanceResource {
                 = instanceLookupService.findInstanceById(instanceId)
                 .orElseThrow(() -> new WebApplicationException(Response.status(Response.Status.NOT_FOUND).build()));
 
-        Set<HostName> hostsUsedByApplicationInstance = getHostsUsedByApplicationInstance(applicationInstance);
-        Map<HostName, HostStatus> hostStatusMap = getHostStatusMap(hostsUsedByApplicationInstance,
-                                                                   statusService.forApplicationInstance(instanceId));
-        Map<HostName, String> hostStatusStringMap = OrchestratorUtil.mapValues(hostStatusMap, HostStatus::name);
-        return InstanceStatusResponse.create(applicationInstance, hostStatusStringMap);
+        Set<HostName> suspendedHosts = statusService.getSuspendedHostsByApplication().apply(applicationInstance.reference());
+        Map<HostName, String> hostStatusMap = getHostsUsedByApplicationInstance(applicationInstance)
+                .stream()
+                .collect(Collectors.toMap(hostName -> hostName,
+                                          hostName -> suspendedHosts.contains(hostName) ? HostStatus.ALLOWED_TO_BE_DOWN.name()
+                                                                                        : HostStatus.NO_REMARKS.name()));
+        return InstanceStatusResponse.create(applicationInstance, hostStatusMap);
     }
 
     @GET
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/StatusService.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/StatusService.java
index c9e846906e3..37f814ef03a 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/StatusService.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/StatusService.java
@@ -68,8 +68,8 @@ public interface StatusService {
     /**
      * Returns a not necessarily consistent mapping from applications to their set of suspended hosts.
      *
-     * If the lock for an application is held when this is acquired, the view of that application's hosts
-     * is consistent and up to date for as long as the lock is held.
+     * If the lock for an application is held when this mapping is acquired, new sets returned for that application
+     * are consistent and up to date for as long as the lock is held. (The sets themselves don't reflect changes.)
      */
     Function<ApplicationInstanceReference, Set<HostName>> getSuspendedHostsByApplication();
 
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java
index 67b14f29764..0ed0d6ad31a 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java
@@ -21,6 +21,7 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 
@@ -91,6 +92,13 @@ public class ZookeeperStatusService implements StatusService {
         }
     }
 
+    @Override
+    public Function<ApplicationInstanceReference, Set<HostName>> getSuspendedHostsByApplication() {
+        Map<ApplicationInstanceReference, Set<HostName>> suspendedHostsByApplication = getValidCache();
+        return application -> suspendedHostsByApplication.computeIfAbsent(application, this::hostsDownFor);
+    }
+
+
     /**
      *  1) locks the status service for an application instance.
      *  2) fails all operations in this thread when the session is lost,