Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 16 additions & 3 deletions src/pika_kv.cc
Original file line number Diff line number Diff line change
Expand Up @@ -206,13 +206,26 @@ void DelCmd::DoInitial() {
}

void DelCmd::Do() {
// 先获取所有键的类型
std::map<std::string, std::string> key_type_map;
for (const auto& key : keys_) {
std::string type;
int result = GetKeyType(key, type, db_);
if (result > 0) {
key_type_map[key] = type;
}
}

// 然后删除键
int64_t count = db_->storage()->Del(keys_);

if (count >= 0) {
res_.AppendInteger(count);
s_ = rocksdb::Status::OK();
std::vector<std::string>::const_iterator it;
for (it = keys_.begin(); it != keys_.end(); it++) {
RemSlotKey(*it, db_);

// 使用预先获取的类型信息从 slotKey 中移除
for (const auto& [key, type] : key_type_map) {
RemSlotKeyByType(type, key, db_);
}
Comment on lines 208 to 229
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Race condition between type lookup and deletion; consider using the new Del overload.

There's a TOCTOU (time-of-check-to-time-of-use) race: between GetKeyType (lines 211-217) and Del (line 220), another thread could delete or modify the key, causing:

  1. A key to be deleted but not removed from the slot (if deleted between GetKeyType and Del)
  2. Incorrect type being used for slot removal (if key type changed)

Additionally, the new Storage::Del(keys, &key_types) overload (added in this PR) collects types atomically with deletion. Consider using it to eliminate the race and the redundant type lookups:

🔧 Suggested fix using the new Del overload
 void DelCmd::Do() {
-  // 先获取所有键的类型
-  std::map<std::string, std::string> key_type_map;
-  for (const auto& key : keys_) {
-    std::string type;
-    int result = GetKeyType(key, type, db_);
-    if (result > 0) {
-      key_type_map[key] = type;
-    }
-  }
-  
-  // 然后删除键
-  int64_t count = db_->storage()->Del(keys_);
+  // Delete keys and collect their types atomically
+  std::vector<std::pair<std::string, std::string>> key_types;
+  int64_t count = db_->storage()->Del(keys_, &key_types);
   
   if (count >= 0) {
     res_.AppendInteger(count);
     s_ = rocksdb::Status::OK();
     
-    // 使用预先获取的类型信息从 slotKey 中移除
-    for (const auto& [key, type] : key_type_map) {
-      RemSlotKeyByType(type, key, db_);
+    // Remove from slotKey using the collected type information
+    for (const auto& [key, type] : key_types) {
+      if (!type.empty()) {
+        RemSlotKeyByType(type, key, db_);
+      }
     }
   } else {
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/pika_kv.cc` around lines 208 - 229, DelCmd::Do currently does a separate
GetKeyType loop then calls db_->storage()->Del(keys_) which causes a TOCTOU
race; replace that two-step approach by calling the new Storage::Del(keys_,
&key_types) overload (collecting deletion count and types atomically) instead of
GetKeyType and db_->storage()->Del(keys_); then use the returned key_types map
to call RemSlotKeyByType(type, key, db_) for each entry, keep the existing
res_.AppendInteger(count) and s_ = rocksdb::Status::OK() semantics, and remove
the old GetKeyType loop and the separate Del call.

} else {
res_.SetRes(CmdRes::kErrOther, "delete error");
Expand Down
13 changes: 4 additions & 9 deletions src/pika_slot_command.cc
Original file line number Diff line number Diff line change
Expand Up @@ -635,15 +635,15 @@ void RemSlotKeyByType(const std::string& type, const std::string& key, const std
std::vector<std::string> members;
members.emplace_back(type + key);
rocksdb::Status s = db->storage()->SRem(slot_key, members, &res);
if (!s.ok()) {
if (!s.ok() && !s.IsNotFound()) {
LOG(ERROR) << "srem key[" << key << "] from slotKey[" << slot_key << "] failed, error: " << s.ToString();
return;
}

if (hastag) {
std::string tag_key = GetSlotsTagKey(crc);
s = db->storage()->SRem(tag_key, members, &res);
if (!s.ok()) {
if (!s.ok() && !s.IsNotFound()) {
LOG(ERROR) << "srem key[" << key << "] from tagKey[" << tag_key << "] failed, error: " << s.ToString();
return;
}
Expand Down Expand Up @@ -761,17 +761,12 @@ void RemSlotKey(const std::string& key, const std::shared_ptr<DB>& db) {
int GetKeyType(const std::string& key, std::string& key_type, const std::shared_ptr<DB>& db) {
enum storage::DataType type;
rocksdb::Status s = db->storage()->GetType(key, type);
if (!s.ok()) {
LOG(WARNING) << "Get key type error: " << key << " " << s.ToString();
if (!s.ok() || type == storage::DataType::kNones) {
LOG(WARNING) << "Get key type error: " << key << " type: " << static_cast<int>(type);
key_type = "";
return -1;
}
auto key_type_char = storage::DataTypeToTag(type);
if (key_type_char == DataTypeToTag(storage::DataType::kNones)) {
LOG(WARNING) << "Get key type error: " << key;
key_type = "";
return -1;
}
key_type = key_type_char;
return 1;
}
Expand Down
1 change: 0 additions & 1 deletion src/pstd/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ cmake_minimum_required(VERSION 3.18)
set (CMAKE_CXX_STANDARD 17)
project (pstd)

# 强制使用用户自定的 memcmp
add_compile_options("-fno-builtin-memcmp -pipe")


Expand Down
3 changes: 3 additions & 0 deletions src/storage/include/storage/storage.h
Original file line number Diff line number Diff line change
Expand Up @@ -1012,6 +1012,9 @@ class Storage {
int64_t Exists(const std::vector<std::string>& keys);

// Return the key exists type count
// Del keys and return deleted keys' types
int64_t Del(const std::vector<std::string>& keys, std::vector<std::pair<std::string, std::string>>* key_types);

// return param type_status: return every type status
int64_t IsExist(const Slice& key, std::map<DataType, Status>* type_status);

Expand Down
17 changes: 16 additions & 1 deletion src/storage/src/storage.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1193,11 +1193,22 @@ int32_t Storage::Expire(const Slice& key, int64_t ttl_millsec) {
}


int64_t Storage::Del(const std::vector<std::string>& keys) {
int64_t Storage::Del(const std::vector<std::string>& keys, std::vector<std::pair<std::string, std::string>>* key_types) {
Status s;
int64_t count = 0;
for (const auto& key : keys) {
auto& inst = GetDBInstance(key);
// 先获取类型
if (key_types != nullptr) {
enum DataType type;
inst->GetType(key, type);
if (type != DataType::kNones) {
key_types->emplace_back(key, std::string(1, DataTypeToTag(type)));
} else {
key_types->emplace_back(key, "");
}
}
// 删除键
s = inst->Del(key);
if (s.ok()) {
count++;
Expand All @@ -1206,6 +1217,10 @@ int64_t Storage::Del(const std::vector<std::string>& keys) {
return count;
}

int64_t Storage::Del(const std::vector<std::string>& keys) {
return Del(keys, nullptr);
}

int64_t Storage::Exists(const std::vector<std::string>& keys) {
int64_t count = 0;
Status s;
Expand Down
Loading