chore: preparation step for lock fingerprints

The main change here is introduction of the strong type LockTag
that differentiates from a string_view key.

Also, some testing improvements to improve the footprint of the next PR.

Signed-off-by: Roman Gershman <[email protected]>

Author: romange

src/core/heap_size.h

@@ -29,8 +29,14 @@ namespace dfly {
 
 namespace heap_size_detail {
 
+template <class, class = void> struct has_marked_stackonly : std::false_type {};
+
+template <class T>
+struct has_marked_stackonly<T, std::void_t<typename T::is_stackonly>> : std::true_type {};
+
 template <typename T> constexpr bool StackOnlyType() {
-  return std::is_trivial_v<T> || std::is_same_v<T, std::string_view>;
+  return std::is_trivial_v<T> || std::is_same_v<T, std::string_view> ||
+         has_marked_stackonly<T>::value;
 }
 
 template <typename T, typename = void> struct has_used_mem : std::false_type {};

src/server/cluster/cluster_config.cc

@@ -55,32 +55,8 @@ bool ClusterConfig::IsEmulated() {
   return cluster_mode == ClusterMode::kEmulatedCluster;
 }
 
-string_view ClusterConfig::KeyTag(string_view key) {
-  auto options = KeyLockArgs::GetLockTagOptions();
-
-  if (!absl::StartsWith(key, options.prefix)) {
-    return key;
-  }
-
-  const size_t start = key.find(options.open_locktag);
-  if (start == key.npos) {
-    return key;
-  }
-
-  size_t end = start;
-  for (unsigned i = 0; i <= options.skip_n_end_delimiters; ++i) {
-    size_t next = end + 1;
-    end = key.find(options.close_locktag, next);
-    if (end == key.npos || end == next) {
-      return key;
-    }
-  }
-
-  return key.substr(start + 1, end - start - 1);
-}
-
 SlotId ClusterConfig::KeySlot(string_view key) {
-  string_view tag = KeyTag(key);
+  string_view tag = LockTagOptions::instance().Tag(key);
   return crc16(tag.data(), tag.length()) & kMaxSlotNum;
 }
 

src/server/cluster/cluster_config.h

@@ -64,12 +64,9 @@ class ClusterConfig {
   }
 
   static bool IsShardedByTag() {
-    return IsEnabledOrEmulated() || KeyLockArgs::GetLockTagOptions().enabled;
+    return IsEnabledOrEmulated() || LockTagOptions::instance().enabled;
   }
 
-  // If the key contains the {...} pattern, return only the part between { and }
-  static std::string_view KeyTag(std::string_view key);
-
   // Returns an instance with `config` if it is valid.
   // Returns heap-allocated object as it is too big for a stack frame.
   static std::shared_ptr<ClusterConfig> CreateFromConfig(std::string_view my_id,

src/server/cluster/cluster_config_test.cc

@@ -27,63 +27,67 @@ class ClusterConfigTest : public BaseFamilyTest {
   const string kMyId = "my-id";
 };
 
+inline string_view GetTag(string_view key) {
+  return LockTagOptions::instance().Tag(key);
+}
+
 TEST_F(ClusterConfigTest, KeyTagTest) {
   SetTestFlag("lock_on_hashtags", "true");
 
-  EXPECT_EQ(ClusterConfig::KeyTag("{user1000}.following"), "user1000");
+  EXPECT_EQ(GetTag("{user1000}.following"), "user1000");
 
-  EXPECT_EQ(ClusterConfig::KeyTag("foo{{bar}}zap"), "{bar");
+  EXPECT_EQ(GetTag("foo{{bar}}zap"), "{bar");
 
-  EXPECT_EQ(ClusterConfig::KeyTag("foo{bar}{zap}"), "bar");
+  EXPECT_EQ(GetTag("foo{bar}{zap}"), "bar");
 
   string_view key = " foo{}{bar}";
-  EXPECT_EQ(key, ClusterConfig::KeyTag(key));
+  EXPECT_EQ(key, GetTag(key));
 
   key = "{}foo{bar}{zap}";
-  EXPECT_EQ(key, ClusterConfig::KeyTag(key));
+  EXPECT_EQ(key, GetTag(key));
 
   SetTestFlag("locktag_delimiter", ":");
   TEST_InvalidateLockTagOptions();
 
   key = "{user1000}.following";
-  EXPECT_EQ(ClusterConfig::KeyTag(key), key);
+  EXPECT_EQ(GetTag(key), key);
 
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue1:123"), "queue1");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123"), "queue");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123:456:789:1000"), "queue");
+  EXPECT_EQ(GetTag("bull:queue1:123"), "queue1");
+  EXPECT_EQ(GetTag("bull:queue:1:123"), "queue");
+  EXPECT_EQ(GetTag("bull:queue:1:123:456:789:1000"), "queue");
 
   key = "bull::queue:1:123";
-  EXPECT_EQ(ClusterConfig::KeyTag(key), key);
+  EXPECT_EQ(GetTag(key), key);
 
   SetTestFlag("locktag_delimiter", ":");
   SetTestFlag("locktag_skip_n_end_delimiters", "0");
   SetTestFlag("locktag_prefix", "bull");
   TEST_InvalidateLockTagOptions();
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:123"), "queue");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:123:456:789:1000"), "queue");
+  EXPECT_EQ(GetTag("bull:queue:123"), "queue");
+  EXPECT_EQ(GetTag("bull:queue:123:456:789:1000"), "queue");
 
   key = "not-bull:queue1:123";
