chore: preparation step for lock fingerprints (#2899)

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 <roman@dragonflydb.io>

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,43 +27,47 @@ 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(GetTag("bull:queue1:123"), "queue1");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123"), "queue");
+  EXPECT_EQ(GetTag("bull:queue:1:123"), "queue");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123:456:789:1000"), "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(GetTag("bull:queue:123"), "queue");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:123:456:789:1000"), "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");
@@ -71,19 +75,19 @@ TEST_F(ClusterConfigTest, KeyTagTest) {
   TEST_InvalidateLockTagOptions();
 
   key = "bull:queue1:123";
-  EXPECT_EQ(ClusterConfig::KeyTag(key), key);
+  EXPECT_EQ(GetTag(key), key);
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123"), "queue:1");
+  EXPECT_EQ(GetTag("bull:queue:1:123"), "queue:1");
-  EXPECT_EQ(ClusterConfig::KeyTag("bull:queue:1:123:456:789:1000"), "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) {
+std::string_view LockTagOptions::Tag(std::string_view key) const {
-  if (GetLockTagOptions().enabled) {
+  if (!absl::StartsWith(key, prefix)) {
-    return ClusterConfig::KeyTag(key);
+    return key;
   }
 
-  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.substr(start + 1, end - start - 1);
 }
 
 atomic_uint64_t used_mem_peak(0);
@@ -455,4 +471,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;
+
+  // Returns the tag according to the rules defined by this options object.
+  std::string_view Tag(std::string_view key) const;
+
+  static const LockTagOptions& instance();
 };
 
 struct KeyLockArgs {
-  static LockTagOptions GetLockTagOptions();
-
-  // Before acquiring and releasing keys, one must "normalize" them via GetLockKey().
-  static std::string_view GetLockKey(std::string_view key);
-
   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());
+    LockTag tag(lock_args.args.front());
-    lock_acquired = lt.Acquire(key, mode);
+    lock_acquired = lt.Acquire(tag, mode);
-    uniq_keys_ = {key};  // needed only for tests.
+    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]);
+      LockTag tag(lock_args.args[i]);
-      if (uniq_keys_.insert(s).second) {
+      if (uniq_keys_.insert(string_view(tag)).second) {
-        lock_acquired &= lt.Acquire(s, mode);
+        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) {
+void DbSlice::ReleaseNormalized(IntentLock::Mode mode, DbIndex db_index, LockTag tag) {
-  DCHECK_EQ(key, KeyLockArgs::GetLockKey(key));
   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());
+    string_view key = lock_args.args.front();
-    ReleaseNormalized(mode, lock_args.db_index, key);
+    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]);
+      LockTag tag(lock_args.args[i]);
-      if (uniq_keys_.insert(s).second) {
+      if (uniq_keys_.insert(string_view(tag)).second) {
-        lt.Release(s, mode);
+        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().
+  // Releases a single tag.
-  void ReleaseNormalized(IntentLock::Mode m, DbIndex db_index, std::string_view key);
+  void ReleaseNormalized(IntentLock::Mode m, DbIndex db_index, LockTag tag);
 
  private:
   void PreUpdate(DbIndex db_ind, Iterator it);

src/server/engine_shard_set.cc

@@ -173,21 +173,19 @@ class RoundRobinSharder {
   static fb2::Mutex mutex_;
 };
 
-bool HasContendedLocks(unsigned shard_id, Transaction* trx, const DbTable* table) {
+bool HasContendedLocks(ShardId shard_id, Transaction* trx, const DbTable* table) {
-  auto is_contended = [table](string_view key) {
+  auto is_contended = [table](LockTag tag) { return table->trans_locks.Find(tag)->IsContended(); };
-    return table->trans_locks.Find(key)->IsContended();
-  };
 
   if (trx->IsMulti()) {
     auto keys = trx->GetMultiKeys();
     for (string_view key : keys) {
-      if (Shard(key, shard_set->size()) == shard_id && is_contended(key))
+      if (Shard(key, shard_set->size()) == shard_id && is_contended(LockTag{key}))
         return true;
     }
   } else {
     KeyLockArgs lock_args = trx->GetLockArgs(shard_id);
     for (size_t i = 0; i < lock_args.args.size(); i += lock_args.key_step) {
-      if (is_contended(KeyLockArgs::GetLockKey(lock_args.args[i])))
+      if (is_contended(LockTag{lock_args.args[i]}))
         return true;
     }
   }
@@ -868,7 +866,7 @@ void EngineShardSet::TEST_EnableCacheMode() {
 
 ShardId Shard(string_view v, ShardId shard_num) {
   if (ClusterConfig::IsShardedByTag()) {
-    v = ClusterConfig::KeyTag(v);
+    v = LockTagOptions::instance().Tag(v);
   }
 
   XXH64_hash_t hash = XXH64(v.data(), v.size(), 120577240643ULL);

src/server/main_service.cc

@@ -594,7 +594,7 @@ void ClusterHtmlPage(const http::QueryArgs& args, HttpContext* send, ClusterFami
                                                 : "Disabled");
 
   if (ClusterConfig::IsEnabledOrEmulated()) {
-    print_kb("Lock on hashtags", KeyLockArgs::GetLockTagOptions().enabled);
+    print_kb("Lock on hashtags", LockTagOptions::instance().enabled);
   }
 
   if (ClusterConfig::IsEnabled()) {
@@ -945,19 +945,18 @@ OpStatus CheckKeysDeclared(const ConnectionState::ScriptInfo& eval_info, const C
 
   // TODO: Switch to transaction internal locked keys once single hop multi transactions are merged
   // const auto& locked_keys = trans->GetMultiKeys();
-  const auto& locked_keys = eval_info.keys;
+  const auto& locked_tags = eval_info.lock_tags;
 
   const auto& key_index = *key_index_res;
   for (unsigned i = key_index.start; i < key_index.end; ++i) {
-    string_view key = KeyLockArgs::GetLockKey(ArgS(args, i));
+    LockTag tag{ArgS(args, i)};
-    if (!locked_keys.contains(key)) {
+    if (!locked_tags.contains(tag)) {
-      VLOG(1) << "Key " << key << " is not declared for command " << cid->name();
+      VLOG(1) << "Key " << string_view(tag) << " is not declared for command " << cid->name();
       return OpStatus::KEY_NOTFOUND;
     }
   }
 
-  if (key_index.bonus &&
+  if (key_index.bonus && !locked_tags.contains(LockTag{ArgS(args, *key_index.bonus)}))
-      !locked_keys.contains(KeyLockArgs::GetLockKey(ArgS(args, *key_index.bonus))))
     return OpStatus::KEY_NOTFOUND;
 
   return OpStatus::OK;
@@ -1867,7 +1866,7 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
   // we can do it once during script insertion into script mgr.
   auto& sinfo = cntx->conn_state.script_info;
   sinfo = make_unique<ConnectionState::ScriptInfo>();
-  sinfo->keys.reserve(eval_args.keys.size());
+  sinfo->lock_tags.reserve(eval_args.keys.size());
 
   optional<ShardId> sid;
 
@@ -1875,7 +1874,7 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
   for (size_t i = 0; i < eval_args.keys.size(); ++i) {
     string_view key = ArgS(eval_args.keys, i);
     slot_checker.Add(key);
-    sinfo->keys.insert(KeyLockArgs::GetLockKey(key));
+    sinfo->lock_tags.insert(LockTag(key));
 
     ShardId cur_sid = Shard(key, shard_count());
     if (i == 0) {

src/server/string_family_test.cc

@@ -791,20 +791,16 @@ TEST_F(StringFamilyTest, SetWithHashtagsNoCluster) {
   SetTestFlag("lock_on_hashtags", "false");
   ResetService();
 
-  auto fb = ExpectConditionWithSuspension(
+  auto fb = ExpectUsedKeys({"{key}1"});
-      [&]() { return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}1"}; });
   EXPECT_EQ(Run({"set", "{key}1", "val1"}), "OK");
   fb.Join();
   EXPECT_FALSE(service_->IsLocked(0, "{key}1"));
 
-  fb = ExpectConditionWithSuspension(
+  fb = ExpectUsedKeys({"{key}2"});
-      [&]() { return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}2"}; });
   EXPECT_EQ(Run({"set", "{key}2", "val2"}), "OK");
   fb.Join();
 
-  fb = ExpectConditionWithSuspension([&]() {
+  fb = ExpectUsedKeys({"{key}1", "{key}2"});
-    return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}1", "{key}2"};
-  });
   EXPECT_THAT(Run({"mget", "{key}1", "{key}2"}), RespArray(ElementsAre("val1", "val2")));
   fb.Join();
   EXPECT_NE(1, GetDebugInfo().shards_count);
@@ -815,19 +811,15 @@ TEST_F(StringFamilyTest, SetWithHashtagsWithEmulatedCluster) {
   SetTestFlag("lock_on_hashtags", "false");
   ResetService();
 
-  auto fb = ExpectConditionWithSuspension(
+  auto fb = ExpectUsedKeys({"{key}1"});
-      [&]() { return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}1"}; });
   EXPECT_EQ(Run({"set", "{key}1", "val1"}), "OK");
   fb.Join();
 
-  fb = ExpectConditionWithSuspension(
+  fb = ExpectUsedKeys({"{key}2"});
-      [&]() { return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}2"}; });
   EXPECT_EQ(Run({"set", "{key}2", "val2"}), "OK");
   fb.Join();
 
-  fb = ExpectConditionWithSuspension([&]() {
+  fb = ExpectUsedKeys({"{key}1", "{key}2"});
-    return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}1", "{key}2"};
-  });
   EXPECT_THAT(Run({"mget", "{key}1", "{key}2"}), RespArray(ElementsAre("val1", "val2")));
   fb.Join();
   EXPECT_EQ(1, GetDebugInfo().shards_count);
@@ -838,16 +830,15 @@ TEST_F(StringFamilyTest, SetWithHashtagsWithHashtagLock) {
   SetTestFlag("lock_on_hashtags", "true");
   ResetService();
 
-  auto condition = [&]() { return GetLastUsedKeys() == absl::flat_hash_set<string>{"key"}; };
+  auto fb = ExpectUsedKeys({"key"});
-  auto fb = ExpectConditionWithSuspension(condition);
   EXPECT_EQ(Run({"set", "{key}1", "val1"}), "OK");
   fb.Join();
 
-  fb = ExpectConditionWithSuspension(condition);
+  fb = ExpectUsedKeys({"key"});
   EXPECT_EQ(Run({"set", "{key}2", "val2"}), "OK");
   fb.Join();
 
-  fb = ExpectConditionWithSuspension(condition);
+  fb = ExpectUsedKeys({"key"});
   EXPECT_THAT(Run({"mget", "{key}1", "{key}2"}), RespArray(ElementsAre("val1", "val2")));
   fb.Join();
   EXPECT_EQ(1, GetDebugInfo().shards_count);
@@ -858,10 +849,7 @@ TEST_F(StringFamilyTest, MultiSetWithHashtagsDontLockHashtags) {
   SetTestFlag("lock_on_hashtags", "false");
   ResetService();
 
-  auto condition = [&]() {
+  auto fb = ExpectUsedKeys({"{key}1", "{key}2", "{key}3"});
-    return GetLastUsedKeys() == absl::flat_hash_set<string>{"{key}1", "{key}2", "{key}3"};
-  };
-  auto fb = ExpectConditionWithSuspension(condition);
 
   EXPECT_EQ(Run({"multi"}), "OK");
   EXPECT_EQ(Run({"set", "{key}1", "val1"}), "QUEUED");
@@ -876,8 +864,7 @@ TEST_F(StringFamilyTest, MultiSetWithHashtagsLockHashtags) {
   SetTestFlag("lock_on_hashtags", "true");
   ResetService();
 
-  auto condition = [&]() { return GetLastUsedKeys() == absl::flat_hash_set<string>{"key"}; };
+  auto fb = ExpectUsedKeys({"key"});
-  auto fb = ExpectConditionWithSuspension(condition);
 
   EXPECT_EQ(Run({"multi"}), "OK");
   EXPECT_EQ(Run({"set", "{key}1", "val1"}), "QUEUED");

src/server/table.cc

@@ -4,8 +4,6 @@
 
 #include "server/table.h"
 
-#include <xxhash.h>
-
 #include "base/flags.h"
 #include "base/logging.h"
 #include "server/cluster/cluster_config.h"
@@ -60,36 +58,23 @@ SlotStats& SlotStats::operator+=(const SlotStats& o) {
   return *this;
 }
 
-size_t LockTable::Size() const {
+std::optional<const IntentLock> LockTable::Find(LockTag tag) const {
-  return locks_.size();
+  LockFp fp = tag.Fingerprint();
-}
-
-LockFp LockTable::Fingerprint(string_view key) {
-  return XXH64(key.data(), key.size(), 0x1C69B3F74AC4AE35UL);
-}
-
-std::optional<const IntentLock> LockTable::Find(string_view key) const {
-  DCHECK_EQ(KeyLockArgs::GetLockKey(key), key);
-
-  LockFp fp = Fingerprint(key);
   if (auto it = locks_.find(fp); it != locks_.end())
     return it->second;
   return std::nullopt;
 }
 
-bool LockTable::Acquire(string_view key, IntentLock::Mode mode) {
+bool LockTable::Acquire(LockTag tag, IntentLock::Mode mode) {
-  DCHECK_EQ(KeyLockArgs::GetLockKey(key), key);
+  LockFp fp = tag.Fingerprint();
-  LockFp fp = Fingerprint(key);
   auto [it, inserted] = locks_.try_emplace(fp);
   return it->second.Acquire(mode);
 }
 
-void LockTable::Release(string_view key, IntentLock::Mode mode) {
+void LockTable::Release(LockTag tag, IntentLock::Mode mode) {
-  DCHECK_EQ(KeyLockArgs::GetLockKey(key), key);
+  LockFp fp = tag.Fingerprint();
-
-  LockFp fp = Fingerprint(key);
   auto it = locks_.find(fp);
-  DCHECK(it != locks_.end()) << key;
+  DCHECK(it != locks_.end()) << string_view(tag);
 
   it->second.Release(mode);
   if (it->second.IsFree())

src/server/table.h

@@ -25,7 +25,6 @@ using PrimeValue = detail::PrimeValue;
 
 using PrimeTable = DashTable<PrimeKey, PrimeValue, detail::PrimeTablePolicy>;
 using ExpireTable = DashTable<PrimeKey, ExpirePeriod, detail::ExpireTablePolicy>;
-using LockFp = uint64_t;  // a key fingerprint used by the LockTable.
 
 /// Iterators are invalidated when new keys are added to the table or some entries are deleted.
 /// Iterators are still valid if a different entry in the table was mutated.
@@ -79,18 +78,17 @@ struct DbTableStats {
   DbTableStats& operator+=(const DbTableStats& o);
 };
 
-// Table for recording locks that uses string_views where possible. LockTable falls back to
+// Table for recording locks. Keys used with the lock table should be normalized with LockTag.
-// strings for locks that are used by multiple transactions. Keys used with the lock table
-// should be normalized with GetLockKey
 class LockTable {
  public:
-  size_t Size() const;
+  size_t Size() const {
-  std::optional<const IntentLock> Find(std::string_view key) const;
+    return locks_.size();
+  }
+  std::optional<const IntentLock> Find(LockTag tag) const;
+  std::optional<const IntentLock> Find(LockFp fp) const;
 
-  bool Acquire(std::string_view key, IntentLock::Mode mode);
+  bool Acquire(LockTag tag, IntentLock::Mode mode);
-  void Release(std::string_view key, IntentLock::Mode mode);
+  void Release(LockTag tag, IntentLock::Mode mode);
-
-  static LockFp Fingerprint(std::string_view key);
 
   auto begin() const {
     return locks_.cbegin();

src/server/test_utils.cc

@@ -677,6 +677,19 @@ fb2::Fiber BaseFamilyTest::ExpectConditionWithSuspension(const std::function<boo
   return fb;
 }
 
+util::fb2::Fiber BaseFamilyTest::ExpectUsedKeys(const std::vector<std::string_view>& keys) {
+  absl::flat_hash_set<string> own_keys;
+  for (const auto& k : keys) {
+    own_keys.insert(string(k));
+  }
+  auto cb = [=] {
+    auto last_keys = GetLastUsedKeys();
+    return last_keys == own_keys;
+  };
+
+  return ExpectConditionWithSuspension(std::move(cb));
+}
+
 void BaseFamilyTest::SetTestFlag(string_view flag_name, string_view new_value) {
   auto* flag = absl::FindCommandLineFlag(flag_name);
   CHECK_NE(flag, nullptr);

src/server/test_utils.h

@@ -148,6 +148,7 @@ class BaseFamilyTest : public ::testing::Test {
   static void ExpectConditionWithinTimeout(const std::function<bool()>& condition,
                                            absl::Duration timeout = absl::Seconds(10));
   util::fb2::Fiber ExpectConditionWithSuspension(const std::function<bool()>& condition);
+  util::fb2::Fiber ExpectUsedKeys(const std::vector<std::string_view>& keys);
 
   static unsigned NumLocked();
 

src/server/transaction.cc

@@ -84,7 +84,7 @@ uint16_t trans_id(const Transaction* ptr) {
 
 bool CheckLocks(const DbSlice& db_slice, IntentLock::Mode mode, const KeyLockArgs& lock_args) {
   for (size_t i = 0; i < lock_args.args.size(); i += lock_args.key_step) {
-    string_view s = KeyLockArgs::GetLockKey(lock_args.args[i]);
+    string_view s = lock_args.args[i];
     if (!db_slice.CheckLock(mode, lock_args.db_index, s))
       return false;
   }
@@ -262,7 +262,7 @@ void Transaction::LaunderKeyStorage(CmdArgVec* keys) {
   m_keys.reserve(keys->size());
 
   for (MutableSlice key : *keys) {
-    string_view key_s = KeyLockArgs::GetLockKey(facade::ToSV(key));
+    string_view key_s = string_view(LockTag{facade::ToSV(key)});
     // Insert copied string view, not original. This is why "try insert" is not allowed
     if (!m_keys_set.contains(key_s))
       m_keys_set.insert(m_keys.emplace_back(key_s));
@@ -633,7 +633,7 @@ bool Transaction::RunInShard(EngineShard* shard, bool txq_ooo) {
     if (auto* bcontroller = shard->blocking_controller(); bcontroller) {
       if (awaked_prerun || was_suspended) {
         CHECK_EQ(largs.key_step, 1u);
-        bcontroller->FinalizeWatched(largs.args, this);
+        bcontroller->FinalizeWatched(GetShardArgs(idx), this);
       }
 
       // Wake only if no tx queue head is currently running
@@ -1310,7 +1310,7 @@ void Transaction::UnlockMultiShardCb(absl::Span<const std::string_view> sharded_
     shard->shard_lock()->Release(IntentLock::EXCLUSIVE);
   } else {
     for (const auto& key : sharded_keys) {
-      shard->db_slice().ReleaseNormalized(*multi_->lock_mode, db_index_, key);
+      shard->db_slice().ReleaseNormalized(*multi_->lock_mode, db_index_, LockTag{key});
     }
   }