1 files changed, 23 insertions, 19 deletions
diff --git a/vespajlib/src/main/java/com/yahoo/concurrent/classlock/ClassLocking.java b/vespajlib/src/main/java/com/yahoo/concurrent/classlock/ClassLocking.java
index 8014148cc30..61e6850ee2e 100644
--- a/vespajlib/src/main/java/com/yahoo/concurrent/classlock/ClassLocking.java
+++ b/vespajlib/src/main/java/com/yahoo/concurrent/classlock/ClassLocking.java
@@ -2,36 +2,40 @@ package com.yahoo.concurrent.classlock;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * @author valerijf
  */
 public class ClassLocking {
     private final Map<Class<?>, ClassLock> classLocks = new HashMap<>();
-    private final Object monitor = new Object();
 
-    public ClassLock lock(Class<?> clazz) {
-        synchronized (monitor) {
-            while(classLocks.containsKey(clazz)) {
-                try {
-                    monitor.wait();
-                } catch (InterruptedException ignored) { }
-            }
+    public synchronized ClassLock lock(Class<?> clazz) {
+        return tryLock(clazz, () -> true);
+    }
 
-            ClassLock classLock = new ClassLock(this, clazz);
-            classLocks.put(clazz, classLock);
-            return classLock;
+    public synchronized ClassLock tryLock(Class<?> clazz, Supplier<Boolean> continueRetrying) {
+        while(classLocks.containsKey(clazz)) {
+            try {
+                wait();
+            } catch (InterruptedException ignored) {
+                if (! continueRetrying.get()) {
+                    throw new LockInterruptException();
+                }
+            }
         }
+
+        ClassLock classLock = new ClassLock(this, clazz);
+        classLocks.put(clazz, classLock);
+        return classLock;
     }
 
-    void unlock(Class<?> clazz, ClassLock classLock) {
-        synchronized (monitor) {
-            if (classLock.equals(classLocks.get(clazz))) {
-                classLocks.remove(clazz);
-                monitor.notifyAll();
-            } else {
-                throw new IllegalArgumentException("Lock has already been released");
-            }
+    synchronized void unlock(Class<?> clazz, ClassLock classLock) {
+        if (classLock.equals(classLocks.get(clazz))) {
+            classLocks.remove(clazz);
+            notifyAll();
+        } else {
+            throw new IllegalArgumentException("Lock has already been released");
         }
     }
 }