-  EXPECT_EQ(ClusterConfig::KeyTag(key), key);
+  EXPECT_EQ(GetTag(key), key);
 
   SetTestFlag("locktag_delimiter", ":");
   SetTestFlag("locktag_skip_n_end_delimiters", "1");
   SetTestFlag("locktag_prefix", "bull");
   TEST_InvalidateLockTagOptions();
 
   key = "bull:queue1:123";
-  EXPECT_EQ(ClusterConfig::KeyTag(key), key);
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123"), "queue:1");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123:456:789:1000"), "queue:1");
+  EXPECT_EQ(GetTag(key), key);
+  EXPECT_EQ(GetTag("bull:queue:1:123"), "queue:1");
+  EXPECT_EQ(GetTag("bull:queue:1:123:456:789:1000"), "queue:1");
 
   key = "bull::queue:1:123";
-  EXPECT_EQ(ClusterConfig::KeyTag(key), key);
+  EXPECT_EQ(GetTag(key), key);
 
   SetTestFlag("locktag_delimiter", "|");
   SetTestFlag("locktag_skip_n_end_delimiters", "2");
   SetTestFlag("locktag_prefix", "");
   TEST_InvalidateLockTagOptions();
 
-  EXPECT_EQ(ClusterConfig::KeyTag("|a|b|c|d|e"), "a|b|c");
+  EXPECT_EQ(GetTag("|a|b|c|d|e"), "a|b|c");
 }
 
 TEST_F(ClusterConfigTest, ConfigSetInvalidEmpty) {

src/server/common.cc

@@ -52,8 +52,10 @@ using namespace std;
 using namespace util;
 
 namespace {
+
 // Thread-local cache with static linkage.
 thread_local std::optional<LockTagOptions> locktag_lock_options;
+
 }  // namespace
 
 void TEST_InvalidateLockTagOptions() {
@@ -63,7 +65,7 @@ void TEST_InvalidateLockTagOptions() {
       [](ShardId shard, ProactorBase* proactor) { locktag_lock_options = nullopt; });
 }
 
-/* static */ LockTagOptions KeyLockArgs::GetLockTagOptions() {
+const LockTagOptions& LockTagOptions::instance() {
   if (!locktag_lock_options.has_value()) {
     string delimiter = absl::GetFlag(FLAGS_locktag_delimiter);
     if (delimiter.empty()) {
@@ -87,12 +89,26 @@ void TEST_InvalidateLockTagOptions() {
   return *locktag_lock_options;
 }
 
-string_view KeyLockArgs::GetLockKey(string_view key) {
-  if (GetLockTagOptions().enabled) {
-    return ClusterConfig::KeyTag(key);
+std::string_view LockTagOptions::Tag(std::string_view key) const {
+  if (!absl::StartsWith(key, prefix)) {
+    return key;
+  }
+
+  const size_t start = key.find(open_locktag);
+  if (start == key.npos) {
+    return key;
+  }
+
+  size_t end = start;
+  for (unsigned i = 0; i <= skip_n_end_delimiters; ++i) {
+    size_t next = end + 1;
+    end = key.find(close_locktag, next);
+    if (end == key.npos || end == next) {
+      return key;
+    }
   }
 
-  return key;
+  return key.substr(start + 1, end - start - 1);
 }
 
 atomic_uint64_t used_mem_peak(0);
@@ -445,4 +461,15 @@ std::ostream& operator<<(std::ostream& os, ArgSlice list) {
   return os << "]";
 }
 
+LockTag::LockTag(std::string_view key) {
+  if (LockTagOptions::instance().enabled)
+    str_ = LockTagOptions::instance().Tag(key);
+  else
+    str_ = key;
+}
+
+LockFp LockTag::Fingerprint() const {
+  return XXH64(str_.data(), str_.size(), 0x1C69B3F74AC4AE35UL);
+}
+
 }  // namespace dfly

src/server/common.h

@@ -48,6 +48,7 @@ using RdbTypeFreqMap = absl::flat_hash_map<unsigned, size_t>;
 constexpr DbIndex kInvalidDbId = DbIndex(-1);
 constexpr ShardId kInvalidSid = ShardId(-1);
 constexpr DbIndex kMaxDbId = 1024;  // Reasonable starting point.
+using LockFp = uint64_t;            // a key fingerprint used by the LockTable.
 
 class CommandId;
 class Transaction;
@@ -59,14 +60,14 @@ struct LockTagOptions {
   char close_locktag = '}';
   unsigned skip_n_end_delimiters = 0;
   std::string prefix;
-};
 
-struct KeyLockArgs {
-  static LockTagOptions GetLockTagOptions();
+  // Returns the tag according to the rules defined by this options object.
+  std::string_view Tag(std::string_view key) const;
 
-  // Before acquiring and releasing keys, one must "normalize" them via GetLockKey().
-  static std::string_view GetLockKey(std::string_view key);
+  static const LockTagOptions& instance();
+};
 
+struct KeyLockArgs {
   DbIndex db_index = 0;
   ArgSlice args;
   unsigned key_step = 1;
@@ -117,6 +118,33 @@ struct OpArgs {
   }
 };
 
+// A strong type for a lock tag. Helps to disambiguide between keys and the parts of the
+// keys that are used for locking.
+class LockTag {
+  std::string_view str_;
+
+ public:
+  using is_stackonly = void;  // marks that this object does not use heap.
+
+  LockTag() = default;
+  explicit LockTag(std::string_view key);
+
+  explicit operator std::string_view() const {
+    return str_;
+  }
+
+  LockFp Fingerprint() const;
+
+  // To make it hashable.
+  template <typename H> friend H AbslHashValue(H h, const LockTag& tag) {
+    return H::combine(std::move(h), tag.str_);
+  }
+
+  bool operator==(const LockTag& o) const {
+    return str_ == o.str_;
+  }
+};
+
 // Record non auto journal command with own txid and dbid.
 void RecordJournal(const OpArgs& op_args, std::string_view cmd, ArgSlice args,
                    uint32_t shard_cnt = 1, bool multi_commands = false);

src/server/conn_context.cc

@@ -236,7 +236,7 @@ size_t ConnectionState::ExecInfo::UsedMemory() const {
 }
 
 size_t ConnectionState::ScriptInfo::UsedMemory() const {
-  return dfly::HeapSize(keys) + async_cmds_heap_mem;
+  return dfly::HeapSize(lock_tags) + async_cmds_heap_mem;
 }
 
 size_t ConnectionState::SubscribeInfo::UsedMemory() const {

src/server/conn_context.h

@@ -100,7 +100,7 @@ struct ConnectionState {
   struct ScriptInfo {
     size_t UsedMemory() const;
 
-    absl::flat_hash_set<std::string_view> keys;  // declared keys
+    absl::flat_hash_set<LockTag> lock_tags;  // declared tags
 
     size_t async_cmds_heap_mem = 0;     // bytes used by async_cmds
     size_t async_cmds_heap_limit = 0;   // max bytes allowed for async_cmds

src/server/db_slice.cc

@@ -198,7 +198,7 @@ unsigned PrimeEvictionPolicy::Evict(const PrimeTable::HotspotBuckets& eb, PrimeT
     string scratch;
     string_view key = last_slot_it->first.GetSlice(&scratch);
     // do not evict locked keys
-    if (lt.Find(KeyLockArgs::GetLockKey(key)).has_value())
+    if (lt.Find(LockTag(key)).has_value())
       return 0;
 
     // log the evicted keys to journal.
@@ -998,16 +998,16 @@ bool DbSlice::Acquire(IntentLock::Mode mode, const KeyLockArgs& lock_args) {
   bool lock_acquired = true;
 
   if (lock_args.args.size() == 1) {
-    string_view key = KeyLockArgs::GetLockKey(lock_args.args.front());
-    lock_acquired = lt.Acquire(key, mode);
-    uniq_keys_ = {key};  // needed only for tests.
+    LockTag tag(lock_args.args.front());
+    lock_acquired = lt.Acquire(tag, mode);
+    uniq_keys_ = {string_view(tag)};  // needed only for tests.
   } else {
     uniq_keys_.clear();
 
     for (size_t i = 0; i < lock_args.args.size(); i += lock_args.key_step) {
-      string_view s = KeyLockArgs::GetLockKey(lock_args.args[i]);
-      if (uniq_keys_.insert(s).second) {
-        lock_acquired &= lt.Acquire(s, mode);
+      LockTag tag(lock_args.args[i]);
+      if (uniq_keys_.insert(string_view(tag)).second) {
+        lock_acquired &= lt.Acquire(tag, mode);
       }
     }
   }
@@ -1018,13 +1018,12 @@ bool DbSlice::Acquire(IntentLock::Mode mode, const KeyLockArgs& lock_args) {
   return lock_acquired;
 }
 
-void DbSlice::ReleaseNormalized(IntentLock::Mode mode, DbIndex db_index, std::string_view key) {
-  DCHECK_EQ(key, KeyLockArgs::GetLockKey(key));
+void DbSlice::ReleaseNormalized(IntentLock::Mode mode, DbIndex db_index, LockTag tag) {
   DVLOG(2) << "Release " << IntentLock::ModeName(mode) << " "
-           << " for " << key;
+           << " for " << string_view(tag);
 
   auto& lt = db_arr_[db_index]->trans_locks;
-  lt.Release(KeyLockArgs::GetLockKey(key), mode);
+  lt.Release(tag, mode);
 }
 
 void DbSlice::Release(IntentLock::Mode mode, const KeyLockArgs& lock_args) {
@@ -1034,15 +1033,15 @@ void DbSlice::Release(IntentLock::Mode mode, const KeyLockArgs& lock_args) {
 
   DVLOG(2) << "Release " << IntentLock::ModeName(mode) << " for " << lock_args.args[0];
   if (lock_args.args.size() == 1) {
-    string_view key = KeyLockArgs::GetLockKey(lock_args.args.front());
-    ReleaseNormalized(mode, lock_args.db_index, key);
+    string_view key = lock_args.args.front();
+    ReleaseNormalized(mode, lock_args.db_index, LockTag{key});
   } else {
     auto& lt = db_arr_[lock_args.db_index]->trans_locks;
     uniq_keys_.clear();
     for (size_t i = 0; i < lock_args.args.size(); i += lock_args.key_step) {
-      string_view s = KeyLockArgs::GetLockKey(lock_args.args[i]);
-      if (uniq_keys_.insert(s).second) {
-        lt.Release(s, mode);
+      LockTag tag(lock_args.args[i]);
+      if (uniq_keys_.insert(string_view(tag)).second) {
+        lt.Release(tag, mode);
       }
     }
   }
@@ -1051,9 +1050,9 @@ void DbSlice::Release(IntentLock::Mode mode, const KeyLockArgs& lock_args) {
 
 bool DbSlice::CheckLock(IntentLock::Mode mode, DbIndex dbid, string_view key) const {
   const auto& lt = db_arr_[dbid]->trans_locks;
-  string_view s = KeyLockArgs::GetLockKey(key);
+  LockTag tag(key);
 
-  auto lock = lt.Find(s);
+  auto lock = lt.Find(tag);
   if (lock) {
     return lock->Check(mode);
   }
@@ -1322,7 +1321,7 @@ void DbSlice::FreeMemWithEvictionStep(DbIndex db_ind, size_t increase_goal_bytes
           // check if the key is locked by looking up transaction table.
           const auto& lt = db_table->trans_locks;
           string_view key = evict_it->first.GetSlice(&tmp);
-          if (lt.Find(KeyLockArgs::GetLockKey(key)).has_value())
+          if (lt.Find(LockTag(key)).has_value())
             continue;
 
           if (auto journal = owner_->journal(); journal) {

src/server/db_slice.h

@@ -477,8 +477,8 @@ class DbSlice {
   void PerformDeletion(Iterator del_it, DbTable* table);
   void PerformDeletion(PrimeIterator del_it, DbTable* table);
 
-  // Releases a single key. `key` must have been normalized by GetLockKey().
-  void ReleaseNormalized(IntentLock::Mode m, DbIndex db_index, std::string_view key);
+  // Releases a single tag.
+  void ReleaseNormalized(IntentLock::Mode m, DbIndex db_index, LockTag tag);
 
  private:
   void PreUpdate(DbIndex db_ind, Iterator it);