1 files changed, 16 insertions, 36 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 1c930b04989..e9734359647 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
@@ -4,15 +4,14 @@ package com.yahoo.vespa.hosted.controller.archive;
 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.ArchiveBucket;
+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.VespaManagedArchiveBucket;
 import com.yahoo.vespa.hosted.controller.persistence.CuratorDb;
 
 import java.net.URI;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
@@ -39,62 +38,43 @@ public class CuratorArchiveBucketDb {
 
     public Optional<URI> archiveUriFor(ZoneId zoneId, TenantName tenant, boolean createIfMissing) {
         return getBucketNameFromCache(zoneId, tenant)
-                .or(() -> findAndUpdateArchiveUriCache(zoneId, tenant, buckets(zoneId)))
                 .or(() -> createIfMissing ? Optional.of(assignToBucket(zoneId, tenant)) : Optional.empty())
                 .map(bucketName -> archiveService.bucketURI(zoneId, bucketName));
     }
 
     private String assignToBucket(ZoneId zoneId, TenantName tenant) {
         try (var lock = curatorDb.lockArchiveBuckets(zoneId)) {
-            Set<ArchiveBucket> zoneBuckets = new HashSet<>(buckets(zoneId));
+            ArchiveBuckets archiveBuckets = buckets(zoneId);
+            updateArchiveUriCache(zoneId, archiveBuckets);
 
-            return findAndUpdateArchiveUriCache(zoneId, tenant, zoneBuckets) // Some other thread might have assigned it before we grabbed the lock
+            return getBucketNameFromCache(zoneId, tenant) // Some other thread might have assigned it before we grabbed the lock
                     .orElseGet(() -> {
                         // If not, find an existing bucket with space
-                        Optional<ArchiveBucket> unfilledBucket = zoneBuckets.stream()
+                        VespaManagedArchiveBucket bucketToAssignTo = archiveBuckets.vespaManaged().stream()
                                 .filter(bucket -> archiveService.canAddTenantToBucket(zoneId, bucket))
-                                .findAny();
+                                .findAny()
+                                // Or create a new one
+                                .orElseGet(() -> archiveService.createArchiveBucketFor(zoneId));
 
-                        // And place the tenant in that bucket.
-                        if (unfilledBucket.isPresent()) {
-                            var unfilled = unfilledBucket.get();
+                        ArchiveBuckets updated = archiveBuckets.with(bucketToAssignTo.withTenant(tenant));
+                        curatorDb.writeArchiveBuckets(zoneId, updated);
+                        updateArchiveUriCache(zoneId, updated);
 
-                            zoneBuckets.remove(unfilled);
-                            zoneBuckets.add(unfilled.withTenant(tenant));
-                            curatorDb.writeArchiveBuckets(zoneId, zoneBuckets);
-
-                            return unfilled.bucketName();
-                        }
-
-                        // We'll have to create a new bucket
-                        var newBucket = archiveService.createArchiveBucketFor(zoneId).withTenant(tenant);
-                        zoneBuckets.add(newBucket);
-                        curatorDb.writeArchiveBuckets(zoneId, zoneBuckets);
-                        updateArchiveUriCache(zoneId, zoneBuckets);
-                        return newBucket.bucketName();
+                        return bucketToAssignTo.bucketName();
                     });
         }
     }
 
-    public Set<ArchiveBucket> buckets(ZoneId zoneId) {
+    public ArchiveBuckets buckets(ZoneId zoneId) {
         return curatorDb.readArchiveBuckets(zoneId);
     }
 
-    private Optional<String> findAndUpdateArchiveUriCache(ZoneId zoneId, TenantName tenant, Set<ArchiveBucket> zoneBuckets) {
-        Optional<String> bucketName = zoneBuckets.stream()
-                .filter(bucket -> bucket.tenants().contains(tenant))
-                .findAny()
-                .map(ArchiveBucket::bucketName);
-        if (bucketName.isPresent()) updateArchiveUriCache(zoneId, zoneBuckets);
-        return bucketName;
-    }
-
     private Optional<String> getBucketNameFromCache(ZoneId zoneId, TenantName tenantName) {
         return Optional.ofNullable(archiveUriCache.get(zoneId)).map(map -> map.get(tenantName));
     }
 
-    private void updateArchiveUriCache(ZoneId zoneId, Set<ArchiveBucket> zoneBuckets) {
-        Map<TenantName, String> bucketNameByTenant = zoneBuckets.stream()
+    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())))
                 .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));