Replace image atomically

7 files changed, 95 insertions, 38 deletions

diff --git a/config-provisioning/src/main/java/com/yahoo/config/provision/DockerImage.java b/config-provisioning/src/main/java/com/yahoo/config/provision/DockerImage.java
index 5ea2286c316..cc1ae60ba0a 100644
--- a/config-provisioning/src/main/java/com/yahoo/config/provision/DockerImage.java
+++ b/config-provisioning/src/main/java/com/yahoo/config/provision/DockerImage.java
@@ -14,16 +14,18 @@ import java.util.Optional;
 // TODO: Rename to ContainerImage. Compatibility with older config-models must be preserved.
 public class DockerImage {
 
-    public static final DockerImage EMPTY = new DockerImage("", "", Optional.empty());
+    public static final DockerImage EMPTY = new DockerImage("", "", Optional.empty(), Optional.empty());
 
     private final String registry;
     private final String repository;
     private final Optional<String> tag;
+    private final Optional<DockerImage> replacedBy;
 
-    DockerImage(String registry, String repository, Optional<String> tag) {
+    DockerImage(String registry, String repository, Optional<String> tag, Optional<DockerImage> replacedBy) {
         this.registry = Objects.requireNonNull(registry, "registry must be non-null");
         this.repository = Objects.requireNonNull(repository, "repository must be non-null");
         this.tag = Objects.requireNonNull(tag, "tag must be non-null");
+        this.replacedBy = Objects.requireNonNull(replacedBy);
     }
 
     /** Returns the registry-part of this, i.e. the host/port of the registry. */
@@ -38,7 +40,7 @@ public class DockerImage {
 
     /** Returns the registry and repository for this image, excluding its tag */
     public String untagged() {
-        return new DockerImage(registry, repository, Optional.empty()).asString();
+        return new DockerImage(registry, repository, Optional.empty(), replacedBy).asString();
     }
 
     /** Returns this image's tag, if any */
@@ -51,14 +53,24 @@ public class DockerImage {
         return tag.map(Version::new).orElse(Version.emptyVersion);
     }
 
+    /** The image that replaces this, if any */
+    public Optional<DockerImage> replacedBy() {
+        return replacedBy;
+    }
+
     /** Returns a copy of this tagged with the given version */
     public DockerImage withTag(Version version) {
-        return new DockerImage(registry, repository, Optional.of(version.toFullString()));
+        return new DockerImage(registry, repository, Optional.of(version.toFullString()), replacedBy);
     }
 
     /** Returns a copy of this with registry set to given value */
     public DockerImage withRegistry(String registry) {
-        return new DockerImage(registry, repository, tag);
+        return new DockerImage(registry, repository, tag, replacedBy);
+    }
+
+    /** Returns a copy of this with replacement image set to given value */
+    public DockerImage withReplacedBy(DockerImage image) {
+        return new DockerImage(registry, repository, tag, Optional.of(image).filter(i -> !i.equals(EMPTY)));
     }
 
     public String asString() {
@@ -78,16 +90,17 @@ public class DockerImage {
         DockerImage that = (DockerImage) o;
         return registry.equals(that.registry) &&
                repository.equals(that.repository) &&
-               tag.equals(that.tag);
+               tag.equals(that.tag) &&
+               replacedBy.equals(that.replacedBy);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(registry, repository, tag);
+        return Objects.hash(registry, repository, tag, replacedBy);
     }
 
     public static DockerImage from(String registry, String repository) {
-        return new DockerImage(registry, repository, Optional.empty());
+        return new DockerImage(registry, repository, Optional.empty(), Optional.empty());
     }
 
     public static DockerImage fromString(String s) {
@@ -101,11 +114,11 @@ public class DockerImage {
         if (repository.isEmpty()) throw new IllegalArgumentException("Repository must be non-empty in '" + s + "'");
 
         int tagStart = repository.indexOf(':');
-        if (tagStart < 0) return new DockerImage(registry, repository, Optional.empty());
+        if (tagStart < 0) return new DockerImage(registry, repository, Optional.empty(), Optional.empty());
 
         String tag = repository.substring(tagStart + 1);
         repository = repository.substring(0, tagStart);
-        return new DockerImage(registry, repository, Optional.of(tag));
+        return new DockerImage(registry, repository, Optional.of(tag), Optional.empty());
     }
 
 }
diff --git a/config-provisioning/src/main/resources/configdefinitions/config.provisioning.node-repository.def b/config-provisioning/src/main/resources/configdefinitions/config.provisioning.node-repository.def
index 6409bbd1966..091f518a683 100644
--- a/config-provisioning/src/main/resources/configdefinitions/config.provisioning.node-repository.def
+++ b/config-provisioning/src/main/resources/configdefinitions/config.provisioning.node-repository.def
@@ -4,6 +4,10 @@ namespace=config.provisioning
 # Default container image to use for nodes.
 containerImage string default="registry.example.com:9999/myorg/vespa"
 
+# A replacement container image to use for nodes. This is used to avoid nodes seeing different wanted images in the
+# transition to a new default image or registry.
+containerImageReplacement string default=""
+
 # Whether to cache data read from ZooKeeper in-memory.
 useCuratorClientCache bool default=false
 
diff --git a/config-provisioning/src/test/java/com/yahoo/config/provision/DockerImageTest.java b/config-provisioning/src/test/java/com/yahoo/config/provision/DockerImageTest.java
index ca41f4628cc..f05ebb6e517 100644
--- a/config-provisioning/src/test/java/com/yahoo/config/provision/DockerImageTest.java
+++ b/config-provisioning/src/test/java/com/yahoo/config/provision/DockerImageTest.java
@@ -8,6 +8,7 @@ import java.util.Map;
 import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 /**
@@ -19,10 +20,10 @@ public class DockerImageTest {
     public void parse() {
         Map<String, DockerImage> tests = Map.of(
                 "", DockerImage.EMPTY,
-                "registry.example.com:9999/vespa/vespa:7.42", new DockerImage("registry.example.com:9999", "vespa/vespa", Optional.of("7.42")),
-                "registry.example.com/vespa/vespa:7.42", new DockerImage("registry.example.com", "vespa/vespa", Optional.of("7.42")),
-                "registry.example.com:9999/vespa/vespa", new DockerImage("registry.example.com:9999", "vespa/vespa", Optional.empty()),
-                "registry.example.com/vespa/vespa", new DockerImage("registry.example.com", "vespa/vespa", Optional.empty())
+                "registry.example.com:9999/vespa/vespa:7.42", new DockerImage("registry.example.com:9999", "vespa/vespa", Optional.of("7.42"), Optional.empty()),
+                "registry.example.com/vespa/vespa:7.42", new DockerImage("registry.example.com", "vespa/vespa", Optional.of("7.42"), Optional.empty()),
+                "registry.example.com:9999/vespa/vespa", new DockerImage("registry.example.com:9999", "vespa/vespa", Optional.empty(), Optional.empty()),
+                "registry.example.com/vespa/vespa", new DockerImage("registry.example.com", "vespa/vespa", Optional.empty(), Optional.empty())
         );
         tests.forEach((value, expected) -> {
             DockerImage parsed = DockerImage.fromString(value);
@@ -52,4 +53,10 @@ public class DockerImageTest {
         }
     }
 
+    @Test
+    public void empty_replacement() {
+        DockerImage image = new DockerImage("foo.example.com", "vespa/vespa", Optional.of("7.42"), Optional.empty());
+        assertTrue(image.withReplacedBy(DockerImage.EMPTY).replacedBy().isEmpty());
+    }
+
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index d8ea19b7677..845d2763501 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -127,7 +127,8 @@ public class NodeRepository extends AbstractComponent {
              Clock.systemUTC(),
              zone,
              new DnsNameResolver(),
-             DockerImage.fromString(config.containerImage()),
+             DockerImage.fromString(config.containerImage())
+                        .withReplacedBy(DockerImage.fromString(config.containerImageReplacement())),
              flagSource,
              config.useCuratorClientCache(),
              zone.environment().isProduction() && !zone.getCloud().dynamicProvisioning() ? 1 : 0,
@@ -144,7 +145,7 @@ public class NodeRepository extends AbstractComponent {
                           Clock clock,
                           Zone zone,
                           NameResolver nameResolver,
-                          DockerImage dockerImage,
+                          DockerImage containerImage,
                           FlagSource flagSource,
                           boolean useCuratorClientCache,
                           int spareCount,
@@ -164,7 +165,7 @@ public class NodeRepository extends AbstractComponent {
         this.osVersions = new OsVersions(this);
         this.infrastructureVersions = new InfrastructureVersions(db);
         this.firmwareChecks = new FirmwareChecks(db, clock);
-        this.containerImages = new ContainerImages(db, dockerImage);
+        this.containerImages = new ContainerImages(db, containerImage, flagSource);
         this.jobControl = new JobControl(new JobControlFlags(db, flagSource));
         this.applications = new Applications(db);
         this.spareCount = spareCount;
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
index 92518239258..45156c57481 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
@@ -6,6 +6,9 @@ import com.google.common.base.Suppliers;
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.curator.Lock;
+import com.yahoo.vespa.flags.BooleanFlag;
+import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
 
 import java.time.Duration;
@@ -23,12 +26,12 @@ import java.util.logging.Logger;
  */
 public class ContainerImages {
 
-    private static final Duration defaultCacheTtl = Duration.ofMinutes(1);
+    private static final Duration cacheTtl = Duration.ofMinutes(1);
     private static final Logger log = Logger.getLogger(ContainerImages.class.getName());
 
     private final CuratorDatabaseClient db;
     private final DockerImage defaultImage;
-    private final Duration cacheTtl;
+    private final BooleanFlag replaceImage;
 
     /**
      * The container image is read on every request to /nodes/v2/node/[fqdn]. Cache current images to avoid
@@ -37,14 +40,10 @@ public class ContainerImages {
      */
     private volatile Supplier<Map<NodeType, DockerImage>> images;
 
-    public ContainerImages(CuratorDatabaseClient db, DockerImage defaultImage) {
-        this(db, defaultImage, defaultCacheTtl);
-    }
-
-    ContainerImages(CuratorDatabaseClient db, DockerImage defaultImage, Duration cacheTtl) {
+    public ContainerImages(CuratorDatabaseClient db, DockerImage defaultImage, FlagSource flagSource) {
         this.db = db;
         this.defaultImage = defaultImage;
-        this.cacheTtl = cacheTtl;
+        this.replaceImage = Flags.REGIONAL_CONTAINER_REGISTRY.bindTo(flagSource);
         createCache();
     }
 
@@ -58,15 +57,14 @@ public class ContainerImages {
         return images.get();
     }
 
-    /** Returns the current docker image for given node type, or the type for corresponding child nodes
-     * if it is a Docker host, or default */
+    /** Returns the container image to use for given node type */
     public DockerImage imageFor(NodeType type) {
         NodeType typeToUseForLookup = type.isHost() ? type.childNodeType() : type;
         DockerImage image = getImages().get(typeToUseForLookup);
         if (image == null) {
-            return defaultImage;
+            image = defaultImage;
         }
-        return image.withRegistry(defaultImage.registry()); // Always use the registry configured for this zone.
+        return rewriteRegistry(image);
     }
 
     /** Set the docker image for nodes of given type */
@@ -84,4 +82,13 @@ public class ContainerImages {
         }
     }
 
+    /** Rewrite the registry part of given image, using this zone's default image */
+    private DockerImage rewriteRegistry(DockerImage image) {
+        DockerImage zoneImage = defaultImage;
+        if (zoneImage.replacedBy().isPresent() && replaceImage.value()) {
+            zoneImage = zoneImage.replacedBy().get();
+        }
+        return image.withRegistry(zoneImage.registry());
+    }
+
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
index 320feed5d66..94b670f8e5e 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
@@ -4,6 +4,7 @@ package com.yahoo.vespa.hosted.provision.provisioning;
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import org.junit.Test;
 
@@ -48,4 +49,26 @@ public class ContainerImagesTest {
         }
     }
 
+    @Test
+    public void image_replacement() {
+        var flagSource = new InMemoryFlagSource();
+        var defaultImage = DockerImage.fromString("foo.example.com/vespa/vespa")
+                                      .withReplacedBy(DockerImage.fromString("bar.example.com/vespa/vespa"));
+        var tester = new ProvisioningTester.Builder().defaultImage(defaultImage).flagSource(flagSource).build();
+        var hosts = tester.makeReadyNodes(2, "default", NodeType.host);
+        tester.activateTenantHosts();
+
+        // Default image is used initially
+        for (var host : hosts) {
+            assertEquals(defaultImage, tester.nodeRepository().containerImages().imageFor(host.type()));
+        }
+
+        // Enabling flag switches to replacement
+        flagSource.withBooleanFlag(Flags.REGIONAL_CONTAINER_REGISTRY.id(), true);
+        for (var host : hosts) {
+            assertEquals(defaultImage.replacedBy().get().asString(),
+                         tester.nodeRepository().containerImages().imageFor(host.type()).asString());
+        }
+    }
+
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
index dd8b04ebf69..994f35afe60 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
@@ -89,11 +89,12 @@ public class ProvisioningTester {
     private int nextHost = 0;
     private int nextIP = 0;
 
-    public ProvisioningTester(Curator curator,
+    private ProvisioningTester(Curator curator,
                               NodeFlavors nodeFlavors,
                               HostResourcesCalculator resourcesCalculator,
                               Zone zone,
                               NameResolver nameResolver,
+                              DockerImage containerImage,
                               Orchestrator orchestrator,
                               HostProvisioner hostProvisioner,
                               LoadBalancerServiceMock loadBalancerService,
@@ -109,7 +110,7 @@ public class ProvisioningTester {
                                                  clock,
                                                  zone,
                                                  nameResolver,
-                                                 DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa"),
+                                                 containerImage,
                                                  flagSource,
                                                  true,
                                                  spareCount,
@@ -596,9 +597,9 @@ public class ProvisioningTester {
         private NameResolver nameResolver;
         private Orchestrator orchestrator;
         private HostProvisioner hostProvisioner;
-        private LoadBalancerServiceMock loadBalancerService;
         private FlagSource flagSource;
         private int spareCount = 0;
+        private DockerImage defaultImage = DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa");
 
         public Builder curator(Curator curator) {
             this.curator = curator;
@@ -636,6 +637,11 @@ public class ProvisioningTester {
             return this;
         }
 
+        public Builder defaultImage(DockerImage defaultImage) {
+            this.defaultImage = defaultImage;
+            return this;
+        }
+
         public Builder orchestrator(Orchestrator orchestrator) {
             this.orchestrator = orchestrator;
             return this;
@@ -646,11 +652,6 @@ public class ProvisioningTester {
             return this;
         }
 
-        public Builder loadBalancerService(LoadBalancerServiceMock loadBalancerService) {
-            this.loadBalancerService = loadBalancerService;
-            return this;
-        }
-
         public Builder flagSource(FlagSource flagSource) {
             this.flagSource = flagSource;
             return this;
@@ -678,9 +679,10 @@ public class ProvisioningTester {
                                           resourcesCalculator,
                                           Optional.ofNullable(zone).orElseGet(Zone::defaultZone),
                                           Optional.ofNullable(nameResolver).orElseGet(() -> new MockNameResolver().mockAnyLookup()),
+                                          defaultImage,
                                           orchestrator,
                                           hostProvisioner,
-                                          Optional.ofNullable(loadBalancerService).orElseGet(LoadBalancerServiceMock::new),
+                                          new LoadBalancerServiceMock(),
                                           Optional.ofNullable(flagSource).orElseGet(InMemoryFlagSource::new),
                                           spareCount);
         }