Sync shutdown, invalidate doom on no singletons, fix unit test

3 files changed, 36 insertions, 19 deletions

diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/CuratorWrapper.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/CuratorWrapper.java
index 31c476b33ae..2cba6ee6efe 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/CuratorWrapper.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/CuratorWrapper.java
@@ -5,6 +5,7 @@ import com.yahoo.component.annotation.Inject;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.path.Path;
 import com.yahoo.vespa.curator.api.VespaCurator;
+import com.yahoo.yolean.UncheckedInterruptedException;
 import org.apache.curator.framework.state.ConnectionState;
 import org.apache.zookeeper.KeeperException.BadVersionException;
 import org.apache.zookeeper.data.Stat;
@@ -14,6 +15,8 @@ import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Implementation of {@link VespaCurator} which delegates to a {@link Curator}.
@@ -122,7 +125,15 @@ public class CuratorWrapper extends AbstractComponent implements VespaCurator {
 
     @Override
     public void deconstruct() {
-        singletons.close();
+        try {
+            singletons.shutdown().get();
+        }
+        catch (InterruptedException e) {
+            throw new UncheckedInterruptedException(e, true);
+        }
+        catch (ExecutionException e) {
+            throw new RuntimeException(e.getCause());
+        }
     }
 
 }
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/SingletonManager.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/SingletonManager.java
index d116b57f506..96511a5dd1c 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/SingletonManager.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/SingletonManager.java
@@ -29,7 +29,7 @@ import java.util.logging.Logger;
  *
  * @author jonmv
  */
-class SingletonManager implements AutoCloseable {
+class SingletonManager {
 
     private static final Logger logger = Logger.getLogger(SingletonManager.class.getName());
 
@@ -91,13 +91,16 @@ class SingletonManager implements AutoCloseable {
         for (Janitor janitor : janitors.values()) janitor.invalidate();
     }
 
-    @Override
-    public synchronized void close() {
+    public synchronized CompletableFuture<?> shutdown() {
+        CompletableFuture<?>[] futures = new CompletableFuture[registrations.size()];
+        int i = 0;
         for (SingletonWorker singleton : List.copyOf(registrations.keySet())) {
             String id = registrations.get(singleton);
             logger.log(Level.WARNING, singleton + " still registered with id '" + id + "' at shutdown");
-            unregister(singleton);
+            futures[i++] = unregister(singleton);
         }
+        return CompletableFuture.allOf(futures)
+                                .orTimeout(10, TimeUnit.SECONDS);
     }
 
 
@@ -220,6 +223,7 @@ class SingletonManager implements AutoCloseable {
                     if (e == null) e = f;
                     else e.addSuppressed(f);
                 }
+                if (singletons.isEmpty()) doom.set(INVALID);
                 if (e != null) throw e;
             }
         }
@@ -243,9 +247,7 @@ class SingletonManager implements AutoCloseable {
                     else e.addSuppressed(f);
                 }
             }
-            if (singletons.isEmpty()) {
-                unlock();
-            }
+            if (singletons.isEmpty()) doom.set(INVALID);
             if (e != null) throw e;
         }
 
@@ -254,7 +256,7 @@ class SingletonManager implements AutoCloseable {
          * If lock is held, or acquired, ping the ZK cluster to extend our deadline.
          */
         private void renewLease() {
-            if (doom.get() == INVALID || singletons.isEmpty()) {
+            if (doom.get() == INVALID) {
                 doom.set(null);
                 return; // Skip to updateStatus, deactivation, and release the lock.
             }
@@ -287,7 +289,7 @@ class SingletonManager implements AutoCloseable {
          */
         private void updateStatus() {
             Instant ourDoom = doom.get();
-            boolean shouldBeActive = ourDoom != null && ! clock.instant().isAfter(ourDoom);
+            boolean shouldBeActive = ourDoom != null && ourDoom != INVALID && ! clock.instant().isAfter(ourDoom);
             if ( ! active && shouldBeActive) {
                 try {
                     active = true;
diff --git a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorWrapperTest.java b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorWrapperTest.java
index 1ff639ed75a..82b9693b8fe 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorWrapperTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorWrapperTest.java
@@ -19,7 +19,6 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.stream.Collectors;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.assertEquals;
@@ -132,7 +131,7 @@ public class CuratorWrapperTest {
                           metric);
 
             // Singleton is reactivated next tick.
-            singleton.phaser.arriveAndAwaitAdvance();
+            singleton.phaser.awaitAdvance(singleton.phaser.arriveAndDeregister());
             assertTrue(singleton.isActive);
             verifyMetrics(Map.of("activation.count", 3.0,
                                  "activation.millis", 0.0,
@@ -144,7 +143,6 @@ public class CuratorWrapperTest {
 
             // Manager unregisters remaining singletons on shutdown.
             curator.deconstruct();
-            singleton.phaser.arriveAndAwaitAdvance();
             assertFalse(singleton.isActive);
             verifyMetrics(Map.of("activation.count", 3.0,
                                  "activation.millis", 0.0,
@@ -239,6 +237,7 @@ public class CuratorWrapperTest {
                           metric);
 
             newSingleton.shutdown();
+            curator.deconstruct();
             verifyMetrics(Map.of("activation.count", 6.0,
                                  "activation.millis", 0.0,
                                  "activation.failure.count", 1.0,
@@ -247,8 +246,6 @@ public class CuratorWrapperTest {
                                  "is_active", 0.0,
                                  "has_lease", 0.0),
                           metric);
-
-            curator.deconstruct();
         }
     }
 
@@ -313,6 +310,7 @@ public class CuratorWrapperTest {
 
             singleton.phaser.arriveAndDeregister();
             singleton.shutdown();
+            curator.deconstruct();
             assertFalse(singleton.isActive);
             verifyMetrics(Map.of("activation.count", 2.0,
                                  "activation.millis", 0.0,
@@ -321,8 +319,6 @@ public class CuratorWrapperTest {
                                  "is_active", 0.0,
                                  "has_lease", 0.0),
                           metric);
-
-            curator.deconstruct();
         }
     }
 
@@ -331,8 +327,16 @@ public class CuratorWrapperTest {
         boolean isActive;
         Phaser phaser = new Phaser(1);
         @Override public String id() { return "singleton"; } // ... lest anonymous subclasses get different IDs ... ƪ(`▿▿▿▿´ƪ)
-        @Override public void activate() { isActive = true; phaser.arriveAndAwaitAdvance(); }
-        @Override public void deactivate() { isActive = false; phaser.arriveAndAwaitAdvance(); }
+        @Override public void activate() {
+            if (isActive) throw new IllegalStateException("already active");
+            isActive = true;
+            phaser.arriveAndAwaitAdvance();
+        }
+        @Override public void deactivate() {
+            if ( ! isActive) throw new IllegalStateException("already inactive");
+            isActive = false;
+            phaser.arriveAndAwaitAdvance();
+        }
         public void shutdown() { unregister(Duration.ofSeconds(2)); }
     }