Find archive URI for given account from ZK cache or ArchiveService

3 files changed, 78 insertions, 5 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDb.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDb.java
index e9734359647..962bd144a21 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDb.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDb.java
@@ -1,15 +1,20 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.archive;
 
+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.Controller;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.ArchiveBuckets;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.ArchiveService;
+import com.yahoo.vespa.hosted.controller.api.integration.archive.TenantManagedArchiveBucket;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.VespaManagedArchiveBucket;
 import com.yahoo.vespa.hosted.controller.persistence.CuratorDb;
 
 import java.net.URI;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
@@ -22,18 +27,23 @@ import java.util.stream.Collectors;
  */
 public class CuratorArchiveBucketDb {
 
+    private static final Duration ENCLAVE_BUCKET_CACHE_LIFETIME = Duration.ofMinutes(60);
+
     /**
      * Archive URIs are often requested because they are returned in /application/v4 API. Since they
      * never change, it's safe to cache them and only update on misses
      */
     private final Map<ZoneId, Map<TenantName, String>> archiveUriCache = new ConcurrentHashMap<>();
+    private final Map<ZoneId, Map<CloudAccount, TenantManagedArchiveBucket>> tenantArchiveCache = new ConcurrentHashMap<>();
 
     private final ArchiveService archiveService;
     private final CuratorDb curatorDb;
+    private final Clock clock;
 
     public CuratorArchiveBucketDb(Controller controller) {
         this.archiveService = controller.serviceRegistry().archiveService();
         this.curatorDb = controller.curator();
+        this.clock = controller.clock();
     }
 
     public Optional<URI> archiveUriFor(ZoneId zoneId, TenantName tenant, boolean createIfMissing) {
@@ -42,6 +52,30 @@ public class CuratorArchiveBucketDb {
                 .map(bucketName -> archiveService.bucketURI(zoneId, bucketName));
     }
 
+    public Optional<URI> archiveUriFor(ZoneId zoneId, CloudAccount account, boolean searchIfMissing) {
+        Instant updatedAfter = searchIfMissing ? clock.instant().minus(ENCLAVE_BUCKET_CACHE_LIFETIME) : Instant.MIN;
+        return getBucketNameFromCache(zoneId, account, updatedAfter)
+                .or(() -> {
+                    if (!searchIfMissing) return Optional.empty();
+                    try (var lock = curatorDb.lockArchiveBuckets(zoneId)) {
+                        ArchiveBuckets archiveBuckets = buckets(zoneId);
+                        updateArchiveUriCache(zoneId, archiveBuckets);
+
+                        return getBucketNameFromCache(zoneId, account, updatedAfter)
+                                .or(() -> archiveService.findEnclaveArchiveBucket(zoneId, account)
+                                        .map(bucketName -> {
+                                            var bucket = new TenantManagedArchiveBucket(bucketName, account, clock.instant());
+                                            ArchiveBuckets updated = archiveBuckets.with(bucket);
+                                            curatorDb.writeArchiveBuckets(zoneId, updated);
+                                            updateArchiveUriCache(zoneId, updated);
+                                            return bucket;
+                                        }));
+                    }
+                })
+                .map(TenantManagedArchiveBucket::bucketName)
+                .map(bucketName -> archiveService.bucketURI(zoneId, bucketName));
+    }
+
     private String assignToBucket(ZoneId zoneId, TenantName tenant) {
         try (var lock = curatorDb.lockArchiveBuckets(zoneId)) {
             ArchiveBuckets archiveBuckets = buckets(zoneId);
@@ -73,11 +107,20 @@ public class CuratorArchiveBucketDb {
         return Optional.ofNullable(archiveUriCache.get(zoneId)).map(map -> map.get(tenantName));
     }
 
+    private Optional<TenantManagedArchiveBucket> getBucketNameFromCache(ZoneId zoneId, CloudAccount cloudAccount, Instant updatedAfter) {
+        return Optional.ofNullable(tenantArchiveCache.get(zoneId))
+                .map(map -> map.get(cloudAccount))
+                .filter(bucket -> bucket.updatedAt().isAfter(updatedAfter));
+    }
+
     private void updateArchiveUriCache(ZoneId zoneId, ArchiveBuckets archiveBuckets) {
         Map<TenantName, String> bucketNameByTenant = archiveBuckets.vespaManaged().stream()
-                .flatMap(bucket -> bucket.tenants().stream()
-                        .map(tenant -> Map.entry(tenant, bucket.bucketName())))
+                .flatMap(bucket -> bucket.tenants().stream().map(tenant -> Map.entry(tenant, bucket.bucketName())))
                 .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));
         archiveUriCache.put(zoneId, bucketNameByTenant);
+
+        Map<CloudAccount, TenantManagedArchiveBucket> bucketByAccount = archiveBuckets.tenantManaged().stream()
+                .collect(Collectors.toUnmodifiableMap(TenantManagedArchiveBucket::cloudAccount, bucket -> bucket));
+        tenantArchiveCache.put(zoneId, bucketByAccount);
     }
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDbTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDbTest.java
index 9770133dcfb..e5571c0e0ca 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDbTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/archive/CuratorArchiveBucketDbTest.java
@@ -1,16 +1,21 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.archive;
 
+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.test.ManualClock;
+import com.yahoo.vespa.hosted.controller.Controller;
 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.MockArchiveService;
 import com.yahoo.vespa.hosted.controller.api.integration.archive.VespaManagedArchiveBucket;
 import org.apache.curator.shaded.com.google.common.collect.Streams;
 import org.junit.jupiter.api.Test;
 
 import java.net.URI;
+import java.time.Duration;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -58,4 +63,29 @@ public class CuratorArchiveBucketDbTest {
                 Set.of(new VespaManagedArchiveBucket("bucketName", "keyArn").withTenant(TenantName.from("firstInZone"))),
                 bucketDb.buckets(ZoneId.from("prod.us-east-3")).vespaManaged());
     }
+
+    @Test
+    void archiveUriForAccount() {
+        Controller controller = new ControllerTester(SystemName.Public).controller();
+        CuratorArchiveBucketDb bucketDb = new CuratorArchiveBucketDb(controller);
+        MockArchiveService service = (MockArchiveService) controller.serviceRegistry().archiveService();
+        ManualClock clock = (ManualClock) controller.clock();
+
+        CloudAccount acc1 = CloudAccount.from("001122334455");
+        ZoneId z1 = ZoneId.from("prod.us-east-3");
+
+        assertEquals(Optional.empty(), bucketDb.archiveUriFor(z1, acc1, true)); // Initially not set
+        service.setEnclaveArchiveBucket(z1, acc1, "bucket-1");
+        assertEquals(Optional.empty(), bucketDb.archiveUriFor(z1, acc1, false));
+        assertEquals(Optional.of(URI.create("s3://bucket-1/")), bucketDb.archiveUriFor(z1, acc1, true));
+        assertEquals(Optional.of(URI.create("s3://bucket-1/")), bucketDb.archiveUriFor(z1, acc1, false));
+
+        service.setEnclaveArchiveBucket(z1, acc1, "bucket-2");
+        assertEquals(Optional.of(URI.create("s3://bucket-1/")), bucketDb.archiveUriFor(z1, acc1, true)); // Returns old value even with search
+
+        clock.advance(Duration.ofMinutes(61)); // After expiry the cache is expired, new search is performed
+        assertEquals(Optional.of(URI.create("s3://bucket-1/")), bucketDb.archiveUriFor(z1, acc1, false)); // When requesting without search, return previous value even if expired
+        assertEquals(Optional.of(URI.create("s3://bucket-2/")), bucketDb.archiveUriFor(z1, acc1, true));
+        assertEquals(Optional.of(URI.create("s3://bucket-2/")), bucketDb.archiveUriFor(z1, acc1, false));
+    }
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java
index be257daa211..998b371bbf1 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java
@@ -37,11 +37,11 @@ import com.yahoo.vespa.hosted.controller.api.integration.organization.MockIssueH
 import com.yahoo.vespa.hosted.controller.api.integration.resource.CostReportConsumerMock;
 import com.yahoo.vespa.hosted.controller.api.integration.resource.ResourceDatabaseClient;
 import com.yahoo.vespa.hosted.controller.api.integration.resource.ResourceDatabaseClientMock;
+import com.yahoo.vespa.hosted.controller.api.integration.secrets.EndpointSecretManager;
 import com.yahoo.vespa.hosted.controller.api.integration.secrets.GcpSecretStore;
-import com.yahoo.vespa.hosted.controller.api.integration.secrets.NoopGcpSecretStore;
 import com.yahoo.vespa.hosted.controller.api.integration.secrets.NoopEndpointSecretManager;
+import com.yahoo.vespa.hosted.controller.api.integration.secrets.NoopGcpSecretStore;
 import com.yahoo.vespa.hosted.controller.api.integration.secrets.NoopTenantSecretService;
-import com.yahoo.vespa.hosted.controller.api.integration.secrets.EndpointSecretManager;
 import com.yahoo.vespa.hosted.controller.api.integration.stubs.DummyOwnershipIssues;
 import com.yahoo.vespa.hosted.controller.api.integration.stubs.DummySystemMonitor;
 import com.yahoo.vespa.hosted.controller.api.integration.stubs.LoggingDeploymentIssues;
@@ -89,7 +89,7 @@ public class ServiceRegistryMock extends AbstractComponent implements ServiceReg
     private final ArtifactRegistryMock containerRegistry = new ArtifactRegistryMock();
     private final NoopTenantSecretService tenantSecretService = new NoopTenantSecretService();
     private final NoopEndpointSecretManager secretManager = new NoopEndpointSecretManager();
-    private final ArchiveService archiveService = new MockArchiveService();
+    private final ArchiveService archiveService = new MockArchiveService(clock);
     private final MockChangeRequestClient changeRequestClient = new MockChangeRequestClient();
     private final AccessControlService accessControlService = new MockAccessControlService();
     private final HorizonClient horizonClient = new MockHorizonClient();