diff --git a/src/core/heap_size.h b/src/core/heap_size.h index b1184ffbf..eae680c40 100644 --- a/src/core/heap_size.h +++ b/src/core/heap_size.h @@ -29,8 +29,14 @@ namespace dfly { namespace heap_size_detail { +template struct has_marked_stackonly : std::false_type {}; + +template +struct has_marked_stackonly> : std::true_type {}; + template constexpr bool StackOnlyType() { - return std::is_trivial_v || std::is_same_v; + return std::is_trivial_v || std::is_same_v || + has_marked_stackonly::value; } template struct has_used_mem : std::false_type {}; diff --git a/src/server/cluster/cluster_config.cc b/src/server/cluster/cluster_config.cc index 1cbddc92c..0138c68ae 100644 --- a/src/server/cluster/cluster_config.cc +++ b/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; } diff --git a/src/server/cluster/cluster_config.h b/src/server/cluster/cluster_config.h index 3c5844572..e1b64617e 100644 --- a/src/server/cluster/cluster_config.h +++ b/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 CreateFromConfig(std::string_view my_id, diff --git a/src/server/cluster/cluster_config_test.cc b/src/server/cluster/cluster_config_test.cc index f02695854..7b37f7f0d 100644 --- a/src/server/cluster/cluster_config_test.cc +++ b/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(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"); @@ -71,19 +75,19 @@ TEST_F(ClusterConfigTest, KeyTagTest) { 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) { diff --git a/src/server/common.cc b/src/server/common.cc index 9df483c16..4f6084261 100644 --- a/src/server/common.cc +++ b/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 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; } - 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 diff --git a/src/server/common.h b/src/server/common.h index 9dd6344d0..7e344db6f 100644 --- a/src/server/common.h +++ b/src/server/common.h @@ -48,6 +48,7 @@ using RdbTypeFreqMap = absl::flat_hash_map; 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 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); diff --git a/src/server/conn_context.cc b/src/server/conn_context.cc index 8763ecf16..55152c75b 100644 --- a/src/server/conn_context.cc +++ b/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 { diff --git a/src/server/conn_context.h b/src/server/conn_context.h index 1315a7076..ba94a117c 100644 --- a/src/server/conn_context.h +++ b/src/server/conn_context.h @@ -100,7 +100,7 @@ struct ConnectionState { struct ScriptInfo { size_t UsedMemory() const; - absl::flat_hash_set keys; // declared keys + absl::flat_hash_set 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 diff --git a/src/server/db_slice.cc b/src/server/db_slice.cc index daf2cb186..3614179de 100644 --- a/src/server/db_slice.cc +++ b/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) { diff --git a/src/server/db_slice.h b/src/server/db_slice.h index 3fb67d0b3..ccbedd258 100644 --- a/src/server/db_slice.h +++ b/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); diff --git a/src/server/engine_shard_set.cc b/src/server/engine_shard_set.cc index 3236081dd..07a6d208c 100644 --- a/src/server/engine_shard_set.cc +++ b/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) { - auto is_contended = [table](string_view key) { - return table->trans_locks.Find(key)->IsContended(); - }; +bool HasContendedLocks(ShardId shard_id, Transaction* trx, const DbTable* table) { + auto is_contended = [table](LockTag tag) { return table->trans_locks.Find(tag)->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); diff --git a/src/server/main_service.cc b/src/server/main_service.cc index b8f5ccebe..5fc53de25 100644 --- a/src/server/main_service.cc +++ b/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)); - if (!locked_keys.contains(key)) { - VLOG(1) << "Key " << key << " is not declared for command " << cid->name(); + LockTag tag{ArgS(args, i)}; + if (!locked_tags.contains(tag)) { + VLOG(1) << "Key " << string_view(tag) << " is not declared for command " << cid->name(); return OpStatus::KEY_NOTFOUND; } } - if (key_index.bonus && - !locked_keys.contains(KeyLockArgs::GetLockKey(ArgS(args, *key_index.bonus)))) + if (key_index.bonus && !locked_tags.contains(LockTag{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(); - sinfo->keys.reserve(eval_args.keys.size()); + sinfo->lock_tags.reserve(eval_args.keys.size()); optional 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) { diff --git a/src/server/string_family_test.cc b/src/server/string_family_test.cc index 1f85207d3..1bd26f45b 100644 --- a/src/server/string_family_test.cc +++ b/src/server/string_family_test.cc @@ -791,20 +791,16 @@ TEST_F(StringFamilyTest, SetWithHashtagsNoCluster) { SetTestFlag("lock_on_hashtags", "false"); ResetService(); - auto fb = ExpectConditionWithSuspension( - [&]() { return GetLastUsedKeys() == absl::flat_hash_set{"{key}1"}; }); + auto fb = ExpectUsedKeys({"{key}1"}); EXPECT_EQ(Run({"set", "{key}1", "val1"}), "OK"); fb.Join(); EXPECT_FALSE(service_->IsLocked(0, "{key}1")); - fb = ExpectConditionWithSuspension( - [&]() { return GetLastUsedKeys() == absl::flat_hash_set{"{key}2"}; }); + fb = ExpectUsedKeys({"{key}2"}); EXPECT_EQ(Run({"set", "{key}2", "val2"}), "OK"); fb.Join(); - fb = ExpectConditionWithSuspension([&]() { - return GetLastUsedKeys() == absl::flat_hash_set{"{key}1", "{key}2"}; - }); + fb = ExpectUsedKeys({"{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( - [&]() { return GetLastUsedKeys() == absl::flat_hash_set{"{key}1"}; }); + auto fb = ExpectUsedKeys({"{key}1"}); EXPECT_EQ(Run({"set", "{key}1", "val1"}), "OK"); fb.Join(); - fb = ExpectConditionWithSuspension( - [&]() { return GetLastUsedKeys() == absl::flat_hash_set{"{key}2"}; }); + fb = ExpectUsedKeys({"{key}2"}); EXPECT_EQ(Run({"set", "{key}2", "val2"}), "OK"); fb.Join(); - fb = ExpectConditionWithSuspension([&]() { - return GetLastUsedKeys() == absl::flat_hash_set{"{key}1", "{key}2"}; - }); + fb = ExpectUsedKeys({"{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{"key"}; }; - auto fb = ExpectConditionWithSuspension(condition); + auto fb = ExpectUsedKeys({"key"}); 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 = [&]() { - return GetLastUsedKeys() == absl::flat_hash_set{"{key}1", "{key}2", "{key}3"}; - }; - auto fb = ExpectConditionWithSuspension(condition); + auto fb = ExpectUsedKeys({"{key}1", "{key}2", "{key}3"}); 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{"key"}; }; - auto fb = ExpectConditionWithSuspension(condition); + auto fb = ExpectUsedKeys({"key"}); EXPECT_EQ(Run({"multi"}), "OK"); EXPECT_EQ(Run({"set", "{key}1", "val1"}), "QUEUED"); diff --git a/src/server/table.cc b/src/server/table.cc index 4cef8f635..9b06b66c9 100644 --- a/src/server/table.cc +++ b/src/server/table.cc @@ -4,8 +4,6 @@ #include "server/table.h" -#include - #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 { - return locks_.size(); -} - -LockFp LockTable::Fingerprint(string_view key) { - return XXH64(key.data(), key.size(), 0x1C69B3F74AC4AE35UL); -} - -std::optional LockTable::Find(string_view key) const { - DCHECK_EQ(KeyLockArgs::GetLockKey(key), key); - - LockFp fp = Fingerprint(key); +std::optional LockTable::Find(LockTag tag) const { + LockFp fp = tag.Fingerprint(); if (auto it = locks_.find(fp); it != locks_.end()) return it->second; return std::nullopt; } -bool LockTable::Acquire(string_view key, IntentLock::Mode mode) { - DCHECK_EQ(KeyLockArgs::GetLockKey(key), key); - LockFp fp = Fingerprint(key); +bool LockTable::Acquire(LockTag tag, IntentLock::Mode mode) { + LockFp fp = tag.Fingerprint(); auto [it, inserted] = locks_.try_emplace(fp); return it->second.Acquire(mode); } -void LockTable::Release(string_view key, IntentLock::Mode mode) { - DCHECK_EQ(KeyLockArgs::GetLockKey(key), key); - - LockFp fp = Fingerprint(key); +void LockTable::Release(LockTag tag, IntentLock::Mode mode) { + LockFp fp = tag.Fingerprint(); 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()) diff --git a/src/server/table.h b/src/server/table.h index 50701f778..896dc1e14 100644 --- a/src/server/table.h +++ b/src/server/table.h @@ -25,7 +25,6 @@ using PrimeValue = detail::PrimeValue; using PrimeTable = DashTable; using ExpireTable = DashTable; -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 -// strings for locks that are used by multiple transactions. Keys used with the lock table -// should be normalized with GetLockKey +// Table for recording locks. Keys used with the lock table should be normalized with LockTag. class LockTable { public: - size_t Size() const; - std::optional Find(std::string_view key) const; + size_t Size() const { + return locks_.size(); + } + std::optional Find(LockTag tag) const; + std::optional Find(LockFp fp) const; - bool Acquire(std::string_view key, IntentLock::Mode mode); - void Release(std::string_view key, IntentLock::Mode mode); - - static LockFp Fingerprint(std::string_view key); + bool Acquire(LockTag tag, IntentLock::Mode mode); + void Release(LockTag tag, IntentLock::Mode mode); auto begin() const { return locks_.cbegin(); diff --git a/src/server/test_utils.cc b/src/server/test_utils.cc index e832dddf7..3b3197755 100644 --- a/src/server/test_utils.cc +++ b/src/server/test_utils.cc @@ -677,6 +677,19 @@ fb2::Fiber BaseFamilyTest::ExpectConditionWithSuspension(const std::function& keys) { + absl::flat_hash_set 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); diff --git a/src/server/test_utils.h b/src/server/test_utils.h index 6b47324b3..16c031c52 100644 --- a/src/server/test_utils.h +++ b/src/server/test_utils.h @@ -148,6 +148,7 @@ class BaseFamilyTest : public ::testing::Test { static void ExpectConditionWithinTimeout(const std::function& condition, absl::Duration timeout = absl::Seconds(10)); util::fb2::Fiber ExpectConditionWithSuspension(const std::function& condition); + util::fb2::Fiber ExpectUsedKeys(const std::vector& keys); static unsigned NumLocked(); diff --git a/src/server/transaction.cc b/src/server/transaction.cc index 79d6eaafa..476b792b3 100644 --- a/src/server/transaction.cc +++ b/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 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}); } }