Prune certified OS versions by cloud

2 files changed, 40 insertions, 26 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/OsController.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/OsController.java
index 58c3b4da5e4..1ca12cac957 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/OsController.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/OsController.java
@@ -18,6 +18,7 @@ import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.TreeSet;
+import java.util.function.BinaryOperator;
 import java.util.function.Function;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -162,12 +163,19 @@ public record OsController(Controller controller) {
     /** Remove certifications for non-existent OS versions */
     public void removeStaleCertifications(OsVersionStatus currentStatus) {
         try (Mutex lock = curator().lockCertifiedOsVersions()) {
-            Optional<OsVersion> minKnownVersion = currentStatus.versions().keySet().stream()
-                                                               .filter(v -> !v.version().isEmpty())
-                                                               .min(Comparator.naturalOrder());
-            if (minKnownVersion.isEmpty()) return;
+            Map<CloudName, Version> oldestVersionByCloud = currentStatus.versions().keySet().stream()
+                                                                        .filter(v -> !v.version().isEmpty())
+                                                                        .collect(Collectors.toMap(OsVersion::cloud,
+                                                                                                  OsVersion::version,
+                                                                                                  BinaryOperator.minBy(Comparator.naturalOrder())));
+            if (oldestVersionByCloud.isEmpty()) return;
+
             Set<CertifiedOsVersion> certifiedVersions = new HashSet<>(readCertified());
-            if (certifiedVersions.removeIf(cv -> cv.osVersion().version().isBefore(minKnownVersion.get().version()))) {
+            boolean modified = certifiedVersions.removeIf(certifiedVersion -> {
+                Version oldestVersion = oldestVersionByCloud.get(certifiedVersion.osVersion().cloud());
+                return oldestVersion == null || certifiedVersion.osVersion().version().isBefore(oldestVersion);
+            });
+            if (modified) {
                 curator().writeCertifiedOsVersions(certifiedVersions);
             }
         }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
index 6f4052bf0ef..8f4db5cad47 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
@@ -12,10 +12,10 @@ import com.yahoo.vespa.hosted.controller.versions.OsVersionStatus;
 import org.junit.jupiter.api.Test;
 
 import java.time.Duration;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -43,47 +43,53 @@ public class OsVersionStatusUpdaterTest {
 
         // Setting a new target adds it to current status
         Version version1 = Version.fromString("7.1");
-        CloudName cloud = CloudName.DEFAULT;
-        tester.controller().os().upgradeTo(version1, cloud, false, false);
+        CloudName cloud0 = CloudName.DEFAULT;
+        tester.controller().os().upgradeTo(version1, cloud0, false, false);
         statusUpdater.maintain();
 
         var osVersions = tester.controller().os().status().versions();
         assertEquals(3, osVersions.size());
-        assertFalse(osVersions.get(new OsVersion(Version.emptyVersion, cloud)).isEmpty(), "All nodes on unknown version");
-        assertTrue(osVersions.get(new OsVersion(version1, cloud)).isEmpty(), "No nodes on current target");
+        assertFalse(osVersions.get(new OsVersion(Version.emptyVersion, cloud0)).isEmpty(), "All nodes on unknown version");
+        assertTrue(osVersions.get(new OsVersion(version1, cloud0)).isEmpty(), "No nodes on current target");
 
-        CloudName otherCloud = CloudName.AWS;
-        tester.controller().os().upgradeTo(version1, otherCloud, false, false);
+        CloudName cloud1 = CloudName.AWS;
+        Version version2 = Version.fromString("7.0");
+        tester.controller().os().upgradeTo(version2, cloud1, false, false);
         statusUpdater.maintain();
 
         osVersions = tester.controller().os().status().versions();
         assertEquals(4, osVersions.size()); // 2 in cloud, 2 in otherCloud.
-        assertFalse(osVersions.get(new OsVersion(Version.emptyVersion, cloud)).isEmpty(), "All nodes on unknown version");
-        assertTrue(osVersions.get(new OsVersion(version1, cloud)).isEmpty(), "No nodes on current target");
-        assertFalse(osVersions.get(new OsVersion(Version.emptyVersion, otherCloud)).isEmpty(), "All nodes on unknown version");
-        assertTrue(osVersions.get(new OsVersion(version1, otherCloud)).isEmpty(), "No nodes on current target");
+        assertFalse(osVersions.get(new OsVersion(Version.emptyVersion, cloud0)).isEmpty(), "All nodes on unknown version");
+        assertTrue(osVersions.get(new OsVersion(version1, cloud0)).isEmpty(), "No nodes on current target");
+        assertFalse(osVersions.get(new OsVersion(Version.emptyVersion, cloud1)).isEmpty(), "All nodes on unknown version");
+        assertTrue(osVersions.get(new OsVersion(version2, cloud1)).isEmpty(), "No nodes on current target");
 
         // Updating status cleans up stale certifications
         Set<OsVersion> knownVersions = osVersions.keySet().stream()
                                                  .filter(osVersion -> !osVersion.version().isEmpty())
                                                  .collect(Collectors.toSet());
-        List<OsVersion> versionsToCertify = new ArrayList<>(knownVersions);
-        OsVersion futureVersion = new OsVersion(Version.fromString("98.0.2"), cloud); // Keep future version
-        versionsToCertify.addAll(List.of(new OsVersion(Version.fromString("3.11"), cloud),
-                                         futureVersion));
-        for (OsVersion version : versionsToCertify) {
+        // Known versions
+        for (OsVersion version : knownVersions) {
+            tester.controller().os().certify(version.version(), version.cloud(), Version.fromString("1.2.3"));
+        }
+        // Additional versions
+        OsVersion staleVersion0 = new OsVersion(Version.fromString("7.0"), cloud0);        // Only removed for this cloud
+        OsVersion staleVersion1 = new OsVersion(Version.fromString("3.11"), cloud0);       // Stale in both clouds
+        OsVersion staleVersion2 = new OsVersion(Version.fromString("3.11"), cloud1);
+        OsVersion futureVersion = new OsVersion(Version.fromString("98.0.2"), cloud0);     // Keep future version
+        for (OsVersion version : List.of(staleVersion0, staleVersion1, staleVersion2, futureVersion)) {
             tester.controller().os().certify(version.version(), version.cloud(), Version.fromString("1.2.3"));
         }
-        knownVersions.add(futureVersion);
-        assertEquals(knownVersions.size() + 1, certifiedOsVersions(tester).size());
         statusUpdater.maintain();
-        assertEquals(knownVersions, certifiedOsVersions(tester));
+        assertEquals(Stream.concat(knownVersions.stream(), Stream.of(futureVersion)).sorted().toList(),
+                     certifiedOsVersions(tester));
     }
 
-    private static Set<OsVersion> certifiedOsVersions(ControllerTester tester) {
+    private static List<OsVersion> certifiedOsVersions(ControllerTester tester) {
         return tester.controller().os().readCertified().stream()
                      .map(CertifiedOsVersion::osVersion)
-                     .collect(Collectors.toSet());
+                     .sorted()
+                     .toList();
     }
 
 }