implement slot to keys with rocksdb

src/IStorage.h

@@ -31,6 +31,7 @@ public:
     virtual void retrieve(const char *key, size_t cchKey, callbackSingle fn) const = 0;
     virtual size_t clear() = 0;
     virtual bool enumerate(callback fn) const = 0;
+    virtual bool enumerate_hashslot(callback fn, int hashslot) const = 0;
     virtual size_t count() const = 0;
 
     virtual void bulkInsert(char **rgkeys, size_t *rgcbkeys, char **rgvals, size_t *rgcbvals, size_t celem) {

src/StorageCache.h

@@ -49,6 +49,7 @@ public:
     void expand(uint64_t slots);
 
     bool enumerate(IStorage::callback fn) const { return m_spstorage->enumerate(fn); }
+    bool enumerate_hashslot(IStorage::callback fn, int hashslot) const { return m_spstorage->enumerate_hashslot(fn, hashslot); }
 
     void beginWriteBatch();
     void endWriteBatch() { m_spstorage->endWriteBatch(); }

src/db.cpp

@@ -2489,26 +2489,29 @@ void slotToKeyUpdateKey(sds key, int add) {
 void slotToKeyUpdateKeyCore(const char *key, size_t keylen, int add) {
     serverAssert(GlobalLocksAcquired());
 
-    unsigned int hashslot = keyHashSlot(key,keylen);
-    unsigned char buf[64];
-    unsigned char *indexed = buf;
-
     g_pserver->cluster->slots_keys_count[hashslot] += add ? 1 : -1;
-    if (keylen+2 > 64) indexed = (unsigned char*)zmalloc(keylen+2, MALLOC_SHARED);
-    indexed[0] = (hashslot >> 8) & 0xff;
-    indexed[1] = hashslot & 0xff;
-    memcpy(indexed+2,key,keylen);
-    int fModified = false;
-    if (add) {
-        fModified = raxInsert(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL,NULL);
-    } else {
-        fModified = raxRemove(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL);
+
+    if (g_pserver->m_pstorageFactory == nullptr) {
+        unsigned int hashslot = keyHashSlot(key,keylen);
+        unsigned char buf[64];
+        unsigned char *indexed = buf;
+
+        if (keylen+2 > 64) indexed = (unsigned char*)zmalloc(keylen+2, MALLOC_SHARED);
+        indexed[0] = (hashslot >> 8) & 0xff;
+        indexed[1] = hashslot & 0xff;
+        memcpy(indexed+2,key,keylen);
+        int fModified = false;
+        if (add) {
+            fModified = raxInsert(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL,NULL);
+        } else {
+            fModified = raxRemove(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL);
+        }
+        // This assert is disabled when a snapshot depth is >0 because prepOverwriteForSnapshot will add in a tombstone,
+        //  this prevents ensure from adding the key to the dictionary which means the caller isn't aware we're already tracking
+        //  the key.
+        serverAssert(fModified || g_pserver->db[0]->snapshot_depth() > 0);
+        if (indexed != buf) zfree(indexed);
     }
-    // This assert is disabled when a snapshot depth is >0 because prepOverwriteForSnapshot will add in a tombstone,
-    //  this prevents ensure from adding the key to the dictionary which means the caller isn't aware we're already tracking
-    //  the key.
-    serverAssert(fModified || g_pserver->db[0]->snapshot_depth() > 0);
-    if (indexed != buf) zfree(indexed);
 }
 
 void slotToKeyAdd(sds key) {
@@ -2544,47 +2547,66 @@ void slotToKeyFlush(int async) {
  * New objects are returned to represent keys, it's up to the caller to
  * decrement the reference count to release the keys names. */
 unsigned int getKeysInSlot(unsigned int hashslot, robj **keys, unsigned int count) {
-    raxIterator iter;
-    int j = 0;
-    unsigned char indexed[2];
+    if (g_pserver->m_pstorageFactory != nullptr) {
+        int j = 0;
+        g_pserver->db[0]->getStorageCache()->enumerate_hashslot([&](const char *key, size_t cchKey, const void *, size_t )->bool {
+            keys[j++] = createStringObject(key, cchKey);
+            return --count;
+        }, hashslot);
+    } else {
+        raxIterator iter;
+        int j = 0;
+        unsigned char indexed[2];
 
-    indexed[0] = (hashslot >> 8) & 0xff;
-    indexed[1] = hashslot & 0xff;
-    raxStart(&iter,g_pserver->cluster->slots_to_keys);
-    raxSeek(&iter,">=",indexed,2);
-    while(count-- && raxNext(&iter)) {
-        if (iter.key[0] != indexed[0] || iter.key[1] != indexed[1]) break;
-        keys[j++] = createStringObject((char*)iter.key+2,iter.key_len-2);
+        indexed[0] = (hashslot >> 8) & 0xff;
+        indexed[1] = hashslot & 0xff;
+        raxStart(&iter,g_pserver->cluster->slots_to_keys);
+        raxSeek(&iter,">=",indexed,2);
+        while(count-- && raxNext(&iter)) {
+            if (iter.key[0] != indexed[0] || iter.key[1] != indexed[1]) break;
+            keys[j++] = createStringObject((char*)iter.key+2,iter.key_len-2);
+        }
+        raxStop(&iter);
+        return j;
     }
-    raxStop(&iter);
-    return j;
 }
 
 /* Remove all the keys in the specified hash slot.
  * The number of removed items is returned. */
 unsigned int delKeysInSlot(unsigned int hashslot) {
     serverAssert(GlobalLocksAcquired());
-    
-    raxIterator iter;
-    int j = 0;
-    unsigned char indexed[2];
+    if (g_pserver->m_pstorageFactory != nullptr) {
+        int j = 0;
+        g_pserver->db[0]->getStorageCache()->enumerate_hashslot([&](const char *key, size_t cchKey, const void *, size_t )->bool {
+            robj *keyobj = createStringObject(key, cchKey);
+            dbDelete(g_pserver->db[0], keyobj);
+            decrRefCount(keyobj);
+            j++;
+            return true;
+        }, hashslot);
+        return j;
+    } else {
+        raxIterator iter;
+        int j = 0;
+        unsigned char indexed[2];
 
-    indexed[0] = (hashslot >> 8) & 0xff;
-    indexed[1] = hashslot & 0xff;
-    raxStart(&iter,g_pserver->cluster->slots_to_keys);
-    while(g_pserver->cluster->slots_keys_count[hashslot]) {
-        raxSeek(&iter,">=",indexed,2);
-        raxNext(&iter);
+        indexed[0] = (hashslot >> 8) & 0xff;
+        indexed[1] = hashslot & 0xff;
+        raxStart(&iter,g_pserver->cluster->slots_to_keys);
+        while(g_pserver->cluster->slots_keys_count[hashslot]) {
+            raxSeek(&iter,">=",indexed,2);
+            raxNext(&iter);
 
-        auto count = g_pserver->cluster->slots_keys_count[hashslot];
-        robj *key = createStringObject((char*)iter.key+2,iter.key_len-2);
-        dbDelete(g_pserver->db[0],key);
-        serverAssert(count > g_pserver->cluster->slots_keys_count[hashslot]);   // we should have deleted something or we will be in an infinite loop
-        decrRefCount(key);
-        j++;
+            auto count = g_pserver->cluster->slots_keys_count[hashslot];
+            robj *key = createStringObject((char*)iter.key+2,iter.key_len-2);
+            dbDelete(g_pserver->db[0],key);
+            serverAssert(count > g_pserver->cluster->slots_keys_count[hashslot]);   // we should have deleted something or we will be in an infinite loop
+            decrRefCount(key);
+            j++;
+        }
+        raxStop(&iter);
+        return j;
     }
-    raxStop(&iter);
-    return j;
 }
 
 unsigned int countKeysInSlot(unsigned int hashslot) {

src/server.h

@@ -1209,6 +1209,7 @@ public:
     bool FSnapshot() const { return m_spdbSnapshotHOLDER != nullptr; }
 
     std::unique_ptr<const StorageCache> CloneStorageCache() { return std::unique_ptr<const StorageCache>(m_spstorage->clone()); }
+    std::shared_ptr<StorageCache> getStorageCache() { return m_spstorage; }
     void bulkDirectStorageInsert(char **rgKeys, size_t *rgcbKeys, char **rgVals, size_t *rgcbVals, size_t celem);
 
     dict_iter find_cached_threadsafe(const char *key) const;
@@ -1370,6 +1371,7 @@ struct redisDb : public redisDbPersistentDataSnapshot
     using redisDbPersistentData::FRehashing;
     using redisDbPersistentData::FTrackingChanges;
     using redisDbPersistentData::CloneStorageCache;
+    using redisDbPersistentData::getStorageCache();
     using redisDbPersistentData::bulkDirectStorageInsert;
 
 public:

src/storage/rocksdb.cpp

@@ -22,11 +22,15 @@ bool FInternalKey(const char *key, size_t cch)
     return false;
 }
 
+std::string getPrefix(int hashslot)
+{
+    char *hash_char = (char *)&hashslot;
+    return std::string(hash_char + (sizeof(unsigned int) - 2), 2);
+}
+
 std::string prefixKey(const char *key, size_t cchKey)
 {
-    unsigned int hash = keyHashSlot(key, cchKey);
-    char *hash_char = (char *)&hash;
-    return FInternalKey(key, cchKey) ? std::string(key, cchKey) : std::string(hash_char + (sizeof(unsigned int) - 2), 2) + std::string(key, cchKey);
+    return FInternalKey(key, cchKey) ? std::string(key, cchKey) : getPrefix(keyHashSlot(key, cchKey)) + std::string(key, cchKey);
 }
 
 RocksDBStorageProvider::RocksDBStorageProvider(RocksDBStorageFactory *pfactory, std::shared_ptr<rocksdb::DB> &spdb, std::shared_ptr<rocksdb::ColumnFamilyHandle> &spcolfam, const rocksdb::Snapshot *psnapshot, size_t count)
@@ -182,6 +186,31 @@ bool RocksDBStorageProvider::enumerate(callback fn) const
     return !it->Valid();
 }
 
+bool RocksDBStorageProvider::enumerate_hashslot(callback fn, int hashslot) const
+{
+    std::string prefix = getPrefix(hashslot);
+    std::unique_ptr<rocksdb::Iterator> it = std::unique_ptr<rocksdb::Iterator>(m_spdb->NewIterator(ReadOptions(), m_spcolfamily.get()));
+    size_t count = 0;
+    for (it->Seek(prefix.cstr()); it->Valid(); it->Next()) {
+        if (FInternalKey(it->key().data(), it->key().size()))
+            continue;
+        if (strncmp(it->key().data(),prefix.cstr(),2) != 0)
+            break;
+        ++count;
+        bool fContinue = fn(it->key().data()+2, it->key().size()-2, it->value().data(), it->value().size());
+        if (!fContinue)
+            break;
+    }
+    if (!it->Valid() && count != m_count)
+    {
+        if (const_cast<RocksDBStorageProvider*>(this)->m_count != count)
+            printf("WARNING: rocksdb count mismatch");
+        const_cast<RocksDBStorageProvider*>(this)->m_count = count;
+    }
+    assert(it->status().ok()); // Check for any errors found during the scan
+    return !it->Valid();
+}
+
 const IStorage *RocksDBStorageProvider::clone() const
 {
     std::unique_lock<fastlock> l(m_lock);

src/storage/rocksdb.h

@@ -32,6 +32,7 @@ public:
     virtual void retrieve(const char *key, size_t cchKey, callbackSingle fn) const override;
     virtual size_t clear() override;
     virtual bool enumerate(callback fn) const override;
+    virtual bool enumerate_hashslot(callback fn, int hashslot) const override;
 
     virtual const IStorage *clone() const override;