Keep track of locks in OrchestratorContext

5 files changed, 138 insertions, 47 deletions

diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java
index 356d645fca4..17dfb924973 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java
@@ -3,13 +3,13 @@ package com.yahoo.vespa.orchestrator;
 
 import com.yahoo.log.LogLevel;
 import com.yahoo.time.TimeBudget;
+import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientTimeouts;
 
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashMap;
 import java.util.Optional;
 import java.util.logging.Logger;
 
@@ -25,33 +25,39 @@ public class OrchestratorContext implements AutoCloseable {
     private static final Duration DEFAULT_TIMEOUT_FOR_BATCH_OP = Duration.ofSeconds(60);
     private static final Duration TIMEOUT_OVERHEAD = Duration.ofMillis(500);
 
+    private final Optional<OrchestratorContext> parent;
     private final Clock clock;
-    private final boolean partOfMultiAppOp;
     private final TimeBudget timeBudget;
     private final boolean probe;
     private final boolean largeLocks;
-    private final List<Runnable> onClose = new ArrayList<>();
+
+    // The key set is the set of applications locked by this context tree: Only the
+    // root context has a non-empty set. The value is an unlock callback to be called
+    // when the root context is closed.
+    private final HashMap<ApplicationInstanceReference, Runnable> locks = new HashMap<>();
 
     /** Create an OrchestratorContext for operations on multiple applications. */
     public static OrchestratorContext createContextForMultiAppOp(Clock clock, boolean largeLocks) {
-        return new OrchestratorContext(clock, true, TimeBudget.fromNow(clock, DEFAULT_TIMEOUT_FOR_BATCH_OP), false, largeLocks);
+        return new OrchestratorContext(null, clock, TimeBudget.fromNow(clock, DEFAULT_TIMEOUT_FOR_BATCH_OP), false, largeLocks);
     }
 
     /** Create an OrchestratorContext for an operation on a single application. */
     public static OrchestratorContext createContextForSingleAppOp(Clock clock) {
-        return new OrchestratorContext(clock, false, TimeBudget.fromNow(clock, DEFAULT_TIMEOUT_FOR_SINGLE_OP), false, false);
+        return new OrchestratorContext(null, clock, TimeBudget.fromNow(clock, DEFAULT_TIMEOUT_FOR_SINGLE_OP), false, false);
     }
 
-    private OrchestratorContext(Clock clock, boolean partOfMultiAppOp, TimeBudget timeBudget, boolean probe, boolean largeLocks) {
+    private OrchestratorContext(OrchestratorContext parentOrNull,
+                                Clock clock,
+                                TimeBudget timeBudget,
+                                boolean probe,
+                                boolean largeLocks) {
+        this.parent = Optional.ofNullable(parentOrNull);
         this.clock = clock;
-        this.partOfMultiAppOp = partOfMultiAppOp;
         this.timeBudget = timeBudget;
         this.probe = probe;
         this.largeLocks = largeLocks;
     }
 
-    public boolean partOfMultiAppOp() { return partOfMultiAppOp; }
-
     public Duration getTimeLeft() {
         return timeBudget.timeLeftOrThrow().get();
     }
@@ -68,12 +74,44 @@ public class OrchestratorContext implements AutoCloseable {
     /** Whether application locks acquired during probing of a batch suspend should be closed after the non-probe is done. */
     public boolean largeLocks() { return largeLocks; }
 
+    /**
+     * Returns true if 1. large locks is enabled, and 2.
+     * {@link #registerLockAcquisition(ApplicationInstanceReference, Runnable) registerLockAcquisition}
+     * has been invoked on any context below the root context that returned true.
+     */
+    public boolean hasLock(ApplicationInstanceReference application) {
+        return parent.map(p -> p.hasLock(application)).orElseGet(() -> locks.containsKey(application));
+    }
+
+    /**
+     * Returns true if large locks is enabled in the root context, and in case the unlock callback
+     * will be invoked when the root context is closed.
+     */
+    public boolean registerLockAcquisition(ApplicationInstanceReference application, Runnable unlock) {
+        if (parent.isPresent()) {
+            return parent.get().registerLockAcquisition(application, unlock);
+        }
+
+        if (!largeLocks) {
+            return false;
+        }
+
+        if (locks.containsKey(application)) {
+            unlock.run();
+            throw new IllegalStateException("Application " + application + " was already associated with a lock");
+        }
+
+        locks.put(application, unlock);
+
+        return true;
+    }
+
     /** Create OrchestratorContext to use within an application lock. */
     public OrchestratorContext createSubcontextWithinLock() {
         // Move deadline towards past by a fixed amount to ensure there's time to process exceptions and
         // access ZooKeeper before the lock times out.
         TimeBudget subTimeBudget = timeBudget.withDeadline(timeBudget.deadline().get().minus(TIMEOUT_OVERHEAD));
-        return new OrchestratorContext(clock, partOfMultiAppOp, subTimeBudget, probe, false);
+        return new OrchestratorContext(this, clock, subTimeBudget, probe, largeLocks);
     }
 
     /** Create an OrchestratorContext for an operation on a single application, but limited to current timeout. */
@@ -85,25 +123,18 @@ public class OrchestratorContext implements AutoCloseable {
             deadline = maxDeadline;
         }
 
-        return new OrchestratorContext(
-                clock,
-                partOfMultiAppOp,
-                TimeBudget.from(clock, now, Optional.of(Duration.between(now, deadline))),
-                probe,
-                false);
+        TimeBudget timeBudget = TimeBudget.from(clock, now, Optional.of(Duration.between(now, deadline)));
+        return new OrchestratorContext(this, clock, timeBudget, probe, largeLocks);
     }
 
-    public void runOnClose(Runnable runnable) { onClose.add(runnable); }
-
     @Override
     public void close() {
-        int i = onClose.size();
-        while (i --> 0) {
+        locks.forEach((application, unlock) -> {
             try {
-                onClose.get(i).run();
+                unlock.run();
             } catch (RuntimeException e) {
                 logger.log(LogLevel.ERROR, "Failed run on close : " + e.getMessage());
             }
-        }
+        });
     }
 }
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java
index 4bfd2b3867e..68862859615 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService.java
@@ -137,16 +137,15 @@ public class ZookeeperStatusService implements StatusService {
         // then a non-probe. With "large locks", the lock is not release in between -
         // no lock is taken on the non-probe. Instead, the release is done on the multi-application
         // context close.
-        if (!context.largeLocks() || !context.partOfMultiAppOp() || context.isProbe()) {
+        if (context.hasLock(applicationInstanceReference)) {
+            onRegistryClose = () -> {};
+        } else {
             Runnable unlock = acquireLock(context, applicationInstanceReference);
-            if (context.largeLocks() && context.isProbe()) {
-                context.runOnClose(unlock);
+            if (context.registerLockAcquisition(applicationInstanceReference, unlock)) {
                 onRegistryClose = () -> {};
             } else {
                 onRegistryClose = unlock;
             }
-        } else {
-            onRegistryClose = () -> {};
         }
 
         try {
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorContextTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorContextTest.java
new file mode 100644
index 00000000000..607894ee104
--- /dev/null
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorContextTest.java
@@ -0,0 +1,59 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.orchestrator;
+
+import com.yahoo.test.ManualClock;
+import com.yahoo.vespa.applicationmodel.ApplicationInstanceId;
+import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
+import com.yahoo.vespa.applicationmodel.TenantId;
+import org.junit.Test;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author hakonhall
+ */
+public class OrchestratorContextTest {
+    private final ApplicationInstanceReference application = new ApplicationInstanceReference(
+            new TenantId("tenant"),
+            new ApplicationInstanceId("app:dev:us-east-1:default"));
+
+    @Test
+    public void testLargeLocks() {
+        var mutable = new Object() { boolean locked = true; };
+        Runnable unlock = () -> mutable.locked = false;
+
+        try (OrchestratorContext rootContext = OrchestratorContext.createContextForMultiAppOp(new ManualClock(), true)) {
+            try (OrchestratorContext probeContext = rootContext.createSubcontextForSingleAppOp(true)) {
+                assertFalse(probeContext.hasLock(application));
+                assertTrue(probeContext.registerLockAcquisition(application, unlock));
+
+                assertTrue(probeContext.hasLock(application));
+                assertTrue(mutable.locked);
+            }
+
+            try (OrchestratorContext nonProbeContext = rootContext.createSubcontextForSingleAppOp(false)) {
+                assertTrue(nonProbeContext.hasLock(application));
+                assertTrue(mutable.locked);
+            }
+
+            assertTrue(mutable.locked);
+        }
+        assertFalse(mutable.locked);
+    }
+
+    @Test
+    public void testLargeLocksDisabled() {
+        var mutable = new Object() { boolean locked = true; };
+        Runnable unlock = () -> mutable.locked = false;
+
+        try (OrchestratorContext rootContext = OrchestratorContext.createContextForMultiAppOp(new ManualClock(), false)) {
+            try (OrchestratorContext probeContext = rootContext.createSubcontextForSingleAppOp(true)) {
+                assertFalse(probeContext.hasLock(application));
+                assertFalse(probeContext.registerLockAcquisition(application, unlock));
+            }
+        }
+
+        assertTrue(mutable.locked);
+    }
+}
\ No newline at end of file
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
index 73e479d290b..a2c99b86ae2 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
@@ -16,8 +16,8 @@ import com.yahoo.vespa.applicationmodel.ServiceInstance;
 import com.yahoo.vespa.applicationmodel.ServiceStatus;
 import com.yahoo.vespa.applicationmodel.ServiceType;
 import com.yahoo.vespa.applicationmodel.TenantId;
-import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.curator.mock.MockCurator;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactory;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactoryMock;
@@ -27,21 +27,16 @@ import com.yahoo.vespa.orchestrator.policy.BatchHostStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.policy.HostedVespaClusterPolicy;
 import com.yahoo.vespa.orchestrator.policy.HostedVespaPolicy;
-import com.yahoo.vespa.orchestrator.status.ApplicationInstanceStatus;
 import com.yahoo.vespa.orchestrator.status.HostStatus;
 import com.yahoo.vespa.orchestrator.status.MutableStatusRegistry;
 import com.yahoo.vespa.orchestrator.status.StatusService;
 import com.yahoo.vespa.orchestrator.status.ZookeeperStatusService;
 import com.yahoo.vespa.service.monitor.ServiceModel;
-import org.apache.curator.framework.recipes.atomic.AtomicValue;
-import org.apache.curator.framework.recipes.atomic.DistributedAtomicLong;
-import org.apache.curator.framework.recipes.locks.InterProcessLock;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.InOrder;
 
-import java.time.Clock;
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
@@ -64,7 +59,6 @@ import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -331,6 +325,8 @@ public class OrchestratorImplTest {
 
     @Test
     public void testLargeLocks() throws Exception {
+        flagSource.withBooleanFlag(Flags.ENABLE_LARGE_ORCHESTRATOR_LOCKS.id(), true);
+
         var tenantId = new TenantId("tenant");
         var applicationInstanceId = new ApplicationInstanceId("app:dev:us-east-1:default");
         var applicationInstanceReference = new ApplicationInstanceReference(tenantId, applicationInstanceId);
@@ -371,9 +367,9 @@ public class OrchestratorImplTest {
         // First invocation is probe, second is not.
         assertEquals(2, contexts.size());
         assertTrue(contexts.get(0).isProbe());
-        assertTrue(contexts.get(0).partOfMultiAppOp());
+        assertTrue(contexts.get(0).largeLocks());
         assertFalse(contexts.get(1).isProbe());
-        assertTrue(contexts.get(1).partOfMultiAppOp());
+        assertTrue(contexts.get(1).largeLocks());
 
         verify(applicationApiFactory, times(2)).create(any(), any(), any());
         verify(policy, times(2)).grantSuspensionRequest(any(), any());
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService2Test.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService2Test.java
index fc644970fa5..8f530f4abf3 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService2Test.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusService2Test.java
@@ -42,8 +42,8 @@ public class ZookeeperStatusService2Test {
     @Test
     public void verifyLocks() throws Exception {
         when(context.isProbe()).thenReturn(true);
-        when(context.largeLocks()).thenReturn(false);
-        when(context.partOfMultiAppOp()).thenReturn(true);
+        when(context.hasLock(any())).thenReturn(false);
+        when(context.registerLockAcquisition(any(), any())).thenReturn(false);
 
         when(curator.createMutex(any())).thenReturn(mutex);
         when(mutex.acquire(anyLong(), any())).thenReturn(true);
@@ -57,7 +57,8 @@ public class ZookeeperStatusService2Test {
         verify(curator, times(1)).createMutex(any());
         verify(mutex, times(1)).acquire(anyLong(), any());
         verify(mutex, times(1)).release();
-        verify(context, times(0)).runOnClose(any());
+        verify(context, times(1)).hasLock(any());
+        verify(context, times(1)).registerLockAcquisition(any(), any());
         verifyNoMoreInteractions(mutex);
 
         // Now the non-probe suspension
@@ -71,15 +72,16 @@ public class ZookeeperStatusService2Test {
         verify(mutex, times(2)).acquire(anyLong(), any());
         verify(mutex, times(2)).release();
         ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
-        verify(context, times(0)).runOnClose(runnableCaptor.capture());
+        verify(context, times(2)).hasLock(any());
+        verify(context, times(2)).registerLockAcquisition(any(), any());
         verifyNoMoreInteractions(mutex);
     }
 
     @Test
     public void verifyLargeLocks() throws Exception {
         when(context.isProbe()).thenReturn(true);
-        when(context.largeLocks()).thenReturn(true);
-        when(context.partOfMultiAppOp()).thenReturn(true);
+        when(context.hasLock(any())).thenReturn(false);
+        when(context.registerLockAcquisition(any(), any())).thenReturn(true);
 
         when(curator.createMutex(any())).thenReturn(mutex);
         when(mutex.acquire(anyLong(), any())).thenReturn(true);
@@ -93,12 +95,15 @@ public class ZookeeperStatusService2Test {
         verify(curator, times(1)).createMutex(any());
         verify(mutex, times(1)).acquire(anyLong(), any());
         verify(mutex, times(0)).release();
-        verify(context, times(1)).runOnClose(any());
+        verify(context, times(1)).hasLock(any());
+        verify(context, times(1)).registerLockAcquisition(any(), any());
         verifyNoMoreInteractions(mutex);
 
         // Now the non-probe suspension
 
         when(context.isProbe()).thenReturn(false);
+        when(context.hasLock(any())).thenReturn(true);
+        when(context.registerLockAcquisition(any(), any())).thenReturn(false);
 
         try (MutableStatusRegistry registry = zookeeperStatusService.lockApplicationInstance_forCurrentThreadOnly(context, reference)) {
             // nothing
@@ -107,15 +112,16 @@ public class ZookeeperStatusService2Test {
         // No (additional) acquire, and no releases.
         verify(mutex, times(1)).acquire(anyLong(), any());
         verify(mutex, times(0)).release();
-        ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
-        verify(context, times(1)).runOnClose(runnableCaptor.capture());
+        verify(context, times(2)).hasLock(any());
+        verify(context, times(1)).registerLockAcquisition(any(), any());
         verifyNoMoreInteractions(mutex);
 
         // Verify the context runnable releases the mutex
+        ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(context, times(1)).registerLockAcquisition(any(), runnableCaptor.capture());
         assertEquals(1, runnableCaptor.getAllValues().size());
         runnableCaptor.getAllValues().forEach(Runnable::run);
         verify(mutex, times(1)).acquire(anyLong(), any());
         verify(mutex, times(1)).release();
-
     }
 }
\ No newline at end of file