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

2 files changed, 23 insertions, 10 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;