refactoring BaseKVStoreClient to support custom routing mechanism

base-kv/base-kv-store-client/src/main/java/com/baidu/bifromq/basekv/client/IBaseKVStoreClient.java

@@ -13,6 +13,7 @@
 
 package com.baidu.bifromq.basekv.client;
 
+import com.baidu.bifromq.basekv.proto.Boundary;
 import com.baidu.bifromq.basekv.proto.KVRangeStoreDescriptor;
 import com.baidu.bifromq.basekv.store.proto.BootstrapReply;
 import com.baidu.bifromq.basekv.store.proto.BootstrapRequest;
@@ -32,10 +33,11 @@
 import com.baidu.bifromq.basekv.store.proto.TransferLeadershipRequest;
 import com.baidu.bifromq.baserpc.IConnectable;
 import io.reactivex.rxjava3.core.Observable;
+import java.util.NavigableMap;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
-public interface IBaseKVStoreClient extends IKVRangeRouter, IConnectable {
+public interface IBaseKVStoreClient extends IConnectable {
     static BaseKVStoreClientBuilder newBuilder() {
         return new BaseKVStoreClientBuilder();
     }
@@ -44,6 +46,10 @@ static BaseKVStoreClientBuilder newBuilder() {
 
     Observable<Set<KVRangeStoreDescriptor>> describe();
 
+    Observable<NavigableMap<Boundary, KVRangeSetting>> effectiveRouter();
+
+    NavigableMap<Boundary, KVRangeSetting> latestEffectiveRouter();
+
     CompletableFuture<BootstrapReply> bootstrap(String storeId, BootstrapRequest request);
 
     CompletableFuture<RecoverReply> recover(String storeId, RecoverRequest request);

base-kv/base-kv-store-client/src/main/java/com/baidu/bifromq/basekv/client/KVRangeRouterUtil.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2024. The BifroMQ Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and limitations under the License.
+ */
+
+package com.baidu.bifromq.basekv.client;
+
+import static com.baidu.bifromq.basekv.utils.BoundaryUtil.compareEndKeys;
+import static com.baidu.bifromq.basekv.utils.BoundaryUtil.endKey;
+
+import com.baidu.bifromq.basekv.proto.Boundary;
+import com.baidu.bifromq.basekv.utils.BoundaryUtil;
+import com.google.common.collect.Lists;
+import com.google.protobuf.ByteString;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.Optional;
+
+public class KVRangeRouterUtil {
+
+    public static Optional<KVRangeSetting> findByKey(ByteString key,
+                                                     NavigableMap<Boundary, KVRangeSetting> effectiveRouter) {
+        Map.Entry<Boundary, KVRangeSetting> entry =
+            effectiveRouter.floorEntry(Boundary.newBuilder().setStartKey(key).build());
+        if (entry != null) {
+            KVRangeSetting setting = entry.getValue();
+            if (BoundaryUtil.inRange(key, entry.getKey())) {
+                return Optional.of(setting);
+            }
+        }
+        return Optional.empty();
+    }
+
+    public static List<KVRangeSetting> findByBoundary(Boundary boundary,
+                                                      NavigableMap<Boundary, KVRangeSetting> effectiveRouter) {
+        if (effectiveRouter.isEmpty()) {
+            return Collections.emptyList();
+        }
+        if (!boundary.hasStartKey() && !boundary.hasEndKey()) {
+            return Lists.newArrayList(effectiveRouter.values());
+        }
+        if (!boundary.hasStartKey()) {
+            Boundary boundaryEnd = Boundary.newBuilder()
+                .setStartKey(boundary.getEndKey())
+                .setEndKey(boundary.getEndKey()).build();
+            return Lists.newArrayList(effectiveRouter.headMap(boundaryEnd, false).values());
+        }
+        if (!boundary.hasEndKey()) {
+            Boundary boundaryStart = Boundary.newBuilder()
+                .setStartKey(boundary.getStartKey())
+                .setEndKey(boundary.getStartKey()).build();
+            Boundary floorBoundary = effectiveRouter.floorKey(boundaryStart);
+            return Lists.newArrayList(
+                effectiveRouter.tailMap(floorBoundary,
+                        compareEndKeys(endKey(floorBoundary), boundary.getStartKey()) > 0)
+                    .values());
+        }
+        Boundary boundaryStart = Boundary.newBuilder()
+            .setStartKey(boundary.getStartKey())
+            .setEndKey(boundary.getStartKey()).build();
+        Boundary boundaryEnd = Boundary.newBuilder()
+            .setStartKey(boundary.getEndKey())
+            .setEndKey(boundary.getEndKey()).build();
+        Boundary floorBoundary = effectiveRouter.floorKey(boundaryStart);
+
+        return Lists.newArrayList(
+            effectiveRouter.subMap(floorBoundary, compareEndKeys(endKey(floorBoundary), boundary.getStartKey()) > 0,
+                boundaryEnd, false).values());
+    }
+}

base-kv/base-kv-store-client/src/main/java/com/baidu/bifromq/basekv/client/scheduler/MutationCallScheduler.java

@@ -13,8 +13,10 @@
 
 package com.baidu.bifromq.basekv.client.scheduler;
 
-import com.baidu.bifromq.basekv.client.KVRangeSetting;
+import static com.baidu.bifromq.basekv.client.KVRangeRouterUtil.findByKey;
+
 import com.baidu.bifromq.basekv.client.IBaseKVStoreClient;
+import com.baidu.bifromq.basekv.client.KVRangeSetting;
 import com.baidu.bifromq.basescheduler.BatchCallScheduler;
 import com.google.protobuf.ByteString;
 import java.time.Duration;
@@ -45,7 +47,7 @@ public MutationCallScheduler(String name,
 
     @Override
     protected final Optional<MutationCallBatcherKey> find(ReqT subCall) {
-        Optional<KVRangeSetting> rangeSetting = storeClient.findByKey(rangeKey(subCall));
+        Optional<KVRangeSetting> rangeSetting = findByKey(rangeKey(subCall), storeClient.latestEffectiveRouter());
         return rangeSetting.map(setting -> new MutationCallBatcherKey(setting.id, setting.leader, setting.ver));
     }
 

base-kv/base-kv-store-client/src/main/java/com/baidu/bifromq/basekv/client/scheduler/QueryCallScheduler.java

@@ -13,8 +13,10 @@
 
 package com.baidu.bifromq.basekv.client.scheduler;
 
-import com.baidu.bifromq.basekv.client.KVRangeSetting;
+import static com.baidu.bifromq.basekv.client.KVRangeRouterUtil.findByKey;
+
 import com.baidu.bifromq.basekv.client.IBaseKVStoreClient;
+import com.baidu.bifromq.basekv.client.KVRangeSetting;
 import com.baidu.bifromq.basescheduler.BatchCallScheduler;
 import com.google.protobuf.ByteString;
 import java.time.Duration;
@@ -51,7 +53,7 @@ protected String selectStore(KVRangeSetting setting, ReqT request) {
 
     @Override
     protected final Optional<QueryCallBatcherKey> find(ReqT req) {
-        return storeClient.findByKey(rangeKey(req)).map(
+        return findByKey(rangeKey(req), storeClient.latestEffectiveRouter()).map(
             range -> new QueryCallBatcherKey(range.id, selectStore(range, req), selectQueue(req), range.ver));
     }