[server] Eliminate global lock for KV tablet creation/deletion/load

fluss-server/src/main/java/org/apache/fluss/server/TabletManagerBase.java

@@ -51,8 +51,6 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -80,8 +78,6 @@ public enum TabletType {
 
     protected final Configuration conf;
 
-    protected final Lock tabletCreationOrDeletionLock = new ReentrantLock();
-
     // TODO make this parameter configurable.
     private final int recoveryThreads;
     private final TabletType tabletType;

fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java

@@ -49,6 +49,7 @@
 import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocatorUtil;
 import org.apache.fluss.utils.FileUtils;
 import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.function.SupplierWithException;
 import org.apache.fluss.utils.types.Tuple2;
 
 import org.rocksdb.RateLimiter;
@@ -69,8 +70,6 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
-import static org.apache.fluss.utils.concurrent.LockUtils.inLock;
-
 /**
  * The entry point to the fluss kv management subsystem. The kv manager is responsible for kv tablet
  * creation, retrieval, and cleaning. All read and write operations to kv tablet are delegated to
@@ -122,6 +121,8 @@ public static RateLimiter getDefaultRateLimiter() {
 
     private final Map<TableBucket, KvTablet> currentKvs = new ConcurrentHashMap<>();
 
+    private final Map<TableBucket, Object> kvLocks = new ConcurrentHashMap<>();
+
     /**
      * For arrow log format. The buffer allocator to allocate memory for arrow write batch of
      * changelog records.
@@ -285,8 +286,8 @@ public KvTablet getOrCreateKv(
             ArrowCompressionInfo arrowCompressionInfo,
             @Nullable Runnable flushCompleteListener)
             throws Exception {
-        return inLock(
-                tabletCreationOrDeletionLock,
+        return inKvLock(
+                tableBucket,
                 () -> {
                     if (currentKvs.containsKey(tableBucket)) {
                         return currentKvs.get(tableBucket);
@@ -356,9 +357,20 @@ public Optional<KvTablet> getKv(TableBucket tableBucket) {
     }
 
     public void dropKv(TableBucket tableBucket) {
-        KvTablet dropKvTablet =
-                inLock(tabletCreationOrDeletionLock, () -> currentKvs.remove(tableBucket));
+        inKvLock(
+                tableBucket,
+                () -> {
+                    try {
+                        doDropKv(tableBucket);
+                    } finally {
+                        kvLocks.remove(tableBucket);
+                    }
+                    return null;
+                });
+    }
 
+    private void doDropKv(TableBucket tableBucket) {
+        KvTablet dropKvTablet = currentKvs.remove(tableBucket);
         if (dropKvTablet != null) {
             TablePath tablePath = dropKvTablet.getTablePath();
             try {
@@ -397,6 +409,38 @@ public KvTablet loadKv(
         Tuple2<PhysicalTablePath, TableBucket> pathAndBucket = FlussPaths.parseTabletDir(tabletDir);
         PhysicalTablePath physicalTablePath = pathAndBucket.f0;
         TableBucket tableBucket = pathAndBucket.f1;
+        return inKvLock(
+                tableBucket,
+                () ->
+                        doLoadKv(
+                                tabletDir,
+                                physicalTablePath,
+                                tableBucket,
+                                schemaGetter,
+                                flushCompleteListener));
+    }
+
+    private KvTablet doLoadKv(
+            File tabletDir,
+            PhysicalTablePath physicalTablePath,
+            TableBucket tableBucket,
+            SchemaGetter schemaGetter,
+            @Nullable Runnable flushCompleteListener)
+            throws Exception {
+        KvTablet currentKv = currentKvs.get(tableBucket);
+        if (currentKv != null) {
+            throw new IllegalStateException(
+                    String.format(
+                            "Duplicate kv tablet directories for bucket %s are found in both %s and %s. "
+                                    + "Recover server from this "
+                                    + "failure by manually deleting one of the two kv directories for this bucket. "
+                                    + "It is recommended to delete the bucket in the kv tablet directory that is "
+                                    + "known to have failed recently.",
+                            tableBucket,
+                            tabletDir.getAbsolutePath(),
+                            currentKv.getKvTabletDir().getAbsolutePath()));
+        }
+
         // get the log tablet for the kv tablet
         LogTablet logTablet =
                 logManager
@@ -443,19 +487,7 @@ public KvTablet loadKv(
                         kvFlushScheduler,
                         flushCompleteListener,
                         autoIncrementManager);
-        if (this.currentKvs.containsKey(tableBucket)) {
-            throw new IllegalStateException(
-                    String.format(
-                            "Duplicate kv tablet directories for bucket %s are found in both %s and %s. "
-                                    + "Recover server from this "
-                                    + "failure by manually deleting one of the two kv directories for this bucket. "
-                                    + "It is recommended to delete the bucket in the kv tablet directory that is "
-                                    + "known to have failed recently.",
-                            tableBucket,
-                            tabletDir.getAbsolutePath(),
-                            currentKvs.get(tableBucket).getKvTabletDir().getAbsolutePath()));
-        }
-        this.currentKvs.put(tableBucket, kvTablet);
+        currentKvs.put(tableBucket, kvTablet);
 
         return kvTablet;
     }
@@ -527,4 +559,17 @@ public void reconfigure(Configuration newConfig) throws ConfigException {
                     "Failed to reconfigure shared RocksDB rate limiter: " + e.getMessage(), e);
         }
     }
+
+    private <T, E extends Exception> T inKvLock(
+            TableBucket tableBucket, SupplierWithException<T, E> action) throws E {
+        while (true) {
+            Object lock = kvLocks.computeIfAbsent(tableBucket, ignored -> new Object());
+            synchronized (lock) {
+                if (kvLocks.get(tableBucket) != lock) {
+                    continue;
+                }
+                return action.get();
+            }
+        }
+    }
 }

fluss-server/src/test/java/org/apache/fluss/server/kv/KvManagerTest.java

@@ -67,6 +67,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -355,6 +356,27 @@ void testDropKv(String partitionName) throws Exception {
         assertThat(kvManager.getKv(tableBucket1)).isPresent();
     }
 
+    @Test
+    void testKvRunConcurrentlyForDifferentBuckets() throws Exception {
+        initTableBuckets(null);
+        BlockingSchemaGetter blockingSchemaGetter = new BlockingSchemaGetter();
+        ExecutorService executor = Executors.newFixedThreadPool(2);
+        Future<KvTablet> blockedCreation =
+                executor.submit(
+                        () -> getOrCreateKv(tablePath1, null, tableBucket1, blockingSchemaGetter));
+        try {
+            blockingSchemaGetter.awaitBlocked();
+
+            Future<KvTablet> otherBucketCreation =
+                    executor.submit(() -> getOrCreateKv(tablePath2, null, tableBucket2));
+            assertThat(otherBucketCreation.get(10, TimeUnit.SECONDS)).isNotNull();
+        } finally {
+            blockingSchemaGetter.unblock();
+            blockedCreation.get(10, TimeUnit.SECONDS);
+            executor.shutdownNow();
+        }
+    }
+
     @Test
     void testGetNonExistentKv() {
         initTableBuckets(null);
@@ -545,4 +567,35 @@ private void verifyMultiGet(KvTablet kvTablet, byte[] key, byte[] expectedValue)
         List<byte[]> gotValues = kvTablet.multiGet(Collections.singletonList(key));
         assertThat(gotValues).containsExactly(expectedValue);
     }
+
+    private static final class BlockingSchemaGetter extends TestingSchemaGetter {
+        private final CountDownLatch blocked = new CountDownLatch(1);
+        private final CountDownLatch unblock = new CountDownLatch(1);
+
+        private BlockingSchemaGetter() {
+            super(new SchemaInfo(DATA1_SCHEMA_PK, 1));
+        }
+
+        @Override
+        public SchemaInfo getLatestSchemaInfo() {
+            blocked.countDown();
+            try {
+                if (!unblock.await(30, TimeUnit.SECONDS)) {
+                    throw new FlussRuntimeException("Timed out waiting to unblock schema lookup.");
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new FlussRuntimeException("Interrupted while blocking schema lookup.", e);
+            }
+            return super.getLatestSchemaInfo();
+        }
+
+        private void awaitBlocked() throws InterruptedException {
+            assertThat(blocked.await(10, TimeUnit.SECONDS)).isTrue();
+        }
+
+        private void unblock() {
+            unblock.countDown();
+        }
+    }
 }