optimize loading tenant state speed during startup and range reset

Author: popduke

bifromq-inbox/bifromq-inbox-rpc-definition/src/main/java/com/baidu/bifromq/inbox/util/KeyUtil.java

@@ -88,6 +88,19 @@ public static boolean isMetadataKey(ByteString key) {
         return key.size() == inboxKeyPrefixLength(key);
     }
 
+    public static boolean hasInboxKeyPrefix(ByteString key) {
+        if (key.size() > SCHEMA_VER.size() + Integer.BYTES) {
+            int tenantPrefixLength = tenantPrefixLength(key);
+            if (key.size() > tenantPrefixLength + Integer.BYTES) {
+                int inboxPrefixLength = inboxPrefixLength(key);
+                if (key.size() >= inboxPrefixLength + Long.BYTES) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     public static boolean isQoS0MessageKey(ByteString key) {
         // QOS0 MessageKey: <INBOX_KEY_PREFIX><QOS0INBOX_SIGN><SEQ>
         int inboxKeyPrefixLen = inboxKeyPrefixLength(key);

bifromq-inbox/bifromq-inbox-rpc-definition/src/test/java/com/baidu/bifromq/inbox/util/KeyUtilTest.java

@@ -38,6 +38,11 @@ public void buildAndParse() {
         assertTrue(isInboxKey(inboxKeyPrefix));
         assertTrue(isInboxKey(qos0MsgKey));
         assertTrue(isInboxKey(bufferMsgKey));
+        assertFalse(KeyUtil.hasInboxKeyPrefix(tenantPrefix));
+        assertFalse(KeyUtil.hasInboxKeyPrefix(inboxPrefix));
+        assertTrue(KeyUtil.hasInboxKeyPrefix(inboxKeyPrefix));
+        assertTrue(KeyUtil.hasInboxKeyPrefix(qos0MsgKey));
+        assertTrue(KeyUtil.hasInboxKeyPrefix(bufferMsgKey));
 
         assertTrue(inboxPrefix.startsWith(tenantPrefix));
         assertTrue(inboxKeyPrefix.startsWith(inboxPrefix));

bifromq-inbox/bifromq-inbox-store/src/main/java/com/baidu/bifromq/inbox/store/InboxStoreCoProc.java

@@ -14,11 +14,13 @@
 package com.baidu.bifromq.inbox.store;
 
 import static com.baidu.bifromq.inbox.util.KeyUtil.bufferMsgKey;
+import static com.baidu.bifromq.inbox.util.KeyUtil.hasInboxKeyPrefix;
 import static com.baidu.bifromq.inbox.util.KeyUtil.inboxKeyPrefix;
 import static com.baidu.bifromq.inbox.util.KeyUtil.inboxKeyUpperBound;
 import static com.baidu.bifromq.inbox.util.KeyUtil.isBufferMessageKey;
 import static com.baidu.bifromq.inbox.util.KeyUtil.isMetadataKey;
 import static com.baidu.bifromq.inbox.util.KeyUtil.isQoS0MessageKey;
+import static com.baidu.bifromq.inbox.util.KeyUtil.parseInboxKeyPrefix;
 import static com.baidu.bifromq.inbox.util.KeyUtil.parseSeq;
 import static com.baidu.bifromq.inbox.util.KeyUtil.parseTenantId;
 import static com.baidu.bifromq.inbox.util.KeyUtil.qos0InboxMsgKey;
@@ -109,6 +111,7 @@ final class InboxStoreCoProc implements IKVRangeCoProc {
     // make it configurable?
     private static final int MAX_GC_BATCH_SIZE = 10000;
     private static long initTime;
+    private final KVRangeId id;
     private final ISettingProvider settingProvider;
     private final IEventCollector eventCollector;
     private final TenantsState tenantStates;
@@ -120,12 +123,14 @@ final class InboxStoreCoProc implements IKVRangeCoProc {
                      ISettingProvider settingProvider,
                      IEventCollector eventCollector,
                      Supplier<IKVReader> rangeReaderProvider) {
+        this.id = id;
         initTime = HLC.INST.getPhysical();
         this.settingProvider = settingProvider;
         this.eventCollector = eventCollector;
         this.rangeReaderProvider = rangeReaderProvider;
         this.tenantStates = new TenantsState(rangeReaderProvider.get(),
             "clusterId", clusterId, "storeId", storeId, "rangeId", KVRangeIdUtil.toString(id));
+        log.debug("Loading tenant states: rangeId={}", KVRangeIdUtil.toString(id));
         load();
     }
 
@@ -212,6 +217,7 @@ public Supplier<RWCoProcOutput> mutate(RWCoProcInput input, IKVReader reader, IK
     @Override
     public void reset(Boundary boundary) {
         tenantStates.reset();
+        log.debug("Reloading tenant states: rangeId={}", KVRangeIdUtil.toString(id));
         load();
     }
 
@@ -987,17 +993,33 @@ private void commitToInbox(ByteString scopedInboxId,
     private void load() {
         IKVReader reader = rangeReaderProvider.get();
         IKVIterator itr = reader.iterator();
+        int probe = 0;
         for (itr.seekToFirst(); itr.isValid(); ) {
             if (isMetadataKey(itr.key())) {
+                probe = 0;
                 try {
                     tenantStates.upsert(parseTenantId(itr.key()), InboxMetadata.parseFrom(itr.value()));
                 } catch (InvalidProtocolBufferException e) {
                     log.error("Unexpected error", e);
                 } finally {
-                    itr.seek(inboxKeyUpperBound(itr.key()));
+                    itr.next();
+                    probe++;
+                }
+            } else {
+                if (probe < 20) {
+                    itr.next();
+                    probe++;
+                } else {
+                    if (hasInboxKeyPrefix(itr.key())) {
+                        itr.seek(inboxKeyUpperBound(parseInboxKeyPrefix(itr.key())));
+                    } else {
+                        itr.next();
+                        probe++;
+                    }
                 }
             }
         }
+        log.debug("Tenant states loaded: rangeId={}", KVRangeIdUtil.toString(id));
     }
 
     private boolean hasExpired(InboxMetadata metadata, long nowTS) {