Improve DebugAwareFunctionalLock (#3098)

Author: devinrsmith

engine/updategraph/src/main/java/io/deephaven/engine/updategraph/UpdateGraphLock.java

@@ -8,15 +8,14 @@
 import io.deephaven.internal.log.LoggerFactory;
 import io.deephaven.io.logger.Logger;
 import io.deephaven.util.FunctionalInterfaces;
-import io.deephaven.util.SafeCloseable;
 import io.deephaven.util.locks.AwareFunctionalLock;
 import org.apache.commons.lang3.exception.ExceptionUtils;
 import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
@@ -105,13 +104,8 @@ public static void installInstrumentation(@Nullable final Instrumentation instru
         rwLock = new ReentrantReadWriteLock(true);
         readLock = rwLock.readLock();
         writeLock = rwLock.writeLock();
-        if (allowUnitTestMode) {
-            sharedLock = new DebugAwareFunctionalLock(new SharedLock());
-            exclusiveLock = new DebugAwareFunctionalLock(new ExclusiveLock());
-        } else {
-            sharedLock = new SharedLock();
-            exclusiveLock = new ExclusiveLock();
-        }
+        sharedLock = new SharedLock();
+        exclusiveLock = allowUnitTestMode ? new DebugAwareFunctionalLock(new ExclusiveLock()) : new ExclusiveLock();
         this.allowUnitTestMode = allowUnitTestMode;
     }
 
@@ -306,10 +300,10 @@ public final Condition newCondition() {
 
     // region DebugLock
     class DebugAwareFunctionalLock implements AwareFunctionalLock {
-        private final AwareFunctionalLock delegate;
-        private final Deque<Throwable> lockingContext = new ArrayDeque<>();
+        private final ExclusiveLock delegate;
+        private final Deque<Throwable> lockingContext = new ConcurrentLinkedDeque<>();
 
-        DebugAwareFunctionalLock(AwareFunctionalLock delegate) {
+        DebugAwareFunctionalLock(ExclusiveLock delegate) {
             this.delegate = delegate;
         }
 
@@ -321,19 +315,19 @@ public boolean isHeldByCurrentThread() {
         @Override
         public void lock() {
             delegate.lock();
-            lockingContext.push(new Throwable());
+            pushContext();
         }
 
         @Override
         public void lockInterruptibly() throws InterruptedException {
             delegate.lockInterruptibly();
-            lockingContext.push(new Throwable());
+            pushContext();
         }
 
         @Override
         public boolean tryLock() {
             if (delegate.tryLock()) {
-                lockingContext.push(new Throwable());
+                pushContext();
                 return true;
             }
             return false;
@@ -342,16 +336,16 @@ public boolean tryLock() {
         @Override
         public boolean tryLock(long time, @NotNull TimeUnit unit) throws InterruptedException {
             if (delegate.tryLock(time, unit)) {
-                lockingContext.push(new Throwable());
+                pushContext();
                 return true;
             }
             return false;
         }
 
         @Override
         public void unlock() {
-            delegate.unlock();
             lockingContext.pop();
+            delegate.unlock();
         }
 
         @NotNull
@@ -360,37 +354,20 @@ public Condition newCondition() {
             return delegate.newCondition();
         }
 
-        @Override
-        public <EXCEPTION_TYPE extends Exception> void doLocked(
-                @NotNull FunctionalInterfaces.ThrowingRunnable<EXCEPTION_TYPE> runnable) throws EXCEPTION_TYPE {
-            delegate.doLocked(runnable);
-        }
-
-        @Override
-        public <EXCEPTION_TYPE extends Exception> void doLockedInterruptibly(
-                @NotNull FunctionalInterfaces.ThrowingRunnable<EXCEPTION_TYPE> runnable)
-                throws InterruptedException, EXCEPTION_TYPE {
-            delegate.doLockedInterruptibly(runnable);
-        }
-
-        @Override
-        public <RESULT_TYPE, EXCEPTION_TYPE extends Exception> RESULT_TYPE computeLocked(
-                @NotNull FunctionalInterfaces.ThrowingSupplier<RESULT_TYPE, EXCEPTION_TYPE> supplier)
-                throws EXCEPTION_TYPE {
-            return delegate.computeLocked(supplier);
-        }
-
-        @Override
-        public <RESULT_TYPE, EXCEPTION_TYPE extends Exception> RESULT_TYPE computeLockedInterruptibly(
-                @NotNull FunctionalInterfaces.ThrowingSupplier<RESULT_TYPE, EXCEPTION_TYPE> supplier)
-                throws InterruptedException, EXCEPTION_TYPE {
-            return delegate.computeLockedInterruptibly(supplier);
-        }
-
         String getDebugMessage() {
             final Throwable item = lockingContext.peek();
             return item == null ? "locking context is empty" : ExceptionUtils.getStackTrace(item);
         }
+
+        private void pushContext() {
+            // Implementation must have already acquired lock
+            try {
+                lockingContext.push(new Throwable());
+            } catch (Throwable t) {
+                delegate.unlock();
+                throw t;
+            }
+        }
     }
     // endregion DebugLock
 

engine/updategraph/src/main/java/io/deephaven/engine/updategraph/UpdateGraphProcessor.java

@@ -575,16 +575,16 @@ public void enableUnitTestMode() {
     }
 
     private void assertLockAvailable(@NotNull final String action) {
-        if (!UpdateGraphProcessor.DEFAULT.exclusiveLock().tryLock()) {
+        final AwareFunctionalLock exclusiveLock = UpdateGraphProcessor.DEFAULT.exclusiveLock();
+        if (!exclusiveLock.tryLock()) {
             log.error().append("Lock is held when ").append(action).append(", with previous holder: ")
                     .append(unitTestModeHolder).endl();
             ThreadDump.threadDump(System.err);
-            UpdateGraphLock.DebugAwareFunctionalLock lock =
-                    (UpdateGraphLock.DebugAwareFunctionalLock) UpdateGraphProcessor.DEFAULT.exclusiveLock();
+            UpdateGraphLock.DebugAwareFunctionalLock lock = (UpdateGraphLock.DebugAwareFunctionalLock) exclusiveLock;
             throw new IllegalStateException(
                     "Lock is held when " + action + ", with previous holder: " + lock.getDebugMessage());
         }
-        UpdateGraphProcessor.DEFAULT.exclusiveLock().unlock();
+        exclusiveLock.unlock();
     }
 
     /**