Let Curator own re-entrant locks

3 files changed, 18 insertions, 1 deletions

diff --git a/zkfacade/abi-spec.json b/zkfacade/abi-spec.json
index 05fb985dbaf..f4ad1ab4372 100644
--- a/zkfacade/abi-spec.json
+++ b/zkfacade/abi-spec.json
@@ -85,6 +85,7 @@
       "public java.util.List getChildren(com.yahoo.path.Path)",
       "public java.util.Optional getData(com.yahoo.path.Path)",
       "public java.util.Optional getStat(com.yahoo.path.Path)",
+      "public com.yahoo.vespa.curator.Lock lock(com.yahoo.path.Path, java.time.Duration)",
       "public org.apache.curator.framework.CuratorFramework framework()",
       "public void close()",
       "public java.lang.String zooKeeperEnsembleConnectionSpec()",
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
index 989584c43e4..6d4f6beece1 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
@@ -32,6 +32,7 @@ import java.time.Duration;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.function.Function;
 import java.util.logging.Logger;
@@ -62,6 +63,9 @@ public class Curator implements AutoCloseable {
     private final String zooKeeperEnsembleConnectionSpec;
     private final int zooKeeperEnsembleCount;
 
+    // All lock keys, to allow re-entrancy. This will grow forever, but this should be too slow to be a problem
+    private final ConcurrentHashMap<Path, Lock> locks = new ConcurrentHashMap<>();
+
     /** Creates a curator instance from a comma-separated string of ZooKeeper host:port strings */
     public static Curator create(String connectionSpec) {
         return new Curator(connectionSpec, connectionSpec, Optional.of(ZK_CLIENT_CONFIG_FILE));
@@ -347,6 +351,14 @@ public class Curator implements AutoCloseable {
         }
     }
 
+    /** Create and acquire a re-entrant lock in given path */
+    public Lock lock(Path path, Duration timeout) {
+        create(path);
+        Lock lock = locks.computeIfAbsent(path, (pathArg) -> new Lock(pathArg.getAbsolute(), this));
+        lock.acquire(timeout);
+        return lock;
+    }
+
     /** Returns the curator framework API */
     public CuratorFramework framework() {
         return curatorFramework;
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/Lock.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/Lock.java
index 2e554d39e44..d97d8f5ed71 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Lock.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Lock.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.curator;
 
 import com.google.common.util.concurrent.UncheckedTimeoutException;
+import com.yahoo.path.Path;
 import com.yahoo.transaction.Mutex;
 import org.apache.curator.framework.recipes.locks.InterProcessLock;
 
@@ -10,7 +11,10 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 
 /**
- * A cluster-wide re-entrant mutex which is released on (the last symmetric) close
+ * A cluster-wide re-entrant mutex which is released on (the last symmetric) close.
+ *
+ * Re-entrancy is limited to the instance of this. To ensure re-entrancy callers should access the lock through
+ * {@link Curator#lock(Path, Duration)} instead of constructing this directly.
  *
  * @author bratseth
  */