[#2127] Improvement(core): improve TreeLock for ease of unlock and fix unexpected err when lock failed

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<TreeLockNode> lockNodes;
 
-  // TreeLockNode that has been locked.
-  private final Deque<TreeLockNode> heldLocks = new ArrayDeque<>();
+  // TreeLockNode that has been locked along with its lock type.
+  private final Deque<Pair<TreeLockNode, LockType>> heldLocks = new ConcurrentLinkedDeque<>();
   private LockType lockType;
 
   TreeLock(List<TreeLockNode> 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<TreeLockNode, LockType> 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);

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<TreeLockNode> 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());
   }
 }