Support per-cloud OS upgrades

1 files changed, 82 insertions, 23 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java
index 518394b46fc..8ad21fee2f1 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java
@@ -1,21 +1,27 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.versions;
 
-import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.yahoo.component.Version;
+import com.yahoo.config.provision.Environment;
+import com.yahoo.config.provision.HostName;
+import com.yahoo.config.provision.RegionName;
 import com.yahoo.vespa.hosted.controller.Controller;
+import com.yahoo.vespa.hosted.controller.api.integration.zone.CloudName;
 import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.application.SystemApplication;
 import com.yahoo.vespa.hosted.controller.maintenance.OsUpgrader;
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
+import static java.util.stream.Collectors.collectingAndThen;
+
 /**
  * Information about OS versions in this system.
  *
@@ -23,36 +29,44 @@ import java.util.stream.Collectors;
  */
 public class OsVersionStatus {
 
-    public static final OsVersionStatus empty = new OsVersionStatus(Collections.emptyList());
+    public static final OsVersionStatus empty = new OsVersionStatus(Collections.emptyMap());
 
-    private final List<OsVersion> versions;
+    private final Map<OsVersion, List<Node>> versions;
 
-    public OsVersionStatus(List<OsVersion> versions) {
-        this.versions = ImmutableList.copyOf(versions);
+    /** Public for serialization purpose only. Use {@link OsVersionStatus#compute(Controller)} for an up-to-date status */
+    public OsVersionStatus(Map<OsVersion, List<Node>> versions) {
+        this.versions = ImmutableMap.copyOf(Objects.requireNonNull(versions, "versions must be non-null"));
     }
 
-    /** All known OS versions */
-    public List<OsVersion> versions() {
+    /** All known OS versions and their nodes */
+    public Map<OsVersion, List<Node>> versions() {
         return versions;
     }
 
-    /**
-     * Compute the current OS version status in this status. This is expensive as all config servers in the system
-     * must be queried.
-     */
+    /** Returns node versions that exist in given cloud */
+    public List<Node> nodeVersionsIn(CloudName cloud) {
+        return versions.entrySet().stream()
+                       .filter(entry -> entry.getKey().cloud().equals(cloud))
+                       .flatMap(entry -> entry.getValue().stream())
+                       .collect(collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
+    }
+
+    /** Compute the current OS versions in this system. This is expensive and should be called infrequently */
     public static OsVersionStatus compute(Controller controller) {
-        Map<Version, List<OsVersion.Node>> versions = new HashMap<>();
-        // Always include target version, if set
-        controller.osVersion().ifPresent(version -> versions.put(version, new ArrayList<>()));
+        Map<OsVersion, List<Node>> versions = new HashMap<>();
+
+        // Always include all target versions
+        controller.osVersions().forEach(osVersion -> versions.put(osVersion, new ArrayList<>()));
+
         for (SystemApplication application : SystemApplication.all()) {
             if (application.nodeTypesWithUpgradableOs().isEmpty()) {
-                continue; // Avoid querying applications that do not have nodes with upgradable OS
+                continue; // Avoid querying applications that do not contain nodes with upgradable OS
             }
             for (ZoneId zone : controller.zoneRegistry().zones().controllerUpgraded().ids()) {
                 controller.configServer().nodeRepository().list(zone, application.id()).stream()
                           .filter(node -> OsUpgrader.eligibleForUpgrade(node, application))
-                          .map(node -> new OsVersion.Node(node.hostname(), node.currentOsVersion(), zone.environment(), zone.region()))
-                          .forEach(node -> versions.compute(node.version(), (ignored, nodes) -> {
+                          .map(node -> new Node(node.hostname(), node.currentOsVersion(), zone.environment(), zone.region()))
+                          .forEach(node -> versions.compute(new OsVersion(node.version(), zone.cloud()), (ignored, nodes) -> {
                               if (nodes == null) {
                                   nodes = new ArrayList<>();
                               }
@@ -61,11 +75,56 @@ public class OsVersionStatus {
                           }));
             }
         }
-        return new OsVersionStatus(versions.entrySet()
-                                           .stream()
-                                           .map(kv -> new OsVersion(kv.getKey(), kv.getValue()))
-                                           .sorted(Comparator.comparing(OsVersion::version))
-                                           .collect(Collectors.toList()));
+
+        return new OsVersionStatus(versions);
+    }
+
+    /** A node in this system and its current OS version */
+    public static class Node {
+
+        private final HostName hostname;
+        private final Version version;
+        private final Environment environment;
+        private final RegionName region;
+
+        public Node(HostName hostname, Version version, Environment environment, RegionName region) {
+            this.hostname = Objects.requireNonNull(hostname, "hostname must be non-null");
+            this.version = Objects.requireNonNull(version, "version must be non-null");
+            this.environment = Objects.requireNonNull(environment, "environment must be non-null");
+            this.region = Objects.requireNonNull(region, "region must be non-null");
+        }
+
+        public HostName hostname() {
+            return hostname;
+        }
+
+        public Version version() {
+            return version;
+        }
+
+        public Environment environment() {
+            return environment;
+        }
+
+        public RegionName region() {
+            return region;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            Node node = (Node) o;
+            return Objects.equals(hostname, node.hostname) &&
+                   Objects.equals(version, node.version) &&
+                   environment == node.environment &&
+                   Objects.equals(region, node.region);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(hostname, version, environment, region);
+        }
     }
 
 }