Propagate account archive URIs to NR

3 files changed, 55 insertions, 21 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdater.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdater.java
index ddb1365f2de..1083e545b33 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdater.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdater.java
@@ -1,10 +1,12 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.maintenance;
 
+import com.yahoo.config.provision.CloudAccount;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.ApplicationController;
 import com.yahoo.vespa.hosted.controller.Controller;
+import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.ArchiveUriUpdate;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.ArchiveUris;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.NodeRepository;
@@ -18,6 +20,7 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.logging.Level;
+import java.util.stream.Stream;
 
 /**
  * Updates archive URIs for tenants in all zones.
@@ -42,14 +45,19 @@ public class ArchiveUriUpdater extends ControllerMaintainer {
     @Override
     protected double maintain() {
         Map<ZoneId, Set<TenantName>> tenantsByZone = new HashMap<>();
+        Map<ZoneId, Set<CloudAccount>> accountsByZone = new HashMap<>();
 
-        controller().zoneRegistry().zonesIncludingSystem().reachable().zones().forEach(
-                z -> tenantsByZone.put(z.getVirtualId(), new HashSet<>(INFRASTRUCTURE_TENANTS)));
+        controller().zoneRegistry().zonesIncludingSystem().reachable().zones().forEach(zone -> {
+            tenantsByZone.put(zone.getVirtualId(), new HashSet<>(INFRASTRUCTURE_TENANTS));
+            accountsByZone.put(zone.getVirtualId(), new HashSet<>());
+        });
 
         for (var application : applications.asList()) {
             for (var instance : application.instances().values()) {
                 for (var deployment : instance.deployments().values()) {
                     tenantsByZone.get(deployment.zone()).add(instance.id().tenant());
+                    applications.decideCloudAccountOf(new DeploymentId(instance.id(), deployment.zone()), application.deploymentSpec())
+                            .ifPresent(account -> accountsByZone.get(deployment.zone()).add(account));
                 }
             }
         }
@@ -59,17 +67,29 @@ public class ArchiveUriUpdater extends ControllerMaintainer {
             try {
                 ArchiveUris zoneArchiveUris = nodeRepository.getArchiveUris(zone);
 
-                for (TenantName tenant : tenantsByZone.get(zone)) {
-                    archiveBucketDb.archiveUriFor(zone, tenant, true)
-                            .filter(uri -> !uri.equals(zoneArchiveUris.tenantArchiveUris().get(tenant)))
-                            .ifPresent(uri -> nodeRepository.updateArchiveUri(zone, ArchiveUriUpdate.setArchiveUriFor(tenant, uri)));
-                }
-
-                zoneArchiveUris.tenantArchiveUris().keySet().stream()
-                        .filter(tenant -> !tenantsByZone.get(zone).contains(tenant))
-                        .forEach(tenant -> nodeRepository.updateArchiveUri(zone, ArchiveUriUpdate.deleteArchiveUriFor(tenant)));
-
-                // TODO (freva): Update account archive URIs
+                Stream.of(
+                        // Tenant URIs that need to be added or updated
+                        tenantsByZone.get(zone).stream()
+                                .flatMap(tenant -> archiveBucketDb.archiveUriFor(zone, tenant, true)
+                                        .filter(uri -> !uri.equals(zoneArchiveUris.tenantArchiveUris().get(tenant)))
+                                        .map(uri -> ArchiveUriUpdate.setArchiveUriFor(tenant, uri))
+                                        .stream()),
+                        // Account URIs that need to be added or updated
+                        accountsByZone.get(zone).stream()
+                                .flatMap(account -> archiveBucketDb.archiveUriFor(zone, account, true)
+                                        .filter(uri -> !uri.equals(zoneArchiveUris.accountArchiveUris().get(account)))
+                                        .map(uri -> ArchiveUriUpdate.setArchiveUriFor(account, uri))
+                                        .stream()),
+                        // Tenant URIs that need to be deleted
+                        zoneArchiveUris.tenantArchiveUris().keySet().stream()
+                                .filter(tenant -> !tenantsByZone.get(zone).contains(tenant))
+                                .map(ArchiveUriUpdate::deleteArchiveUriFor),
+                        // Account URIs that need to be deleted
+                        zoneArchiveUris.accountArchiveUris().keySet().stream()
+                                .filter(account -> !accountsByZone.get(zone).contains(account))
+                                .map(ArchiveUriUpdate::deleteArchiveUriFor))
+                        .flatMap(s -> s)
+                        .forEach(update -> nodeRepository.updateArchiveUri(zone, update));
             } catch (Exception e) {
                 log.log(Level.WARNING, "Failed to update archive URI in " + zone + ". Retrying in " + interval() + ". Error: " +
                         Exceptions.toMessageString(e));
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java
index 38ff9967ef6..de186109784 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java
@@ -266,7 +266,7 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
 
     @Override
     public boolean hasZone(ZoneId zoneId, CloudAccount cloudAccount) {
-        return hasZone(zoneId) && cloudAccountZones.getOrDefault(cloudAccount, Set.of()).contains(zoneId);
+        return hasZone(zoneId) && (system.isPublic() || cloudAccountZones.getOrDefault(cloudAccount, Set.of()).contains(zoneId));
     }
 
     @Override
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdaterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdaterTest.java
index ad37c4a985d..d6b59ef860f 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdaterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/ArchiveUriUpdaterTest.java
@@ -5,21 +5,24 @@ import com.yahoo.config.provision.CloudAccount;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.zone.ZoneId;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
+import com.yahoo.vespa.flags.PermanentFlags;
 import com.yahoo.vespa.hosted.controller.ControllerTester;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.ArchiveBuckets;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.ArchiveUriUpdate;
+import com.yahoo.vespa.hosted.controller.api.integration.archive.MockArchiveService;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.VespaManagedArchiveBucket;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.ArchiveUris;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.NodeRepository;
-import com.yahoo.vespa.hosted.controller.api.integration.deployment.JobType;
 import com.yahoo.vespa.hosted.controller.application.SystemApplication;
-import com.yahoo.vespa.hosted.controller.application.pkg.ApplicationPackage;
+import com.yahoo.vespa.hosted.controller.deployment.ApplicationPackageBuilder;
 import com.yahoo.vespa.hosted.controller.deployment.DeploymentContext;
 import com.yahoo.vespa.hosted.controller.deployment.DeploymentTester;
 import org.junit.jupiter.api.Test;
 
 import java.net.URI;
 import java.time.Duration;
+import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 
@@ -38,6 +41,7 @@ public class ArchiveUriUpdaterTest {
 
         var tenant1 = TenantName.from("tenant1");
         var tenant2 = TenantName.from("tenant2");
+        var account1 = CloudAccount.from("001122334455");
         var tenantInfra = SystemApplication.TENANT;
         var application = tester.newDeploymentContext(tenant1.value(), "app1", "instance1");
         ZoneId zone = ZoneId.from("prod", "aws-us-east-1c");
@@ -49,17 +53,20 @@ public class ArchiveUriUpdaterTest {
 
         // Archive service now has URI for tenant1, but tenant1 is not deployed in zone
         setBucketNameInService(Map.of(tenant1, "uri-1"), zone);
+        setAccountBucketNameInService(zone, account1, "bkt-1");
         updater.maintain();
         assertArchiveUris(zone, Map.of(TenantName.from("hosted-vespa"), "s3://bucketName/"), Map.of());
 
-        deploy(application, zone);
+        ((InMemoryFlagSource) tester.controller().flagSource())
+                .withListFlag(PermanentFlags.CLOUD_ACCOUNTS.id(), List.of(account1.value()), String.class);
+        deploy(application, zone, account1);
         updater.maintain();
-        assertArchiveUris(zone, Map.of(tenant1, "s3://uri-1/", tenantInfra, "s3://bucketName/"), Map.of());
+        assertArchiveUris(zone, Map.of(tenant1, "s3://uri-1/", tenantInfra, "s3://bucketName/"), Map.of(account1, "s3://bkt-1/"));
 
         // URI for tenant1 should be updated and removed for tenant2
         setArchiveUriInNodeRepo(Map.of(tenant1, "wrong-uri", tenant2, "uri-2"), zone);
         updater.maintain();
-        assertArchiveUris(zone, Map.of(tenant1, "s3://uri-1/", tenantInfra, "s3://bucketName/"), Map.of());
+        assertArchiveUris(zone, Map.of(tenant1, "s3://uri-1/", tenantInfra, "s3://bucketName/"), Map.of(account1, "s3://bkt-1/"));
     }
 
     private void assertArchiveUris(ZoneId zone, Map<TenantName, String> expectedTenantUris, Map<CloudAccount, String> expectedAccountUris) {
@@ -75,13 +82,20 @@ public class ArchiveUriUpdaterTest {
         tester.controller().curator().writeArchiveBuckets(zone, buckets);
     }
 
+    private void setAccountBucketNameInService(ZoneId zone, CloudAccount cloudAccount, String bucketName) {
+        ((MockArchiveService) tester.controller().serviceRegistry().archiveService()).setEnclaveArchiveBucket(zone, cloudAccount, bucketName);
+    }
+
     private void setArchiveUriInNodeRepo(Map<TenantName, String> archiveUris, ZoneId zone) {
         NodeRepository nodeRepository = tester.controller().serviceRegistry().configServer().nodeRepository();
         archiveUris.forEach((tenant, uri) -> nodeRepository.updateArchiveUri(zone, ArchiveUriUpdate.setArchiveUriFor(tenant, URI.create(uri))));
     }
 
-    private void deploy(DeploymentContext application, ZoneId zone) {
-        application.runJob(JobType.deploymentTo(zone), new ApplicationPackage(new byte[0]));
+    private void deploy(DeploymentContext application, ZoneId zone, CloudAccount cloudAccount) {
+        application.submit(new ApplicationPackageBuilder()
+                .cloudAccount(cloudAccount.value())
+                .region(zone.region().value())
+                .build()).deploy();
     }
 
 }