[BugFix] Locker: rollback partial intensive-lock acquisition (#72423)

Signed-off-by: Kevin Cai <kevin.cai@celerdata.com>
(cherry picked from commit f32f70b)

Author: kevincai

fe/fe-core/src/main/java/com/starrocks/common/util/concurrent/lock/Locker.java

@@ -288,18 +288,26 @@ public void lockTablesWithIntensiveDbLock(Long dbId, List<Long> tableList, LockT
         Preconditions.checkState(lockType.equals(LockType.READ) || lockType.equals(LockType.WRITE));
         List<Long> tableListClone = new ArrayList<>(tableList);
         if (Config.lock_manager_enabled) {
+            LockType intentionType = (lockType == LockType.WRITE)
+                    ? LockType.INTENTION_EXCLUSIVE
+                    : LockType.INTENTION_SHARED;
+            boolean dbLockHeld = false;
+            List<Long> ridLockedList = new ArrayList<>();
             try {
-                if (lockType == LockType.WRITE) {
-                    this.lock(dbId, LockType.INTENTION_EXCLUSIVE, 0);
-                } else {
-                    this.lock(dbId, LockType.INTENTION_SHARED, 0);
-                }
+                this.lock(dbId, intentionType, 0);
+                dbLockHeld = true;
 
                 Collections.sort(tableListClone);
                 for (Long rid : tableListClone) {
                     this.lock(rid, lockType, 0);
+                    ridLockedList.add(rid);
                 }
             } catch (LockException e) {
+                // Roll back any locks already acquired so that a partial failure
+                // (e.g. deadlock victim selection on the second / Nth lock call)
+                // does not leave a stranded DB intention or table lock that
+                // would block subsequent DB-WRITE operations indefinitely.
+                rollbackPartialIntensiveLock(dbId, intentionType, dbLockHeld, ridLockedList, lockType);
                 throw ErrorReportException.report(ErrorCode.ERR_LOCK_ERROR, e.getMessage());
             }
         } else {
@@ -308,6 +316,30 @@ public void lockTablesWithIntensiveDbLock(Long dbId, List<Long> tableList, LockT
         }
     }
 
+    /**
+     * Best-effort rollback of a partial intensive-lock acquisition. Releases each
+     * already-acquired lock and swallows any errors raised by individual release
+     * calls so that the original LockException (the actual cause of the rollback)
+     * surfaces to the caller instead of being masked by a release-side error.
+     */
+    private void rollbackPartialIntensiveLock(Long dbId, LockType intentionType,
+                                              boolean dbLockHeld, List<Long> ridLockedList, LockType tableLockType) {
+        for (Long rid : ridLockedList) {
+            try {
+                this.release(rid, tableLockType);
+            } catch (Exception releaseEx) {
+                LOG.warn("Failed to release table lock {} during partial-acquire rollback", rid, releaseEx);
+            }
+        }
+        if (dbLockHeld) {
+            try {
+                this.release(dbId, intentionType);
+            } catch (Exception releaseEx) {
+                LOG.warn("Failed to release DB intention lock {} during partial-acquire rollback", dbId, releaseEx);
+            }
+        }
+    }
+
     /**
      * No need to release lock explicitly, it will be released automatically when the locker failed.
      */
@@ -389,14 +421,19 @@ public void unLockTableWithIntensiveDbLock(Long dbId, Long tableId, LockType loc
     public void lockTableWithIntensiveDbLock(Long dbId, Long tableId, LockType lockType) {
         Preconditions.checkState(lockType.equals(LockType.READ) || lockType.equals(LockType.WRITE));
         if (Config.lock_manager_enabled) {
+            LockType intentionType = (lockType == LockType.WRITE)
+                    ? LockType.INTENTION_EXCLUSIVE
+                    : LockType.INTENTION_SHARED;
+            boolean dbLockHeld = false;
             try {
-                if (lockType == LockType.WRITE) {
-                    this.lock(dbId, LockType.INTENTION_EXCLUSIVE, 0);
-                } else {
-                    this.lock(dbId, LockType.INTENTION_SHARED, 0);
-                }
+                this.lock(dbId, intentionType, 0);
+                dbLockHeld = true;
                 this.lock(tableId, lockType, 0);
             } catch (LockException e) {
+                // Roll back the DB intention lock if step 1 succeeded but the
+                // table-lock acquisition failed (e.g. deadlock victim, interrupt).
+                // Otherwise the caller would leak an IS/IX lock on the DB.
+                rollbackPartialIntensiveLock(dbId, intentionType, dbLockHeld, Collections.emptyList(), lockType);
                 throw ErrorReportException.report(ErrorCode.ERR_LOCK_ERROR, e.getMessage());
             }
         } else {

fe/fe-core/src/test/java/com/starrocks/common/lock/TestLockInterface.java

@@ -17,13 +17,16 @@
 import com.google.common.collect.Lists;
 import com.starrocks.catalog.Database;
 import com.starrocks.common.Config;
+import com.starrocks.common.ErrorReportException;
 import com.starrocks.common.util.concurrent.QueryableReentrantReadWriteLock;
 import com.starrocks.common.util.concurrent.lock.LockException;
 import com.starrocks.common.util.concurrent.lock.LockManager;
 import com.starrocks.common.util.concurrent.lock.LockParams;
 import com.starrocks.common.util.concurrent.lock.LockType;
 import com.starrocks.common.util.concurrent.lock.Locker;
+import com.starrocks.common.util.concurrent.lock.NotSupportLockException;
 import com.starrocks.server.GlobalStateMgr;
+import mockit.Invocation;
 import mockit.Mock;
 import mockit.MockUp;
 import org.junit.jupiter.api.AfterEach;
@@ -224,6 +227,89 @@ public void testTryLockTablesWithIntensiveDbLock2() throws LockException {
         Assertions.assertTrue(lockManager.isOwner(rid2, locker, LockType.WRITE));
     }
 
+    /**
+     * Regression: lockTableWithIntensiveDbLock acquires a DB intention lock first
+     * and a table lock second. If the table-lock acquisition throws (deadlock
+     * victim, interrupt, etc.) the helper must release the already-held DB
+     * intention lock; otherwise it leaks an IS / IX lock and blocks subsequent
+     * DB-WRITE operations until the FE restarts.
+     */
+    @Test
+    public void testLockTableWithIntensiveDbLockRollsBackDbLockOnTableLockFailure() {
+        long dbId = 100L;
+        long tableId = 101L;
+        Database database = new Database(dbId, "db_rollback_single");
+
+        // Force the second LockManager.lock call (the table-lock acquisition) to fail.
+        // The first call (DB intention) is allowed to proceed normally so the rollback
+        // path actually has something to release.
+        final int[] callCount = {0};
+        new MockUp<LockManager>() {
+            @Mock
+            public void lock(Invocation inv, long rid, Locker locker, LockType lockType, long timeout)
+                    throws LockException {
+                callCount[0]++;
+                if (callCount[0] == 2) {
+                    throw new NotSupportLockException("simulated table-lock failure");
+                }
+                inv.proceed();
+            }
+        };
+
+        Locker locker = new Locker();
+        Assertions.assertThrows(ErrorReportException.class, () ->
+                locker.lockTableWithIntensiveDbLock(database.getId(), tableId, LockType.WRITE));
+
+        LockManager lockManager = GlobalStateMgr.getCurrentState().getLockManager();
+        Assertions.assertFalse(lockManager.isOwner(database.getId(), locker, LockType.INTENTION_EXCLUSIVE),
+                "DB intention lock must be released after partial-acquire rollback");
+        Assertions.assertFalse(lockManager.isOwner(tableId, locker, LockType.WRITE),
+                "Table lock must not be held (acquisition failed)");
+    }
+
+    /**
+     * Regression: lockTablesWithIntensiveDbLock loops over multiple table ids.
+     * If the Nth table-lock acquisition fails after earlier ones succeeded, the
+     * helper must release the DB intention lock AND every already-acquired table
+     * lock; otherwise it strands all of them.
+     */
+    @Test
+    public void testLockTablesWithIntensiveDbLockRollsBackAllAcquiredLocksOnNthTableLockFailure() {
+        long dbId = 200L;
+        long tableId1 = 201L;
+        long tableId2 = 202L;
+        Database database = new Database(dbId, "db_rollback_multi");
+
+        // Call sequence: 1 = DB IX, 2 = tableId1 (succeeds), 3 = tableId2 (fails).
+        // The helper sorts the table list, and 201 < 202, so this ordering is
+        // deterministic.
+        final int[] callCount = {0};
+        new MockUp<LockManager>() {
+            @Mock
+            public void lock(Invocation inv, long rid, Locker locker, LockType lockType, long timeout)
+                    throws LockException {
+                callCount[0]++;
+                if (callCount[0] == 3) {
+                    throw new NotSupportLockException("simulated table-lock failure");
+                }
+                inv.proceed();
+            }
+        };
+
+        Locker locker = new Locker();
+        Assertions.assertThrows(ErrorReportException.class, () ->
+                locker.lockTablesWithIntensiveDbLock(database.getId(),
+                        Lists.newArrayList(tableId1, tableId2), LockType.WRITE));
+
+        LockManager lockManager = GlobalStateMgr.getCurrentState().getLockManager();
+        Assertions.assertFalse(lockManager.isOwner(database.getId(), locker, LockType.INTENTION_EXCLUSIVE),
+                "DB intention lock must be released after partial-acquire rollback");
+        Assertions.assertFalse(lockManager.isOwner(tableId1, locker, LockType.WRITE),
+                "First table lock (acquired before failure) must be released by rollback");
+        Assertions.assertFalse(lockManager.isOwner(tableId2, locker, LockType.WRITE),
+                "Second table lock (acquisition failed) must not be held");
+    }
+
     @Test
     public void testTryLockTablesWithIntensiveDbLock3() throws LockException {
         long rid = 1L;