Store archiveUri in node-repo ZK

8 files changed, 260 insertions, 86 deletions

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 aae64d0563c..f7ed4ebad3a 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
@@ -23,6 +23,7 @@ import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
 import com.yahoo.vespa.hosted.provision.persistence.DnsNameResolver;
 import com.yahoo.vespa.hosted.provision.persistence.JobControlFlags;
 import com.yahoo.vespa.hosted.provision.persistence.NameResolver;
+import com.yahoo.vespa.hosted.provision.provisioning.ArchiveUris;
 import com.yahoo.vespa.hosted.provision.provisioning.ContainerImages;
 import com.yahoo.vespa.hosted.provision.provisioning.FirmwareChecks;
 import com.yahoo.vespa.hosted.provision.provisioning.HostResourcesCalculator;
@@ -57,6 +58,7 @@ public class NodeRepository extends AbstractComponent {
     private final InfrastructureVersions infrastructureVersions;
     private final FirmwareChecks firmwareChecks;
     private final ContainerImages containerImages;
+    private final ArchiveUris archiveUris;
     private final JobControl jobControl;
     private final Applications applications;
     private final LoadBalancers loadBalancers;
@@ -119,6 +121,7 @@ public class NodeRepository extends AbstractComponent {
         this.infrastructureVersions = new InfrastructureVersions(db);
         this.firmwareChecks = new FirmwareChecks(db, clock);
         this.containerImages = new ContainerImages(db, containerImage);
+        this.archiveUris = new ArchiveUris(db);
         this.jobControl = new JobControl(new JobControlFlags(db, flagSource));
         this.applications = new Applications(db);
         this.loadBalancers = new LoadBalancers(db);
@@ -162,6 +165,9 @@ public class NodeRepository extends AbstractComponent {
     /** Returns the container images to use for nodes in this. */
     public ContainerImages containerImages() { return containerImages; }
 
+    /** Returns the archive URIs to use for nodes in this. */
+    public ArchiveUris archiveUris() { return archiveUris; }
+
     /** Returns the status of maintenance jobs managed by this. */
     public JobControl jobControl() { return jobControl; }
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
index 6150ee9f4a0..c2fe063dae6 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
@@ -10,6 +10,7 @@ import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.NodeFlavors;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.path.Path;
 import com.yahoo.transaction.NestedTransaction;
@@ -69,6 +70,7 @@ public class CuratorDatabaseClient {
     private static final Path osVersionsPath = root.append("osVersions");
     private static final Path containerImagesPath = root.append("dockerImages");
     private static final Path firmwareCheckPath = root.append("firmwareCheck");
+    private static final Path archiveUrisPath = root.append("archiveUris");
 
     private static final Duration defaultLockTimeout = Duration.ofMinutes(6);
 
@@ -102,6 +104,7 @@ public class CuratorDatabaseClient {
         db.create(osVersionsPath);
         db.create(containerImagesPath);
         db.create(firmwareCheckPath);
+        db.create(archiveUrisPath);
         db.create(loadBalancersPath);
         provisionIndexCounter.initialize(100);
     }
@@ -461,6 +464,24 @@ public class CuratorDatabaseClient {
         return read(firmwareCheckPath, data -> Instant.ofEpochMilli(Long.parseLong(new String(data))));
     }
 
+    // Archive URIs -----------------------------------------------------------
+
+    public void writeArchiveUris(Map<TenantName, String> archiveUris) {
+        byte[] data = TenantArchiveUriSerializer.toJson(archiveUris);
+        NestedTransaction transaction = new NestedTransaction();
+        CuratorTransaction curatorTransaction = db.newCuratorTransactionIn(transaction);
+        curatorTransaction.add(CuratorOperations.setData(archiveUrisPath.getAbsolute(), data));
+        transaction.commit();
+    }
+
+    public Map<TenantName, String> readArchiveUris() {
+        return read(archiveUrisPath, TenantArchiveUriSerializer::fromJson).orElseGet(Map::of);
+    }
+
+    public Lock lockArchiveUris() {
+        return db.lock(lockPath.append("archiveUris"), defaultLockTimeout);
+    }
+
     // Load balancers -----------------------------------------------------------
 
     public List<LoadBalancerId> readLoadBalancerIds() {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/TenantArchiveUriSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/TenantArchiveUriSerializer.java
new file mode 100644
index 00000000000..3f0c8f90887
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/TenantArchiveUriSerializer.java
@@ -0,0 +1,45 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.persistence;
+
+import com.yahoo.config.provision.TenantName;
+import com.yahoo.slime.Cursor;
+import com.yahoo.slime.Inspector;
+import com.yahoo.slime.ObjectTraverser;
+import com.yahoo.slime.Slime;
+import com.yahoo.slime.SlimeUtils;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * Serializer for archive URIs that are set per tenant.
+ *
+ * @author freva
+ */
+public class TenantArchiveUriSerializer {
+
+    private TenantArchiveUriSerializer() {}
+
+    public static byte[] toJson(Map<TenantName, String> archiveUrisByTenantName) {
+        Slime slime = new Slime();
+        Cursor object = slime.setObject();
+        archiveUrisByTenantName.forEach((tenantName, archiveUri) ->
+                object.setString(tenantName.value(), archiveUri));
+        try {
+            return SlimeUtils.toJsonBytes(slime);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    public static Map<TenantName, String> fromJson(byte[] data) {
+        Map<TenantName, String> archiveUrisByTenantName = new TreeMap<>(); // Use TreeMap to sort by tenant name
+        Inspector inspector = SlimeUtils.jsonToSlime(data).get();
+        inspector.traverse((ObjectTraverser) (key, value) ->
+                archiveUrisByTenantName.put(TenantName.from(key), value.asString()));
+        return archiveUrisByTenantName;
+    }
+
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ArchiveUris.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ArchiveUris.java
new file mode 100644
index 00000000000..4c0b0f284c4
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ArchiveUris.java
@@ -0,0 +1,86 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.provisioning;
+
+import com.yahoo.config.provision.TenantName;
+import com.yahoo.lang.CachedSupplier;
+import com.yahoo.vespa.curator.Lock;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.node.Allocation;
+import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
+
+import java.time.Duration;
+import java.util.Map;
+import java.util.Optional;
+import java.util.TreeMap;
+import java.util.logging.Logger;
+import java.util.regex.Pattern;
+
+/**
+ * Thread safe class to get and set archive URI for given tenants. Archive URIs are stored in ZooKeeper so that
+ * nodes within the same tenant have the same archive URI from all the config servers.
+ *
+ * @author freva
+ */
+public class ArchiveUris {
+
+    private static final Logger log = Logger.getLogger(ArchiveUris.class.getName());
+    private static final Pattern validUriPattern = Pattern.compile("[a-z0-9]+://(?:(?:[a-z0-9]+(?:-[a-z0-9])*)+/)+");
+    private static final Duration cacheTtl = Duration.ofMinutes(1);
+
+    private final CuratorDatabaseClient db;
+    private final CachedSupplier<Map<TenantName, String>> archiveUris;
+
+    public ArchiveUris(CuratorDatabaseClient db) {
+        this.db = db;
+        this.archiveUris = new CachedSupplier<>(db::readArchiveUris, cacheTtl);
+    }
+
+    /** Returns the current archive URI for each tenant */
+    public Map<TenantName, String> getArchiveUris() {
+        return archiveUris.get();
+    }
+
+    /** Returns the archive URI to use for given tenant */
+    public Optional<String> archiveUriFor(TenantName tenant) {
+        return Optional.ofNullable(archiveUris.get().get(tenant));
+    }
+
+    /** Returns the archive URI to use for given node */
+    public Optional<String> archiveUriFor(Node node) {
+        return node.allocation().map(Allocation::owner)
+                .flatMap(app -> archiveUriFor(app.tenant())
+                        .map(uri -> {
+                            StringBuilder sb = new StringBuilder(100).append(uri)
+                                    .append(app.application().value()).append('/')
+                                    .append(app.instance().value()).append('/');
+
+                            for (char c: node.hostname().toCharArray()) {
+                                if (c == '.') break;
+                                sb.append(c);
+                            }
+
+                            return sb.append('/').toString();
+                        }));
+    }
+
+    /** Set the docker image for nodes of given type */
+    public void setArchiveUri(TenantName tenant, Optional<String> archiveUri) {
+        try (Lock lock = db.lockArchiveUris()) {
+            Map<TenantName, String> archiveUris = new TreeMap<>(db.readArchiveUris());
+            if (Optional.ofNullable(archiveUris.get(tenant)).equals(archiveUri)) return; // No change
+
+            archiveUri.map(ArchiveUris::normalizeUri).ifPresentOrElse(uri -> archiveUris.put(tenant, uri),
+                                                                      () -> archiveUris.remove(tenant));
+            db.writeArchiveUris(archiveUris);
+            this.archiveUris.refresh(); // Throw away current cache
+            log.info("Set archive URI for " + tenant + " to " + archiveUri.orElse(null));
+        }
+    }
+
+    static String normalizeUri(String uri) {
+        if (!uri.endsWith("/")) uri = uri + "/";
+        if (!validUriPattern.matcher(uri).matches())
+            throw new IllegalArgumentException("Invalid archive URI: " + uri);
+        return uri;
+    }
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
index eaef044e07c..2c61d8092b6 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
@@ -11,9 +11,6 @@ import com.yahoo.container.jdisc.HttpRequest;
 import com.yahoo.restapi.SlimeJsonResponse;
 import com.yahoo.slime.Cursor;
 import com.yahoo.vespa.applicationmodel.HostName;
-import com.yahoo.vespa.flags.FetchVector;
-import com.yahoo.vespa.flags.FlagSource;
-import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Address;
@@ -179,7 +176,7 @@ class NodesResponse extends SlimeJsonResponse {
         node.reports().toSlime(object, "reports");
         node.modelName().ifPresent(modelName -> object.setString("modelName", modelName));
         node.switchHostname().ifPresent(switchHostname -> object.setString("switchHostname", switchHostname));
-        archiveUri(nodeRepository.flagSource(), node).ifPresent(uri -> object.setString("archiveUri", uri));
+        nodeRepository.archiveUris().archiveUriFor(node).ifPresent(uri -> object.setString("archiveUri", uri));
     }
 
     private void toSlime(ApplicationId id, Cursor object) {
@@ -234,30 +231,4 @@ class NodesResponse extends SlimeJsonResponse {
         return path.substring(lastSlash+1);
     }
 
-    // TODO (freva): Store this in Application or Node
-    static Optional<String> archiveUri(FlagSource flagSource, Node node) {
-        String bucket = Flags.SYNC_HOST_LOGS_TO_S3_BUCKET.bindTo(flagSource)
-                .with(FetchVector.Dimension.NODE_TYPE, node.type().name())
-                .with(FetchVector.Dimension.APPLICATION_ID, node.allocation().map(alloc -> alloc.owner().serializedForm()).orElse(null))
-                .value();
-        if (bucket.isBlank()) return Optional.empty();
-
-        StringBuilder sb = new StringBuilder(100).append("s3://").append(bucket).append('/');
-        if (node.type() == NodeType.tenant) {
-            if (node.allocation().isEmpty()) return Optional.empty();
-            ApplicationId app = node.allocation().get().owner();
-
-            sb.append(app.tenant().value()).append('/').append(app.application().value()).append('/').append(app.instance().value()).append('/');
-        } else {
-            sb.append("hosted-vespa/");
-        }
-
-        for (char c: node.hostname().toCharArray()) {
-            if (c == '.') break;
-            sb.append(c);
-        }
-
-        return Optional.of(sb.append('/').toString());
-    }
-
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/TenantArchiveUriSerializerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/TenantArchiveUriSerializerTest.java
new file mode 100644
index 00000000000..ce32a38b049
--- /dev/null
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/TenantArchiveUriSerializerTest.java
@@ -0,0 +1,27 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.persistence;
+
+import com.yahoo.config.provision.TenantName;
+import org.junit.Test;
+
+import java.util.Map;
+import java.util.TreeMap;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author freva
+ */
+public class TenantArchiveUriSerializerTest {
+
+    @Test
+    public void test_serialization() {
+        Map<TenantName, String> archiveUris = new TreeMap<>();
+        archiveUris.put(TenantName.from("tenant1"), "ftp://host123.test/dir/");
+        archiveUris.put(TenantName.from("tenant2"), "ftp://archive.test/vespa/");
+
+        Map<TenantName, String> serialized = TenantArchiveUriSerializer.fromJson(TenantArchiveUriSerializer.toJson(archiveUris));
+        assertEquals(archiveUris, serialized);
+    }
+
+}
\ No newline at end of file
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ArchiveUrisTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ArchiveUrisTest.java
new file mode 100644
index 00000000000..c6848e2af2a
--- /dev/null
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ArchiveUrisTest.java
@@ -0,0 +1,74 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.provisioning;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.ClusterMembership;
+import com.yahoo.config.provision.Flavor;
+import com.yahoo.config.provision.NodeResources;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.node.Allocation;
+import com.yahoo.vespa.hosted.provision.node.Generation;
+import org.junit.Test;
+
+import java.util.Optional;
+
+import static com.yahoo.vespa.hosted.provision.provisioning.ArchiveUris.normalizeUri;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.fail;
+
+/**
+ * @author freva
+ */
+public class ArchiveUrisTest {
+
+    @Test
+    public void archive_uri() {
+        ApplicationId app = ApplicationId.from("vespa", "music", "main");
+        Node allocated = createNode(app);
+        Node unallocated = createNode(null);
+        ArchiveUris archiveUris = new ProvisioningTester.Builder().build().nodeRepository().archiveUris();
+
+        assertFalse(archiveUris.archiveUriFor(unallocated).isPresent());
+        assertFalse(archiveUris.archiveUriFor(allocated).isPresent());
+
+        archiveUris.setArchiveUri(app.tenant(), Optional.of("scheme://hostname/dir"));
+        assertEquals("scheme://hostname/dir/music/main/h432a/", archiveUris.archiveUriFor(allocated).get());
+    }
+
+    private Node createNode(ApplicationId appId) {
+        Node.Builder nodeBuilder = Node.create("id", "h432a.prod.us-south-1.vespa.domain.tld", new Flavor(NodeResources.unspecified()), Node.State.parked, NodeType.tenant);
+        Optional.ofNullable(appId)
+                .map(app -> new Allocation(app,
+                        ClusterMembership.from("container/default/0/0", Version.fromString("1.2.3"), Optional.empty()),
+                        NodeResources.unspecified(),
+                        Generation.initial(),
+                        false))
+                .ifPresent(nodeBuilder::allocation);
+        return nodeBuilder.build();
+    }
+
+    @Test
+    public void normalize_test() {
+        assertEquals("ftp://domain/legal-dir123/", normalizeUri("ftp://domain/legal-dir123"));
+        assertEquals("ftp://domain/legal-dir123/", normalizeUri("ftp://domain/legal-dir123/"));
+        assertEquals("s3://my-bucket-name/my-tenant-123/", normalizeUri("s3://my-bucket-name/my-tenant-123/"));
+        assertThrows(IllegalArgumentException.class, () -> normalizeUri("domain/dir/"));
+        assertThrows(IllegalArgumentException.class, () -> normalizeUri("ftp:/domain/dir/"));
+        assertThrows(IllegalArgumentException.class, () -> normalizeUri("ftp:/domain//dir/"));
+        assertThrows(IllegalArgumentException.class, () -> normalizeUri("ftp://domain/illegal_dir/"));
+        assertThrows(IllegalArgumentException.class, () -> normalizeUri("ftp://domain/-illegal-dir/"));
+        assertThrows(IllegalArgumentException.class, () -> normalizeUri("ftp://domain/illegal-dir-/"));
+    }
+
+    private static void assertThrows(Class<? extends Throwable> clazz, Runnable runnable) {
+        try {
+            runnable.run();
+            fail("Expected " + clazz);
+        } catch (Throwable e) {
+            if (!clazz.isInstance(e)) throw e;
+        }
+    }
+}
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponseTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponseTest.java
deleted file mode 100644
index 021afd0df7c..00000000000
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponseTest.java
+++ /dev/null
@@ -1,56 +0,0 @@
-// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.hosted.provision.restapi;
-
-import com.yahoo.component.Version;
-import com.yahoo.config.provision.ApplicationId;
-import com.yahoo.config.provision.ClusterMembership;
-import com.yahoo.config.provision.Flavor;
-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 com.yahoo.vespa.hosted.provision.Node;
-import com.yahoo.vespa.hosted.provision.node.Allocation;
-import com.yahoo.vespa.hosted.provision.node.Generation;
-import org.junit.Test;
-
-import java.util.Optional;
-
-import static org.junit.Assert.assertEquals;
-
-/**
- * @author freva
- */
-public class NodesResponseTest {
-    private final InMemoryFlagSource flagSource = new InMemoryFlagSource();
-
-    @Test
-    public void archive_uri() {
-        ApplicationId app = ApplicationId.from("vespa", "music", "main");
-        // Flag not set, no archive uri
-        assertArchiveUri(null, "h432a.prod.us-south-1.vespa.domain.tld", NodeType.tenant, app);
-        assertArchiveUri(null, "cfg1.prod.us-south-1.vespa.domain.tld", NodeType.config, app);
-
-        flagSource.withStringFlag(Flags.SYNC_HOST_LOGS_TO_S3_BUCKET.id(), "vespa-data-bucket");
-        // Flag is set, but node not allocated, only sync non-tenant nodes
-        assertArchiveUri(null, "h432a.prod.us-south-1.vespa.domain.tld", NodeType.tenant, null);
-        assertArchiveUri("s3://vespa-data-bucket/hosted-vespa/cfg1/", "cfg1.prod.us-south-1.vespa.domain.tld", NodeType.config, null);
-
-        // Flag is set and node is allocated
-        assertArchiveUri("s3://vespa-data-bucket/vespa/music/main/h432a/", "h432a.prod.us-south-1.vespa.domain.tld", NodeType.tenant, app);
-        assertArchiveUri("s3://vespa-data-bucket/hosted-vespa/cfg1/", "cfg1.prod.us-south-1.vespa.domain.tld", NodeType.config, app);
-    }
-
-    private void assertArchiveUri(String archiveUri, String hostname, NodeType type, ApplicationId appId) {
-        Node.Builder nodeBuilder = Node.create("id", hostname, new Flavor(NodeResources.unspecified()), Node.State.parked, type);
-        Optional.ofNullable(appId)
-                .map(app -> new Allocation(app,
-                        ClusterMembership.from("container/default/0/0", Version.fromString("1.2.3"), Optional.empty()),
-                        NodeResources.unspecified(),
-                        Generation.initial(),
-                        false))
-                .ifPresent(nodeBuilder::allocation);
-
-        assertEquals(archiveUri, NodesResponse.archiveUri(flagSource, nodeBuilder.build()).orElse(null));
-    }
-}
\ No newline at end of file