[client] Fix NPE thrown by MetadataUpdater when getting bucket count (apache#2054)

Author: swuferhong

fluss-client/src/main/java/org/apache/fluss/client/write/RoundRobinBucketAssigner.java

@@ -31,10 +31,12 @@
 @Internal
 public class RoundRobinBucketAssigner extends DynamicBucketAssigner {
     private final PhysicalTablePath physicalTablePath;
+    private final int bucketNumber;
     private final AtomicInteger counter = new AtomicInteger(new Random().nextInt());
 
-    public RoundRobinBucketAssigner(PhysicalTablePath physicalTablePath) {
+    public RoundRobinBucketAssigner(PhysicalTablePath physicalTablePath, int bucketNumber) {
         this.physicalTablePath = physicalTablePath;
+        this.bucketNumber = bucketNumber;
     }
 
     @Override
@@ -47,8 +49,7 @@ public int assignBucket(Cluster cluster) {
             return bucketsForTable.get(bucket).getBucketId();
         } else {
             // no buckets are available, give a non-available bucket.
-            return MathUtils.toPositive(nextValue)
-                    % cluster.getBucketCount(physicalTablePath.getTablePath());
+            return MathUtils.toPositive(nextValue) % bucketNumber;
         }
     }
 

fluss-client/src/main/java/org/apache/fluss/client/write/StickyBucketAssigner.java

@@ -35,10 +35,12 @@
 public class StickyBucketAssigner extends DynamicBucketAssigner {
 
     private final PhysicalTablePath physicalTablePath;
+    private final int bucketNumber;
     private final AtomicInteger currentBucketId;
 
-    public StickyBucketAssigner(PhysicalTablePath physicalTablePath) {
+    public StickyBucketAssigner(PhysicalTablePath physicalTablePath, int bucketNumber) {
         this.physicalTablePath = physicalTablePath;
+        this.bucketNumber = bucketNumber;
         this.currentBucketId = new AtomicInteger(-1);
     }
 
@@ -73,7 +75,7 @@ private int nextBucket(Cluster cluster, int preBucketId) {
                     cluster.getAvailableBucketsForPhysicalTablePath(physicalTablePath);
             if (availableBuckets.isEmpty()) {
                 int random = MathUtils.toPositive(ThreadLocalRandom.current().nextInt());
-                newBucket = random % cluster.getBucketCount(physicalTablePath.getTablePath());
+                newBucket = random % bucketNumber;
             } else if (availableBuckets.size() == 1) {
                 newBucket = availableBuckets.get(0).getBucketId();
             } else {

fluss-client/src/main/java/org/apache/fluss/client/write/WriterClient.java

@@ -354,9 +354,9 @@ private BucketAssigner createBucketAssigner(
             ConfigOptions.NoKeyAssigner noKeyAssigner =
                     conf.get(ConfigOptions.CLIENT_WRITER_BUCKET_NO_KEY_ASSIGNER);
             if (noKeyAssigner == ROUND_ROBIN) {
-                return new RoundRobinBucketAssigner(physicalTablePath);
+                return new RoundRobinBucketAssigner(physicalTablePath, bucketNumber);
             } else if (noKeyAssigner == STICKY) {
-                return new StickyBucketAssigner(physicalTablePath);
+                return new StickyBucketAssigner(physicalTablePath, bucketNumber);
             } else {
                 throw new IllegalArgumentException(
                         "Unsupported append only row bucket assigner: " + noKeyAssigner);

fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java

@@ -343,7 +343,8 @@ void testAppendWithStickyBucketAssigner() throws Exception {
         int batchSize = 100;
         IndexedRow row = indexedRow(DATA1_ROW_TYPE, new Object[] {1, "a"});
 
-        StickyBucketAssigner bucketAssigner = new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH);
+        StickyBucketAssigner bucketAssigner =
+                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH, 3);
         RecordAccumulator accum =
                 createTestRecordAccumulator(
                         (int) Duration.ofMinutes(1).toMillis(),

fluss-client/src/test/java/org/apache/fluss/client/write/StickyStaticBucketAssignerTest.java

@@ -65,7 +65,7 @@ void testSticky() {
         // init cluster.
         Cluster cluster = updateCluster(Arrays.asList(bucket1, bucket2, bucket3));
         StickyBucketAssigner stickyBucketAssigner =
-                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH);
+                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH, 3);
         int bucketId = stickyBucketAssigner.assignBucket(cluster);
         assertThat(bucketId >= 0 && bucketId < 3).isTrue();
 
@@ -94,7 +94,7 @@ void testBucketIdShouldNotChange() {
         // init cluster.
         Cluster cluster = updateCluster(Arrays.asList(bucket1, bucket2, bucket3));
         StickyBucketAssigner stickyBucketAssigner =
-                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH);
+                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH, 3);
         int bucketId = stickyBucketAssigner.assignBucket(cluster);
         for (int i = 0; i < 3; i++) {
             if (i != bucketId) {
@@ -111,7 +111,7 @@ void testOnlyOneAvailableBuckets() {
         // init cluster.
         Cluster cluster = updateCluster(Collections.singletonList(bucket1));
         StickyBucketAssigner stickyBucketAssigner =
-                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH);
+                new StickyBucketAssigner(DATA1_PHYSICAL_TABLE_PATH, 3);
         int bucketId = stickyBucketAssigner.assignBucket(cluster);
 
         for (int i = 0; i < 100; i++) {
@@ -136,7 +136,7 @@ void testAvailableBucketsTest() {
         Cluster cluster = updateCluster(allBuckets);
 
         // Assure we never choose bucket 1 for tp1 because it is unavailable.
-        StickyBucketAssigner stickyBucketAssigner = new StickyBucketAssigner(tp1);
+        StickyBucketAssigner stickyBucketAssigner = new StickyBucketAssigner(tp1, 3);
         int bucketForTp1 = stickyBucketAssigner.assignBucket(cluster);
         assertThat(bucketForTp1).isNotEqualTo(1);
         for (int i = 0; i < 100; i++) {
@@ -145,7 +145,7 @@ void testAvailableBucketsTest() {
         }
 
         // Assure we always choose bucket 1 for tp2.
-        stickyBucketAssigner = new StickyBucketAssigner(tp2);
+        stickyBucketAssigner = new StickyBucketAssigner(tp2, 3);
         int bucketForTp2 = stickyBucketAssigner.assignBucket(cluster);
         assertThat(bucketForTp2).isEqualTo(1);
         for (int i = 0; i < 100; i++) {
@@ -154,7 +154,7 @@ void testAvailableBucketsTest() {
         }
 
         // Assure that we can still choose one bucket even if there are no available buckets.
-        stickyBucketAssigner = new StickyBucketAssigner(tp3);
+        stickyBucketAssigner = new StickyBucketAssigner(tp3, 3);
         int bucketForTp3 = stickyBucketAssigner.assignBucket(cluster);
         assertThat(bucketForTp3).isIn(0, 1, 2);
         stickyBucketAssigner.onNewBatch(cluster, bucketForTp3);
@@ -165,7 +165,7 @@ void testAvailableBucketsTest() {
     void testMultiThreadToCallOnNewBatch() {
         Cluster cluster = updateCluster(Arrays.asList(bucket1, bucket2, bucket3));
         StickyBucketAssigner stickyBucketAssigner =
-                new StickyBucketAssigner(PhysicalTablePath.of(DATA1_TABLE_PATH));
+                new StickyBucketAssigner(PhysicalTablePath.of(DATA1_TABLE_PATH), 3);
         int bucketId = stickyBucketAssigner.assignBucket(cluster);
         Queue<Integer> bucketIds = new ConcurrentLinkedQueue<>();
         Thread[] threads = new Thread[100];

fluss-common/src/main/java/org/apache/fluss/cluster/Cluster.java

@@ -173,10 +173,6 @@ public TablePath getTablePathOrElseThrow(long tableId) {
                                                 + " in cluster"));
     }
 
-    public int getBucketCount(TablePath tablePath) {
-        return tableInfoByPath.get(tablePath).getNumBuckets();
-    }
-
     /** Get the bucket location for this table-bucket. */
     public Optional<BucketLocation> getBucketLocation(TableBucket tableBucket) {
         return Optional.ofNullable(availableLocationByBucket.get(tableBucket));