add big key detection

Author: caiyu

include/pika_admin.h

@@ -267,7 +267,9 @@ class InfoCmd : public Cmd {
     kInfoAll,
     kInfoDebug,
     kInfoCommandStats,
-    kInfoCache
+    kInfoCache,
+    kInfoBigKeys
+
   };
   InfoCmd(const std::string& name, int arity, uint32_t flag) : Cmd(name, arity, flag) {}
   void Do() override;
@@ -295,6 +297,7 @@ class InfoCmd : public Cmd {
   const static std::string kDebugSection;
   const static std::string kCommandStatsSection;
   const static std::string kCacheSection;
+  const static std::string kBigKeysSection;
 
   void DoInitial() override;
   void Clear() override {
@@ -304,6 +307,7 @@ class InfoCmd : public Cmd {
   }
 
   void InfoServer(std::string& info);
+  void InfoBigKeys(std::string& info);
   void InfoClients(std::string& info);
   void InfoStats(std::string& info);
   void InfoExecCount(std::string& info);

src/pika_admin.cc

@@ -884,7 +884,7 @@ const std::string InfoCmd::kRocksDBSection = "rocksdb";
 const std::string InfoCmd::kDebugSection = "debug";
 const std::string InfoCmd::kCommandStatsSection = "commandstats";
 const std::string InfoCmd::kCacheSection = "cache";
-
+const std::string InfoCmd::kBigKeysSection = "bigkeys";
 
 const std::string ClientCmd::KILLTYPE_NORMAL = "normal";
 const std::string ClientCmd::KILLTYPE_PUBSUB = "pubsub";
@@ -910,6 +910,8 @@ void InfoCmd::DoInitial() {
     keyspace_scan_dbs_ = g_pika_server->GetAllDBName();
   } else if (strcasecmp(argv_[1].data(), kServerSection.data()) == 0) {
     info_section_ = kInfoServer;
+  } else if (strcasecmp(argv_[1].data(), "bigkeys") == 0) {
+    info_section_ = kInfoBigKeys;
   } else if (strcasecmp(argv_[1].data(), kClientsSection.data()) == 0) {
     info_section_ = kInfoClients;
   } else if (strcasecmp(argv_[1].data(), kStatsSection.data()) == 0) {
@@ -994,6 +996,9 @@ void InfoCmd::Do() {
       InfoReplication(info);
       info.append("\r\n");
       InfoKeyspace(info);
+      info.append("\r\n");
+      InfoBigKeys(info);
+      info.append("\r\n");
       break;
     case kInfoAll:
       InfoServer(info);
@@ -1017,6 +1022,9 @@ void InfoCmd::Do() {
       InfoKeyspace(info);
       info.append("\r\n");
       InfoRocksDB(info);
+      info.append("\r\n");
+      InfoBigKeys(info);
+      info.append("\r\n");
       break;
     case kInfoServer:
       InfoServer(info);
@@ -1054,6 +1062,9 @@ void InfoCmd::Do() {
     case kInfoCache:
       InfoCache(info, db_);
       break;
+    case kInfoBigKeys:
+      InfoBigKeys(info);
+      break;
     default:
       // kInfoErr is nothing
       break;
@@ -1094,7 +1105,21 @@ void InfoCmd::InfoServer(std::string& info) {
 
   info.append(tmp_stream.str());
 }
-
+void InfoCmd::InfoBigKeys(std::string& info) {
+  std::stringstream tmp_stream;
+  std::shared_lock db_rwl(g_pika_server->dbs_rw_);
+  for (const auto& db_item : g_pika_server->dbs_) {
+    if (!db_item.second) {
+      continue;
+    }
+    std::vector<storage::BigKeyInfo> bigkeys;
+    db_item.second->storage()->GetBigKeyStatistics(&bigkeys);
+    std::string bigkey_info;
+    storage::FormatBigKeyStatistics(bigkeys, &bigkey_info);
+    tmp_stream << bigkey_info;
+  }
+  info.append(tmp_stream.str());
+}
 void InfoCmd::InfoClients(std::string& info) {
   std::stringstream tmp_stream;
   tmp_stream << "# Clients"

src/storage/include/storage/storage.h

@@ -127,7 +127,15 @@ struct ScoreMember {
 enum BeforeOrAfter { Before, After };
 
 enum DataType { kAll, kStrings, kHashes, kLists, kZSets, kSets, kStreams };
-
+//Big Keys INFO struct
+struct BigKeyInfo {
+  std::string key;
+  DataType type;
+  uint64_t key_size;
+  uint64_t value_size;
+  uint64_t last_update_time;
+};
+void FormatBigKeyStatistics(const std::vector<BigKeyInfo>& bigkeys, std::string* out);
 const char DataTypeTag[] = {'a', 'k', 'h', 'l', 'z', 's', 'x'};
 
 enum class OptionType {
@@ -1100,7 +1108,7 @@ class Storage {
   Status EnableAutoCompaction(const OptionType& option_type, 
                     const std::string& db_type, const std::unordered_map<std::string, std::string>& options);
   void GetRocksDBInfo(std::string& info);
-
+  void GetBigKeyStatistics(std::vector<BigKeyInfo>* bigkeys);
  private:
   std::unique_ptr<RedisStrings> strings_db_;
   std::unique_ptr<RedisHashes> hashes_db_;

src/storage/src/redis.cc

@@ -5,7 +5,9 @@
 
 #include "src/redis.h"
 #include <sstream>
-
+#include <thread>
+#include <fstream>
+#include <glog/logging.h>
 namespace storage {
 
 Redis::Redis(Storage* const s, const DataType& type)
@@ -20,6 +22,7 @@ Redis::Redis(Storage* const s, const DataType& type)
   default_compact_range_options_.exclusive_manual_compaction = false;
   default_compact_range_options_.change_level = true;
   handles_.clear();
+  StartBigKeysLogThread();
 }
 
 Redis::~Redis() {
@@ -33,6 +36,7 @@ Redis::~Redis() {
   if (default_compact_range_options_.canceled) {
     delete default_compact_range_options_.canceled;
   }
+  stop_big_keys_log_ = true;
 }
 
 Status Redis::GetScanStartPoint(const Slice& key, const Slice& pattern, int64_t cursor, std::string* start_point) {
@@ -109,7 +113,16 @@ Status Redis::SetOptions(const OptionType& option_type, const std::unordered_map
   }
   return s;
 }
-
+inline const char* DataTypeName(DataType type) {
+  switch (type) {
+    case kStrings: return "string";
+    case kHashes: return "hash";
+    case kLists: return "list";
+    case kSets: return "set";
+    case kZSets: return "zset";
+    default: return "unknown";
+  }
+}
 void Redis::GetRocksDBInfo(std::string &info, const char *prefix) {
     std::ostringstream string_stream;
     string_stream << "#" << prefix << "RocksDB" << "\r\n";
@@ -198,5 +211,47 @@ void Redis::SetCompactRangeOptions(const bool is_canceled) {
     default_compact_range_options_.canceled->store(is_canceled);
   } 
 }
-
+//big keys
+void Redis::CheckAndRecordBigKeys(const std::string& key, DataType type, uint64_t key_size, uint64_t value_size) {
+  bool is_bigkey = false;
+  if (type == kStrings) {
+    is_bigkey = (key_size >= BIGKEYS_THRESHOLD || value_size >= BIGKEYS_THRESHOLD);
+  } else {
+    is_bigkey = (key_size >= BIGKEYS_THRESHOLD);
+  }
+  if (is_bigkey) {
+    std::lock_guard<std::mutex> lock(big_keys_mutex_);
+    BigKeyInfo info;
+    info.key = key;
+    info.type = type;
+    info.key_size = key_size;
+    info.value_size = value_size;
+    info.last_update_time = static_cast<uint64_t>(time(nullptr));
+    big_keys_info_map_[key] = info;
+  }
+}
+void Redis::StartBigKeysLogThread() {
+  std::thread([this]() {
+    while (!stop_big_keys_log_) {
+      {
+        std::lock_guard<std::mutex> lock(big_keys_mutex_);
+        std::ofstream ofs("log/pika.WARNING", std::ios::app);
+        if (ofs.is_open()) {
+          for (const auto& kv : big_keys_info_map_) {
+            ofs << "[BigKey] key=" << kv.first
+                << " type=" << DataTypeName(kv.second.type)
+                << " key_size=" << kv.second.key_size 
+                << " value_size="<< kv.second.value_size
+                << " Updated Time=" << kv.second.last_update_time
+                << std::endl;
+          }
+          ofs.close();
+        } else {
+          LOG(WARNING) << "Failed to open log/pika.WARNING for writing!";
+        }
+      }
+      std::this_thread::sleep_for(std::chrono::minutes(1));
+    }
+  }).detach();
+}
 }  // namespace storage

src/storage/src/redis.h

@@ -5,7 +5,7 @@
 
 #ifndef SRC_REDIS_H_
 #define SRC_REDIS_H_
-
+#define BIGKEYS_THRESHOLD 1
 #include <memory>
 #include <string>
 #include <vector>
@@ -24,7 +24,8 @@
 namespace storage {
 using Status = rocksdb::Status;
 using Slice = rocksdb::Slice;
-
+//Big Keys INFO struct
+struct BigKeyInfo;
 class Redis {
  public:
   Redis(Storage* storage, const DataType& type);
@@ -115,7 +116,7 @@ class Redis {
   Status SetSmallCompactionDurationThreshold(uint64_t small_compaction_duration_threshold);
   std::vector<rocksdb::ColumnFamilyHandle*> GetHandles(){ return handles_;};
   void GetRocksDBInfo(std::string &info, const char *prefix);
-
+  virtual void GetBigKeyStatistics(std::vector<BigKeyInfo>* bigkeys) = 0;
  protected:
   Storage* const storage_;
   DataType type_;
@@ -126,6 +127,13 @@ class Redis {
   rocksdb::WriteOptions default_write_options_;
   rocksdb::ReadOptions default_read_options_;
   rocksdb::CompactRangeOptions default_compact_range_options_;
+  // big keys
+  std::mutex big_keys_mutex_;
+  std::unordered_map<std::string, BigKeyInfo> big_keys_info_map_;
+  uint64_t big_keys_threshold_;
+  std::atomic<bool> stop_big_keys_log_ = false;
+  void CheckAndRecordBigKeys(const std::string& key, DataType type, uint64_t value_size, uint64_t key_size = 0);
+  void StartBigKeysLogThread();
 
   // For Scan
   std::unique_ptr<LRUCache<std::string, std::string>> scan_cursors_store_;

src/storage/src/redis_hashes.cc

@@ -421,6 +421,11 @@ Status RedisHashes::HIncrby(const Slice& key, const Slice& field, int64_t value,
   }
   s = db_->Write(default_write_options_, &batch);
   UpdateSpecificKeyStatistics(key.ToString(), statistic);
+  int32_t count = 0;
+  Status cnt_s = HLen(key, &count);
+  if (cnt_s.ok()) {
+    CheckAndRecordBigKeys(key.ToString(), kHashes, count);
+  }
   return s;
 }
 
@@ -495,6 +500,11 @@ Status RedisHashes::HIncrbyfloat(const Slice& key, const Slice& field, const Sli
   }
   s = db_->Write(default_write_options_, &batch);
   UpdateSpecificKeyStatistics(key.ToString(), statistic);
+  int32_t count = 0;
+  Status cnt_s = HLen(key, &count);
+  if (cnt_s.ok()) {
+    CheckAndRecordBigKeys(key.ToString(), kHashes, count);
+  }
   return s;
 }
 
@@ -659,6 +669,11 @@ Status RedisHashes::HMSet(const Slice& key, const std::vector<FieldValue>& fvs)
   }
   s = db_->Write(default_write_options_, &batch);
   UpdateSpecificKeyStatistics(key.ToString(), statistic);
+  int32_t count = 0;
+  Status cnt_s = HLen(key, &count);
+  if (cnt_s.ok()) {
+    CheckAndRecordBigKeys(key.ToString(), kHashes, count);
+  }
   return s;
 }
 
@@ -718,6 +733,11 @@ Status RedisHashes::HSet(const Slice& key, const Slice& field, const Slice& valu
   }
   s = db_->Write(default_write_options_, &batch);
   UpdateSpecificKeyStatistics(key.ToString(), statistic);
+  int32_t count = 0;
+  Status cnt_s = HLen(key, &count);
+  if (cnt_s.ok()) {
+    CheckAndRecordBigKeys(key.ToString(), kHashes, count);
+  }
   return s;
 }
 
@@ -768,9 +788,22 @@ Status RedisHashes::HSetnx(const Slice& key, const Slice& field, const Slice& va
   } else {
     return s;
   }
-  return db_->Write(default_write_options_, &batch);
+  s = db_->Write(default_write_options_, &batch);
+  int32_t count = 0;
+  Status cnt_s = HLen(key, &count);
+  if (cnt_s.ok()) {
+    CheckAndRecordBigKeys(key.ToString(), kHashes, count);
+  }
+  return s;
+}
+void RedisHashes::GetBigKeyStatistics(std::vector<BigKeyInfo>* bigkeys) {
+  std::lock_guard<std::mutex> lock(big_keys_mutex_);
+  for (const auto& kv : big_keys_info_map_) {
+    BigKeyInfo info = kv.second;
+    info.key = kv.first;
+    bigkeys->push_back(info);
+  }
 }
-
 Status RedisHashes::HVals(const Slice& key, std::vector<std::string>* values) {
   rocksdb::ReadOptions read_options;
   const rocksdb::Snapshot* snapshot;

src/storage/src/redis_hashes.h

@@ -70,6 +70,7 @@ class RedisHashes : public Redis {
 
   // Iterate all data
   void ScanDatabase();
+  void GetBigKeyStatistics(std::vector<BigKeyInfo>* bigkeys);
 };
 
 }  //  namespace storage