lock per key fixes

Signed-off-by: Peter Alfonsi <[email protected]>
sgup432 · Jun 4, 2024 · 6eb3e05 · 6eb3e05
1 parent 810de9a
commit 6eb3e05
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/...les/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java b/...les/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java
@@ -171,9 +171,11 @@ ICache<K, V> getDiskCache() {
     private void writeLock(ICacheKey<K> key) {
         if (threadLocal.get() != null) {
             LockWrapper lockWrapper = threadLocal.get();
-            lockWrapper.writeLock.close();
-            if (lockWrapper.refCount.decrementAndGet() == 0) {
-                locks.remove(key);
+            if (lockWrapper != null && lockWrapper.writeLock.isHeldByCurrentThread()) {
+                lockWrapper.writeLock.close();
+                if (lockWrapper.refCount.decrementAndGet() == 0) {
+                    locks.remove(key);
+                }
             }
         }
         LockWrapper lockWrapper = locks.computeIfAbsent(key, key1 -> {
@@ -195,18 +197,22 @@ private void readLock(ICacheKey<K> key) {
 
     private void unlockWriteLock(ICacheKey<K> key) {
         LockWrapper lockWrapper = locks.get(key);
-        lockWrapper.writeLock.close();
-        if (lockWrapper.refCount.decrementAndGet() == 0) {
-            locks.remove(key);
+        if (lockWrapper != null && lockWrapper.writeLock.isHeldByCurrentThread()) {
+            lockWrapper.writeLock.close();
+            if (lockWrapper.refCount.decrementAndGet() == 0) {
+                locks.remove(key);
+            }
         }
         threadLocal.remove();
     }
 
     private void unlockReadLock(ICacheKey<K> key) {
         LockWrapper lockWrapper = locks.get(key);
-        lockWrapper.readLock.close();
-        if (lockWrapper.refCount.decrementAndGet() == 0) {
-            locks.remove(key);
+        if (lockWrapper != null && lockWrapper.readLock.isHeldByCurrentThread()) {
+            lockWrapper.readLock.close();
+            if (lockWrapper.refCount.decrementAndGet() == 0) {
+                locks.remove(key);
+            }
         }
         threadLocal.remove();
     }

diff --git a/...ache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java b/...ache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java
@@ -1368,7 +1368,9 @@ public void testNumLocksTiming() throws Exception {
                             System.out.println("Finished iter " + j);
                         }
                     }
-                } catch (Exception ignored) {}
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
                 countDownLatch.countDown();
             });
             threads[i].start();