update prefix to uint16_t

Snapchat · Oct 6, 2023 · e771f08 · e771f08
1 parent 348c78c
commit e771f08
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 9 deletions.
diff --git a/src/storage/rocksdb.cpp b/src/storage/rocksdb.cpp
@@ -26,8 +26,9 @@ bool FInternalKey(const char *key, size_t cch)
 
 std::string getPrefix(unsigned int hashslot)
 {
-    char *hash_char = (char *)&hashslot;
-    return std::string(hash_char, sizeof(unsigned int));
+    HASHSLOT_PREFIX_TYPE slot = (HASHSLOT_PREFIX_TYPE)hashslot;
+    char *hash_char = (char *)&slot;
+    return std::string(hash_char, HASHSLOT_PREFIX_BYTES);
 }
 
 std::string prefixKey(const char *key, size_t cchKey)
@@ -186,7 +187,7 @@ bool RocksDBStorageProvider::enumerate(callback fn) const
         if (FInternalKey(it->key().data(), it->key().size()))
             continue;
         ++count;
-        bool fContinue = fn(it->key().data()+sizeof(unsigned int), it->key().size()-sizeof(unsigned int), it->value().data(), it->value().size());
+        bool fContinue = fn(it->key().data()+HASHSLOT_PREFIX_BYTES, it->key().size()-HASHSLOT_PREFIX_BYTES, it->value().data(), it->value().size());
         if (!fContinue)
             break;
     }
@@ -208,14 +209,14 @@ bool RocksDBStorageProvider::enumerate_hashslot(callback fn, unsigned int hashsl
     for (it->Seek(prefix); it->Valid(); it->Next()) {
         if (FInternalKey(it->key().data(), it->key().size()))
             continue;
-        if (*(unsigned int *)it->key().data() != hashslot)
+        if (*(HASHSLOT_PREFIX_TYPE *)it->key().data() != hashslot)
             break;
         ++count;
-        bool fContinue = fn(it->key().data()+sizeof(unsigned int), it->key().size()-sizeof(unsigned int), it->value().data(), it->value().size());
+        bool fContinue = fn(it->key().data()+HASHSLOT_PREFIX_BYTES, it->key().size()-HASHSLOT_PREFIX_BYTES, it->value().data(), it->value().size());
         if (!fContinue)
             break;
     }
-    bool full_iter = !it->Valid() || (*(unsigned int *)it->key().data() != hashslot);
+    bool full_iter = !it->Valid() || (*(HASHSLOT_PREFIX_TYPE *)it->key().data() != hashslot);
     if (full_iter && count != g_pserver->cluster->slots_keys_count[hashslot])
     {
         printf("WARNING: rocksdb hashslot count mismatch");
@@ -280,7 +281,7 @@ std::vector<std::string> RocksDBStorageProvider::getEvictionCandidates(unsigned
         for (it->Seek(randomHashSlot()); it->Valid() && result.size() < count; it->Next()) {
             if (FInternalKey(it->key().data(), it->key().size()))
                 continue;
-            result.emplace_back(it->key().data() + sizeof(unsigned int), it->key().size() - sizeof(unsigned int));
+            result.emplace_back(it->key().data() + HASHSLOT_PREFIX_BYTES, it->key().size() - HASHSLOT_PREFIX_BYTES);
         }
     } else {
         std::unique_ptr<rocksdb::Iterator> it = std::unique_ptr<rocksdb::Iterator>(m_spdb->NewIterator(ReadOptions(), m_spexpirecolfamily.get()));

diff --git a/src/storage/rocksdb.h b/src/storage/rocksdb.h
@@ -7,6 +7,8 @@
 #include "../fastlock.h"
 
 #define INTERNAL_KEY_PREFIX "\x00\x04\x03\x00\x05\x02\x04"
+#define HASHSLOT_PREFIX_TYPE uint16_t
+#define HASHSLOT_PREFIX_BYTES sizeof(HASHSLOT_PREFIX_TYPE)
 static const char count_key[] = INTERNAL_KEY_PREFIX "__keydb__count\1";
 static const char version_key[] = INTERNAL_KEY_PREFIX "__keydb__version\1";
 static const char meta_key[] = INTERNAL_KEY_PREFIX "__keydb__metadata\1";

diff --git a/src/storage/rocksdbfactory.cpp b/src/storage/rocksdbfactory.cpp
@@ -72,7 +72,7 @@ RocksDBStorageFactory::RocksDBStorageFactory(const char *dbfile, int dbnum, cons
     rocksdb::DB *db = nullptr;
 
     auto options = RocksDbOptions();
-    options.prefix_extractor.reset(rocksdb::NewFixedPrefixTransform(sizeof(unsigned int)));
+    options.prefix_extractor.reset(rocksdb::NewFixedPrefixTransform(HASHSLOT_PREFIX_BYTES));
 
     for (int idb = 0; idb < dbnum; ++idb)
     {
@@ -196,7 +196,7 @@ IStorage *RocksDBStorageFactory::create(int db, key_load_iterator iter, void *pr
                 printf("\tDatabase %d was not shutdown cleanly, recomputing metrics\n", db);
             fFirstRealKey = false;
             if (iter != nullptr)
-                iter(it->key().data()+sizeof(unsigned int), it->key().size()-sizeof(unsigned int), privdata);
+                iter(it->key().data()+HASHSLOT_PREFIX_BYTES, it->key().size()-HASHSLOT_PREFIX_BYTES, privdata);
             ++count;
         }
     }