From c42fb7d499cfc91ed5f728e1221bfe3bc893c0be Mon Sep 17 00:00:00 2001 From: YongXing Date: Sun, 18 Feb 2024 16:33:16 +0800 Subject: [PATCH] [#2127] Improvement(core): improve TreeLock for ease of unlock and fix unexpected err when lock failed (#2141) ### What changes were proposed in this pull request? - improve TreeLock for ease of unlock by keeping lock type - fix unexpected err when lock failed and unlock all nodes that have been locked ### Why are the changes needed? Fix: #2127 ### Does this PR introduce _any_ user-facing change? N/A ### How was this patch tested? UT --- .../datastrato/gravitino/lock/TreeLock.java | 47 ++++++---- .../gravitino/lock/TestTreeLock.java | 87 ++++++++++++++++--- 2 files changed, 106 insertions(+), 28 deletions(-) diff --git a/core/src/main/java/com/datastrato/gravitino/lock/TreeLock.java b/core/src/main/java/com/datastrato/gravitino/lock/TreeLock.java index 14f7ceb9a77..ec8076ff752 100644 --- a/core/src/main/java/com/datastrato/gravitino/lock/TreeLock.java +++ b/core/src/main/java/com/datastrato/gravitino/lock/TreeLock.java @@ -6,9 +6,10 @@ package com.datastrato.gravitino.lock; import com.datastrato.gravitino.NameIdentifier; -import java.util.ArrayDeque; import java.util.Deque; import java.util.List; +import java.util.concurrent.ConcurrentLinkedDeque; +import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -63,8 +64,8 @@ public class TreeLock { // TreeLockNode to be locked private final List lockNodes; - // TreeLockNode that has been locked. - private final Deque heldLocks = new ArrayDeque<>(); + // TreeLockNode that has been locked along with its lock type. + private final Deque> heldLocks = new ConcurrentLinkedDeque<>(); private LockType lockType; TreeLock(List lockNodes, NameIdentifier identifier) { @@ -73,7 +74,9 @@ public class TreeLock { } /** - * Lock the tree lock with the given lock type. + * Lock the tree lock with the given lock type. This method locks all nodes in the list, from the + * root to the leaf, and pushes them onto the deque. If an exception occurs during the locking + * process, it will unlock all nodes that have been locked so far. * * @param lockType The lock type to lock the tree lock. */ @@ -84,8 +87,23 @@ public void lock(LockType lockType) { for (int i = 0; i < length; i++) { TreeLockNode treeLockNode = lockNodes.get(i); LockType type = i == length - 1 ? lockType : LockType.READ; - treeLockNode.lock(type); - heldLocks.push(treeLockNode); + try { + treeLockNode.lock(type); + heldLocks.push(Pair.of(treeLockNode, type)); + if (LOG.isTraceEnabled()) { + LOG.trace("Locked node: {}, lock type: {}", treeLockNode, type); + } + } catch (Exception e) { + LOG.error( + "Failed to lock the treeNode, identifier: {}, node {} of lockNodes: [{}]", + identifier, + treeLockNode, + lockNodes, + e); + // unlock all nodes that have been locked when an exception occurs. + unlock(); + throw e; + } } if (LOG.isTraceEnabled()) { @@ -103,22 +121,19 @@ public void unlock() { throw new IllegalStateException("We must lock the tree lock before unlock it."); } - boolean lastNode = false; while (!heldLocks.isEmpty()) { - LockType type = LockType.READ; - if (!lastNode) { - lastNode = true; - type = lockType; - } - - TreeLockNode current = heldLocks.pop(); - // Unlock the node and decrease the reference count. + Pair pair = heldLocks.pop(); + TreeLockNode current = pair.getLeft(); + LockType type = pair.getRight(); current.unlock(type); + if (LOG.isTraceEnabled()) { + LOG.trace("Unlocked node: {}, lock type: {}", current, type); + } } if (LOG.isTraceEnabled()) { LOG.trace( - "Unlocked the tree lock, ident: {}, lockNodes: [{}], lock type: {}", + "Unlocked the tree lock, identifier: {}, lockNodes: [{}], lock type: {}", identifier, lockNodes, lockType); diff --git a/core/src/test/java/com/datastrato/gravitino/lock/TestTreeLock.java b/core/src/test/java/com/datastrato/gravitino/lock/TestTreeLock.java index 235afd3e859..3fcbb5bcebc 100644 --- a/core/src/test/java/com/datastrato/gravitino/lock/TestTreeLock.java +++ b/core/src/test/java/com/datastrato/gravitino/lock/TestTreeLock.java @@ -6,23 +6,86 @@ package com.datastrato.gravitino.lock; import static com.datastrato.gravitino.lock.TestLockManager.getConfig; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.doThrow; +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; class TestTreeLock { + private LockManager lockManager; + private TreeLock lock; + + @BeforeEach + void setUp() { + lockManager = new LockManager(getConfig()); + lock = lockManager.createTreeLock(TestLockManager.randomNameIdentifier()); + } + + @Test + void testLockAndUnlockWithReadLock() { + assertDoesNotThrow( + () -> { + lock.lock(LockType.READ); + lock.unlock(); + }); + } @Test - void testLock() { - LockManager lockManager = new LockManager(getConfig()); - - for (int i = 0; i < 1000; i++) { - TreeLock lock = lockManager.createTreeLock(TestLockManager.randomNameIdentifier()); - lock.lock(i % 2 == 0 ? LockType.READ : LockType.WRITE); - try { - // Business logic here. - } finally { - lock.unlock(); - } - } + void testLockAndUnlockWithWriteLock() { + assertDoesNotThrow( + () -> { + lock.lock(LockType.WRITE); + lock.unlock(); + }); + } + + @Test + void testUnlockWithoutLock() { + assertThrows( + IllegalStateException.class, + () -> { + lock.unlock(); + }); + } + + @Test + void testMultipleLockAndUnlock() { + assertDoesNotThrow( + () -> { + for (int i = 0; i < 1000; i++) { + lock.lock(i % 2 == 0 ? LockType.READ : LockType.WRITE); + lock.unlock(); + } + }); + } + + @Test + void testLockFailureAndUnlock() { + TreeLockNode mockNode1 = Mockito.mock(TreeLockNode.class); + TreeLockNode mockNode2 = Mockito.mock(TreeLockNode.class); + TreeLockNode mockNode3 = Mockito.mock(TreeLockNode.class); + + // Mock the lock method of the second node to throw an exception + doThrow(new RuntimeException("Mock exception")).when(mockNode2).lock(Mockito.any()); + + List lockNodes = Arrays.asList(mockNode1, mockNode2, mockNode3); + TreeLock treeLock = new TreeLock(lockNodes, TestLockManager.randomNameIdentifier()); + + assertThrows( + RuntimeException.class, + () -> treeLock.lock(LockType.WRITE), + "Expected lock to throw, but it didn't"); + + // Verify that the first node was unlocked + Mockito.verify(mockNode1, Mockito.times(1)).unlock(Mockito.any()); + + // Verify that the second and third nodes were not unlocked + Mockito.verify(mockNode2, Mockito.never()).unlock(Mockito.any()); + Mockito.verify(mockNode3, Mockito.never()).unlock(Mockito.any()); } }