Revert "Take load balancer lock per application"

8 files changed, 111 insertions, 145 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 48a53104f7f..7d875434e1c 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
@@ -1,4 +1,4 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision;
 
 import com.google.inject.Inject;
@@ -17,7 +17,6 @@ import com.yahoo.transaction.Mutex;
 import com.yahoo.transaction.NestedTransaction;
 import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.hosted.provision.lb.LoadBalancer;
-import com.yahoo.vespa.hosted.provision.lb.LoadBalancerId;
 import com.yahoo.vespa.hosted.provision.lb.LoadBalancerInstance;
 import com.yahoo.vespa.hosted.provision.lb.LoadBalancerList;
 import com.yahoo.vespa.hosted.provision.maintenance.InfrastructureVersions;
@@ -51,7 +50,6 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.function.BiFunction;
-import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -196,16 +194,7 @@ public class NodeRepository extends AbstractComponent {
 
     /** Returns a filterable list of all load balancers in this repository */
     public LoadBalancerList loadBalancers() {
-        return loadBalancers((ignored) -> true);
-    }
-
-    /** Returns a filterable list of load balancers belonging to given application */
-    public LoadBalancerList loadBalancers(ApplicationId application) {
-        return loadBalancers((id) -> id.application().equals(application));
-    }
-
-    private LoadBalancerList loadBalancers(Predicate<LoadBalancerId> predicate) {
-        return LoadBalancerList.copyOf(db.readLoadBalancers(predicate).values());
+        return LoadBalancerList.copyOf(database().readLoadBalancers().values());
     }
 
     public List<Node> getNodes(ApplicationId id, Node.State ... inState) { return db.getNodes(id, inState); }
@@ -232,7 +221,7 @@ public class NodeRepository extends AbstractComponent {
         candidates.parentOf(node).ifPresent(trustedNodes::add);
         node.allocation().ifPresent(allocation -> {
             trustedNodes.addAll(candidates.owner(allocation.owner()).asList());
-            loadBalancers.asList().stream()
+            loadBalancers.owner(allocation.owner()).asList().stream()
                          .map(LoadBalancer::instance)
                          .map(LoadBalancerInstance::networks)
                          .forEach(trustedNetworks::addAll);
@@ -304,12 +293,7 @@ public class NodeRepository extends AbstractComponent {
      */
     public List<NodeAcl> getNodeAcls(Node node, boolean children) {
         NodeList candidates = list();
-        LoadBalancerList loadBalancers;
-        if (node.allocation().isPresent()) {
-            loadBalancers = loadBalancers(node.allocation().get().owner());
-        } else {
-            loadBalancers = LoadBalancerList.EMPTY;
-        }
+        LoadBalancerList loadBalancers = loadBalancers();
         if (children) {
             return candidates.childrenOf(node).asList().stream()
                              .map(childNode -> getNodeAcl(childNode, candidates, loadBalancers))
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerList.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerList.java
index 479fd328162..014d3df8d9a 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerList.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerList.java
@@ -1,6 +1,9 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.lb;
 
+import com.yahoo.config.provision.ApplicationId;
+
+import java.time.Instant;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
@@ -15,19 +18,27 @@ import java.util.stream.Stream;
  */
 public class LoadBalancerList implements Iterable<LoadBalancer> {
 
-    public static LoadBalancerList EMPTY = new LoadBalancerList(List.of());
-
     private final List<LoadBalancer> loadBalancers;
 
     private LoadBalancerList(Collection<LoadBalancer> loadBalancers) {
         this.loadBalancers = List.copyOf(Objects.requireNonNull(loadBalancers, "loadBalancers must be non-null"));
     }
 
+    /** Returns the subset of load balancers owned by given application */
+    public LoadBalancerList owner(ApplicationId application) {
+        return of(loadBalancers.stream().filter(lb -> lb.id().application().equals(application)));
+    }
+
     /** Returns the subset of load balancers that are in given state */
     public LoadBalancerList in(LoadBalancer.State state) {
         return of(loadBalancers.stream().filter(lb -> lb.state() == state));
     }
 
+    /** Returns the subset of load balancers that last changed before given instant */
+    public LoadBalancerList changedBefore(Instant instant) {
+        return of(loadBalancers.stream().filter(lb -> lb.changedAt().isBefore(instant)));
+    }
+
     public List<LoadBalancer> asList() {
         return loadBalancers;
     }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java
index 0f363991310..bcb0c901f14 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java
@@ -15,8 +15,6 @@ import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Objects;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 /**
@@ -53,31 +51,37 @@ public class LoadBalancerExpirer extends Maintainer {
 
     /** Move reserved load balancer that have expired to inactive */
     private void expireReserved() {
-        var now = nodeRepository().clock().instant();
-        withLoadBalancersIn(State.reserved, lb -> {
-            var gracePeriod = now.minus(reservedExpiry);
-            if (!lb.changedAt().isBefore(gracePeriod)) return; // Should not move to inactive yet
-            db.writeLoadBalancer(lb.with(State.inactive, now));
-        });
+        try (var lock = db.lockLoadBalancers()) {
+            var now = nodeRepository().clock().instant();
+            var expirationTime = now.minus(reservedExpiry);
+            var expired = nodeRepository().loadBalancers()
+                                          .in(State.reserved)
+                                          .changedBefore(expirationTime);
+            expired.forEach(lb -> db.writeLoadBalancer(lb.with(State.inactive, now)));
+        }
     }
 
     /** Deprovision inactive load balancers that have expired */
     private void removeInactive() {
         var failed = new ArrayList<LoadBalancerId>();
-        var lastException = new AtomicReference<Exception>();
-        var now = nodeRepository().clock().instant();
-        withLoadBalancersIn(State.inactive, lb -> {
-            var gracePeriod = now.minus(inactiveExpiry);
-            if (!lb.changedAt().isBefore(gracePeriod)) return; // Should not be removed yet
-            if (!allocatedNodes(lb.id()).isEmpty()) return;    // Still has nodes, do not remove
-            try {
-                service.remove(lb.id().application(), lb.id().cluster());
-                db.removeLoadBalancer(lb.id());
-            } catch (Exception e){
-                failed.add(lb.id());
-                lastException.set(e);
+        Exception lastException = null;
+        try (var lock = db.lockLoadBalancers()) {
+            var now = nodeRepository().clock().instant();
+            var expirationTime = now.minus(inactiveExpiry);
+            var expired = nodeRepository().loadBalancers()
+                                          .in(State.inactive)
+                                          .changedBefore(expirationTime);
+            for (var lb : expired) {
+                if (!allocatedNodes(lb.id()).isEmpty()) continue; // Defer removal if there are still nodes allocated to application
+                try {
+                    service.remove(lb.id().application(), lb.id().cluster());
+                    db.removeLoadBalancer(lb.id());
+                } catch (Exception e) {
+                    failed.add(lb.id());
+                    lastException = e;
+                }
             }
-        });
+        }
         if (!failed.isEmpty()) {
             log.log(LogLevel.WARNING, String.format("Failed to remove %d load balancers: %s, retrying in %s",
                                                     failed.size(),
@@ -85,27 +89,30 @@ public class LoadBalancerExpirer extends Maintainer {
                                                           .map(LoadBalancerId::serializedForm)
                                                           .collect(Collectors.joining(", ")),
                                                     interval()),
-                    lastException.get());
+                    lastException);
         }
     }
 
     /** Remove reals from inactive load balancers */
     private void pruneReals() {
         var failed = new ArrayList<LoadBalancerId>();
-        var lastException = new AtomicReference<Exception>();
-        withLoadBalancersIn(State.inactive, lb -> {
-            var allocatedNodes = allocatedNodes(lb.id()).stream().map(Node::hostname).collect(Collectors.toSet());
-            var reals = new LinkedHashSet<>(lb.instance().reals());
-            // Remove any real no longer allocated to this application
-            reals.removeIf(real -> !allocatedNodes.contains(real.hostname().value()));
-            try {
-                service.create(lb.id().application(), lb.id().cluster(), reals, true);
-                db.writeLoadBalancer(lb.with(lb.instance().withReals(reals)));
-            } catch (Exception e) {
-                failed.add(lb.id());
-                lastException.set(e);
+        Exception lastException = null;
+        try (var lock = db.lockLoadBalancers()) {
+            var deactivated = nodeRepository().loadBalancers().in(State.inactive);
+            for (var lb : deactivated) {
+                var allocatedNodes = allocatedNodes(lb.id()).stream().map(Node::hostname).collect(Collectors.toSet());
+                var reals = new LinkedHashSet<>(lb.instance().reals());
+                // Remove any real no longer allocated to this application
+                reals.removeIf(real -> !allocatedNodes.contains(real.hostname().value()));
+                try {
+                    service.create(lb.id().application(), lb.id().cluster(), reals, true);
+                    db.writeLoadBalancer(lb.with(lb.instance().withReals(reals)));
+                } catch (Exception e) {
+                    failed.add(lb.id());
+                    lastException = e;
+                }
             }
-        });
+        }
         if (!failed.isEmpty()) {
             log.log(LogLevel.WARNING, String.format("Failed to remove reals from %d load balancers: %s, retrying in %s",
                                                     failed.size(),
@@ -113,21 +120,7 @@ public class LoadBalancerExpirer extends Maintainer {
                                                           .map(LoadBalancerId::serializedForm)
                                                           .collect(Collectors.joining(", ")),
                                                     interval()),
-                    lastException.get());
-        }
-    }
-
-    /** Apply operation to all load balancers that exist in given state, while holding lock */
-    private void withLoadBalancersIn(LoadBalancer.State state, Consumer<LoadBalancer> operation) {
-        try (var legacyLock = db.lockLoadBalancers()) {
-            for (var id : db.readLoadBalancerIds()) {
-                try (var lock = db.lockLoadBalancers(id.application())) {
-                    var loadBalancer = db.readLoadBalancer(id);
-                    if (loadBalancer.isEmpty()) continue;              // Load balancer was removed during loop
-                    if (loadBalancer.get().state() != state) continue; // Wrong state
-                    operation.accept(loadBalancer.get());
-                }
-            }
+                    lastException);
         }
     }
 
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 0a8575578ce..ce0bcc2e337 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
@@ -1,4 +1,4 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. 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.google.common.util.concurrent.UncheckedTimeoutException;
@@ -38,7 +38,6 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.function.Function;
-import java.util.function.Predicate;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -484,16 +483,18 @@ public class CuratorDatabaseClient {
 
     // Load balancers
     public List<LoadBalancerId> readLoadBalancerIds() {
-        return readLoadBalancerIds((ignored) -> true);
+        return curatorDatabase.getChildren(loadBalancersRoot).stream()
+                              .map(LoadBalancerId::fromSerializedForm)
+                              .collect(Collectors.toUnmodifiableList());
     }
 
-    public Map<LoadBalancerId, LoadBalancer> readLoadBalancers(Predicate<LoadBalancerId> filter) {
-        return readLoadBalancerIds(filter).stream()
-                                          .map(this::readLoadBalancer)
-                                          .filter(Optional::isPresent)
-                                          .map(Optional::get)
-                                          .collect(collectingAndThen(toMap(LoadBalancer::id, Function.identity()),
-                                                                     Collections::unmodifiableMap));
+    public Map<LoadBalancerId, LoadBalancer> readLoadBalancers() {
+        return readLoadBalancerIds().stream()
+                                    .map(this::readLoadBalancer)
+                                    .filter(Optional::isPresent)
+                                    .map(Optional::get)
+                                    .collect(collectingAndThen(toMap(LoadBalancer::id, Function.identity()),
+                                                           Collections::unmodifiableMap));
     }
 
     public Optional<LoadBalancer> readLoadBalancer(LoadBalancerId id) {
@@ -521,26 +522,14 @@ public class CuratorDatabaseClient {
         transaction.commit();
     }
 
-    // TODO(mpolden): Remove this and all usages once migration to per-application lock is complete
     public Lock lockLoadBalancers() {
         return lock(lockRoot.append("loadBalancersLock"), defaultLockTimeout);
     }
 
-    public Lock lockLoadBalancers(ApplicationId application) {
-        return lock(lockRoot.append("loadBalancersLock2").append(application.serializedForm()), defaultLockTimeout);
-    }
-
     private Path loadBalancerPath(LoadBalancerId id) {
         return loadBalancersRoot.append(id.serializedForm());
     }
 
-    private List<LoadBalancerId> readLoadBalancerIds(Predicate<LoadBalancerId> predicate) {
-        return curatorDatabase.getChildren(loadBalancersRoot).stream()
-                              .map(LoadBalancerId::fromSerializedForm)
-                              .filter(predicate)
-                              .collect(Collectors.toUnmodifiableList());
-    }
-
     private Transaction.Operation createOrSet(Path path, byte[] data) {
         if (curatorDatabase.exists(path)) {
             return CuratorOperations.setData(path.getAbsolute(), data);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java
index d4d5b46dfdf..1500154aa07 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java
@@ -1,4 +1,4 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. 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.ApplicationId;
@@ -51,12 +51,10 @@ public class LoadBalancerProvisioner {
         this.db = nodeRepository.database();
         this.service = service;
         // Read and write all load balancers to make sure they are stored in the latest version of the serialization format
-        try (var legacyLock = db.lockLoadBalancers()) {
+        try (var lock = db.lockLoadBalancers()) {
             for (var id : db.readLoadBalancerIds()) {
-                try (var lock = db.lockLoadBalancers(id.application())) {
-                    var loadBalancer = db.readLoadBalancer(id);
-                    loadBalancer.ifPresent(db::writeLoadBalancer);
-                }
+                var loadBalancer = db.readLoadBalancer(id);
+                loadBalancer.ifPresent(db::writeLoadBalancer);
             }
         }
     }
@@ -75,10 +73,8 @@ public class LoadBalancerProvisioner {
         if (requestedNodes.type() != NodeType.tenant) return; // Nothing to provision for this node type
         if (!cluster.type().isContainer()) return; // Nothing to provision for this cluster type
         if (application.instance().isTester()) return; // Do not provision for tester instances
-        try (var legacyLock = db.lockLoadBalancers()) {
-            try (var lock = db.lockLoadBalancers(application)) {
-                provision(application, cluster.id(), false, lock);
-            }
+        try (var loadBalancersLock = db.lockLoadBalancers()) {
+            provision(application, cluster.id(), false, loadBalancersLock);
         }
     }
 
@@ -94,17 +90,15 @@ public class LoadBalancerProvisioner {
      */
     public void activate(ApplicationId application, Set<ClusterSpec> clusters,
                          @SuppressWarnings("unused") Mutex applicationLock, NestedTransaction transaction) {
-        try (var legacyLock = db.lockLoadBalancers()) {
-            try (var lock = db.lockLoadBalancers(application)) {
-                var containerClusters = containerClusterOf(clusters);
-                for (var clusterId : containerClusters) {
-                    // Provision again to ensure that load balancer instance is re-configured with correct nodes
-                    provision(application, clusterId, true, lock);
-                }
-                // Deactivate any surplus load balancers, i.e. load balancers for clusters that have been removed
-                var surplusLoadBalancers = surplusLoadBalancersOf(application, containerClusters);
-                deactivate(surplusLoadBalancers, transaction);
+        try (var loadBalancersLock = db.lockLoadBalancers()) {
+            var containerClusters = containerClusterOf(clusters);
+            for (var clusterId : containerClusters) {
+                // Provision again to ensure that load balancer instance is re-configured with correct nodes
+                provision(application, clusterId, true, loadBalancersLock);
             }
+            // Deactivate any surplus load balancers, i.e. load balancers for clusters that have been removed
+            var surplusLoadBalancers = surplusLoadBalancersOf(application, containerClusters);
+            deactivate(surplusLoadBalancers, transaction);
         }
     }
 
@@ -114,17 +108,16 @@ public class LoadBalancerProvisioner {
      */
     public void deactivate(ApplicationId application, NestedTransaction transaction) {
         try (var applicationLock = nodeRepository.lock(application)) {
-            try (var legacyLock = db.lockLoadBalancers()) {
-                try (var lock = db.lockLoadBalancers(application)) {
-                    deactivate(nodeRepository.loadBalancers(application).asList(), transaction);
-                }
+            try (Mutex loadBalancersLock = db.lockLoadBalancers()) {
+                deactivate(nodeRepository.loadBalancers().owner(application).asList(), transaction);
             }
         }
     }
 
     /** Returns load balancers of given application that are no longer referenced by given clusters */
     private List<LoadBalancer> surplusLoadBalancersOf(ApplicationId application, Set<ClusterSpec.Id> activeClusters) {
-        var activeLoadBalancersByCluster = nodeRepository.loadBalancers(application)
+        var activeLoadBalancersByCluster = nodeRepository.loadBalancers()
+                                                         .owner(application)
                                                          .in(LoadBalancer.State.active)
                                                          .asList()
                                                          .stream()
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/LoadBalancersResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/LoadBalancersResponse.java
index 3147d4caded..9f8f4a804d1 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/LoadBalancersResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/LoadBalancersResponse.java
@@ -1,4 +1,4 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.restapi.v2;
 
 import com.yahoo.config.provision.ApplicationId;
@@ -37,14 +37,10 @@ public class LoadBalancersResponse extends HttpResponse {
     }
 
     private List<LoadBalancer> loadBalancers() {
-        LoadBalancerList loadBalancers;
-        var application = application();
-        if (application.isPresent()) {
-            loadBalancers = nodeRepository.loadBalancers(application.get());
-        } else {
-            loadBalancers = nodeRepository.loadBalancers();
-        }
-        return loadBalancers.asList();
+        LoadBalancerList loadBalancers = nodeRepository.loadBalancers();
+        return application().map(loadBalancers::owner)
+                            .map(LoadBalancerList::asList)
+                            .orElseGet(loadBalancers::asList);
     }
 
     @Override
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirerTest.java
index 4b7534b431e..12b48fd7a35 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirerTest.java
@@ -1,4 +1,4 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.maintenance;
 
 import com.yahoo.component.Vtag;
@@ -32,14 +32,14 @@ import static org.junit.Assert.assertTrue;
  */
 public class LoadBalancerExpirerTest {
 
-    private final ProvisioningTester tester = new ProvisioningTester.Builder().build();
+    private ProvisioningTester tester = new ProvisioningTester.Builder().build();
 
     @Test
     public void expire_inactive() {
         LoadBalancerExpirer expirer = new LoadBalancerExpirer(tester.nodeRepository(),
                                                               Duration.ofDays(1),
                                                               tester.loadBalancerService());
-        Supplier<Map<LoadBalancerId, LoadBalancer>> loadBalancers = () -> tester.nodeRepository().database().readLoadBalancers((ignored) -> true);
+        Supplier<Map<LoadBalancerId, LoadBalancer>> loadBalancers = () -> tester.nodeRepository().database().readLoadBalancers();
 
         // Deploy two applications with a total of three load balancers
         ClusterSpec.Id cluster1 = ClusterSpec.Id.from("qrs");
@@ -67,11 +67,10 @@ public class LoadBalancerExpirerTest {
         // Expirer prunes reals before expiration time of load balancer itself
         expirer.maintain();
         assertEquals(Set.of(), tester.loadBalancerService().instances().get(lb1).reals());
-        assertEquals(Set.of(), loadBalancers.get().get(lb1).instance().reals());
+        assertEquals(Set.of(), tester.nodeRepository().loadBalancers().owner(lb1.application()).asList().get(0).instance().reals());
 
         // Expirer defers removal of load balancer until expiration time passes
         expirer.maintain();
-        assertSame(LoadBalancer.State.inactive, loadBalancers.get().get(lb1).state());
         assertTrue("Inactive load balancer not removed", tester.loadBalancerService().instances().containsKey(lb1));
 
         // Expirer removes load balancers once expiration time passes
@@ -86,7 +85,7 @@ public class LoadBalancerExpirerTest {
         // A single cluster is removed
         deployApplication(app2, cluster1);
         expirer.maintain();
-        assertSame(LoadBalancer.State.inactive, loadBalancers.get().get(lb3).state());
+        assertEquals(LoadBalancer.State.inactive, loadBalancers.get().get(lb3).state());
 
         // Expirer defers removal while nodes are still allocated to cluster
         expirer.maintain();
@@ -104,7 +103,7 @@ public class LoadBalancerExpirerTest {
         LoadBalancerExpirer expirer = new LoadBalancerExpirer(tester.nodeRepository(),
                                                               Duration.ofDays(1),
                                                               tester.loadBalancerService());
-        Supplier<Map<LoadBalancerId, LoadBalancer>> loadBalancers = () -> tester.nodeRepository().database().readLoadBalancers((ignored) -> true);
+        Supplier<Map<LoadBalancerId, LoadBalancer>> loadBalancers = () -> tester.nodeRepository().database().readLoadBalancers();
 
 
         // Prepare application
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java
index ee9a582c4db..a26e802dfe8 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java
@@ -1,4 +1,4 @@
-// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 Yahoo Holdings. 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.google.common.collect.Iterators;
@@ -40,14 +40,15 @@ public class LoadBalancerProvisionerTest {
 
     private final ApplicationId app1 = ApplicationId.from("tenant1", "application1", "default");
     private final ApplicationId app2 = ApplicationId.from("tenant2", "application2", "default");
+
     private final ApplicationId infraApp1 = ApplicationId.from("vespa", "tenant-host", "default");
 
-    private final ProvisioningTester tester = new ProvisioningTester.Builder().build();
+    private ProvisioningTester tester = new ProvisioningTester.Builder().build();
 
     @Test
     public void provision_load_balancer() {
-        Supplier<List<LoadBalancer>> lbApp1 = () -> tester.nodeRepository().loadBalancers(app1).asList();
-        Supplier<List<LoadBalancer>> lbApp2 = () -> tester.nodeRepository().loadBalancers(app2).asList();
+        Supplier<List<LoadBalancer>> lbApp1 = () -> tester.nodeRepository().loadBalancers().owner(app1).asList();
+        Supplier<List<LoadBalancer>> lbApp2 = () -> tester.nodeRepository().loadBalancers().owner(app2).asList();
         ClusterSpec.Id containerCluster1 = ClusterSpec.Id.from("qrs1");
         ClusterSpec.Id contentCluster = ClusterSpec.Id.from("content");
 
@@ -79,7 +80,7 @@ public class LoadBalancerProvisionerTest {
         tester.activate(app1, prepare(app1,
                                       clusterRequest(ClusterSpec.Type.container, containerCluster1),
                                       clusterRequest(ClusterSpec.Type.content, contentCluster)));
-        LoadBalancer loadBalancer = tester.nodeRepository().loadBalancers(app1).asList().get(0);
+        LoadBalancer loadBalancer = tester.nodeRepository().loadBalancers().owner(app1).asList().get(0);
         assertEquals(2, loadBalancer.instance().reals().size());
         assertTrue("Failed node is removed", loadBalancer.instance().reals().stream()
                                                          .map(Real::hostname)
@@ -157,7 +158,7 @@ public class LoadBalancerProvisionerTest {
         var nodes = prepare(app1, Capacity.fromCount(2, new NodeResources(1, 4, 10, 0.3), false, true),
                                            true,
                                            clusterRequest(ClusterSpec.Type.container, ClusterSpec.Id.from("qrs")));
-        Supplier<LoadBalancer> lb = () -> tester.nodeRepository().loadBalancers(app1).asList().get(0);
+        Supplier<LoadBalancer> lb = () -> tester.nodeRepository().loadBalancers().owner(app1).asList().get(0);
         assertTrue("Load balancer provisioned with empty reals", tester.loadBalancerService().instances().get(lb.get().id()).reals().isEmpty());
         assignIps(tester.nodeRepository().getNodes(app1));
         tester.activate(app1, nodes);
@@ -188,7 +189,7 @@ public class LoadBalancerProvisionerTest {
                                            clusterRequest(ClusterSpec.Type.container,
                                                           ClusterSpec.Id.from("tenant-host"))));
         assertTrue("No load balancer provisioned", tester.loadBalancerService().instances().isEmpty());
-        assertEquals(List.of(), tester.nodeRepository().loadBalancers(infraApp1).asList());
+        assertEquals(List.of(), tester.nodeRepository().loadBalancers().owner(infraApp1).asList());
     }
 
     @Test
@@ -196,12 +197,12 @@ public class LoadBalancerProvisionerTest {
         tester.activate(app1, prepare(app1, clusterRequest(ClusterSpec.Type.content,
                                                            ClusterSpec.Id.from("tenant-host"))));
         assertTrue("No load balancer provisioned", tester.loadBalancerService().instances().isEmpty());
-        assertEquals(List.of(), tester.nodeRepository().loadBalancers(app1).asList());
+        assertEquals(List.of(), tester.nodeRepository().loadBalancers().owner(app1).asList());
     }
 
     @Test
     public void provision_load_balancer_combined_cluster() {
-        Supplier<List<LoadBalancer>> lbs = () -> tester.nodeRepository().loadBalancers(app1).asList();
+        Supplier<List<LoadBalancer>> lbs = () -> tester.nodeRepository().loadBalancers().owner(app1).asList();
         ClusterSpec.Id cluster = ClusterSpec.Id.from("foo");
 
         var nodes = prepare(app1, clusterRequest(ClusterSpec.Type.combined, cluster));