feat: support blpop/brpop

.github/workflows/pika.yml

@@ -69,7 +69,7 @@ jobs:
         working-directory: ${{github.workspace}}/build
         run: |
           python3 ../tests/integration/pika_replication_test.py
-
+          python3 ../tests/unit/Blpop_Brpop_test.py
 
   build_on_centos:
     # The CMake configure and build commands are platform agnostic and should work equally well on Windows or Mac.

CMakeLists.txt

@@ -22,13 +22,17 @@ elseif (CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
   endif()
 endif()
 
-#Sanitizer for testing stage
-#CMAKE_BUILD_TYPE must be "Debug" if you wanna use sanitizer.
+
+############# You should enable sanitizer if you are developing pika #############
+# Uncomment the following two lines to enable AddressSanitizer to detect memory leaks and other memory-related bugs.
 #set(CMAKE_BUILD_TYPE "Debug")
+#set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fsanitize=address  -O0 -fno-omit-frame-pointer -fno-optimize-sibling-calls")
+
+# [Notice] AddressSanitizer and ThreadSanitizer can not be enabled at the same time.
 
-#set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fsanitize=address")
-# [Notice] address sanitizer and thread sanitizer can not be enabled at the same time.
-#set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fsanitize=thread")
+# Uncomment the following two lines to enable ThreadSanitizer to detect data race and other thread-related issue.
+#set(CMAKE_BUILD_TYPE "Debug")
+#set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fsanitize=thread -O0 -fno-omit-frame-pointer -fno-optimize-sibling-calls")
 
 execute_process(COMMAND uname -p OUTPUT_VARIABLE HOST_ARCH)
 string(TOLOWER ${HOST_ARCH} HOST_ARCH)

include/pika_client_conn.h

@@ -55,6 +55,7 @@ class PikaClientConn : public net::RedisConn {
   bool IsPubSub() { return is_pubsub_; }
   void SetIsPubSub(bool is_pubsub) { is_pubsub_ = is_pubsub; }
   void SetCurrentTable(const std::string& db_name) { current_db_ = db_name; }
+  const std::string& GetCurrentTable() override{ return current_db_; }
   void SetWriteCompleteCallback(WriteCompleteCallback cb) { write_completed_cb_ = std::move(cb); }
 
   net::ServerThread* server_thread() { return server_thread_; }

include/pika_command.h

@@ -16,6 +16,7 @@
 #include "pstd/include/pstd_string.h"
 
 #include "include/pika_slot.h"
+#include "net/src/dispatch_thread.h"
 
 class SyncMasterSlot;
 class SyncSlaveSlot;
@@ -131,13 +132,15 @@ const std::string kCmdNamePKHRScanRange = "pkhrscanrange";
 const std::string kCmdNameLIndex = "lindex";
 const std::string kCmdNameLInsert = "linsert";
 const std::string kCmdNameLLen = "llen";
+const std::string kCmdNameBLPop = "blpop";
 const std::string kCmdNameLPop = "lpop";
 const std::string kCmdNameLPush = "lpush";
 const std::string kCmdNameLPushx = "lpushx";
 const std::string kCmdNameLRange = "lrange";
 const std::string kCmdNameLRem = "lrem";
 const std::string kCmdNameLSet = "lset";
 const std::string kCmdNameLTrim = "ltrim";
+const std::string kCmdNameBRpop = "brpop";
 const std::string kCmdNameRPop = "rpop";
 const std::string kCmdNameRPopLPush = "rpoplpush";
 const std::string kCmdNameRPush = "rpush";
@@ -403,6 +406,18 @@ class CmdRes {
   CmdRet ret_ = kNone;
 };
 
+/**
+ * Current used by:
+ * blpop,brpop
+ */
+struct UnblockTaskArgs {
+  std::string key;
+  std::shared_ptr<Slot> slot;
+  net::DispatchThread* dispatchThread{ nullptr };
+  UnblockTaskArgs(std::string key_, std::shared_ptr<Slot> slot_, net::DispatchThread* dispatchThread_)
+      : key(std::move(key_)), slot(slot_), dispatchThread(dispatchThread_) {}
+};
+
 class Cmd : public std::enable_shared_from_this<Cmd> {
  public:
   enum CmdStage { kNone, kBinlogStage, kExecuteStage };
@@ -444,6 +459,7 @@ class Cmd : public std::enable_shared_from_this<Cmd> {
 
   bool is_read() const;
   bool is_write() const;
+
   bool is_local() const;
   bool is_suspend() const;
   bool is_admin_require() const;

include/pika_list.h

@@ -72,6 +72,49 @@ class LLenCmd : public Cmd {
   void DoInitial() override;
 };
 
+class BlockingBaseCmd : public Cmd {
+ public:
+  BlockingBaseCmd(const std::string& name, int arity, uint16_t flag) : Cmd(name, arity, flag) {}
+
+  //blpop/brpop used start
+  struct WriteBinlogOfPopArgs{
+    BlockKeyType block_type;
+    std::string key;
+    std::shared_ptr<Slot> slot;
+    std::shared_ptr<net::NetConn> conn;
+    WriteBinlogOfPopArgs() = default;
+    WriteBinlogOfPopArgs(BlockKeyType block_type_, const std::string& key_,
+                         std::shared_ptr<Slot> slot_, std::shared_ptr<net::NetConn> conn_)
+        : block_type(block_type_), key(key_), slot(slot_), conn(conn_){}
+  };
+  void BlockThisClientToWaitLRPush(BlockKeyType block_pop_type, std::vector<std::string>& keys, int64_t expire_time);
+  void TryToServeBLrPopWithThisKey(const std::string& key, std::shared_ptr<Slot> slot);
+  static void ServeAndUnblockConns(void* args);
+  static void WriteBinlogOfPop(std::vector<WriteBinlogOfPopArgs>& pop_args);
+  void removeDuplicates(std::vector<std::string> & keys_);
+  //blpop/brpop used functions end
+};
+
+class BLPopCmd final : public BlockingBaseCmd {
+ public:
+  BLPopCmd(const std::string& name, int arity, uint16_t flag) : BlockingBaseCmd(name, arity, flag){};
+  virtual std::vector<std::string> current_key() const {
+    return { keys_ };
+  }
+  virtual void Do(std::shared_ptr<Slot> slot = nullptr);
+  virtual void Split(std::shared_ptr<Slot> slot, const HintKeys& hint_keys){};
+  virtual void Merge(){};
+  virtual Cmd* Clone() override { return new BLPopCmd(*this); }
+  void DoInitial() override;
+  void DoBinlog(const std::shared_ptr<SyncMasterSlot>& slot) override;
+
+ private:
+  std::vector<std::string> keys_;
+  int64_t expire_time_{0};
+  WriteBinlogOfPopArgs binlog_args_;
+  bool is_binlog_deferred_{false};
+};
+
 class LPopCmd : public Cmd {
  public:
   LPopCmd(const std::string& name, int arity, uint16_t flag) : Cmd(name, arity, flag){};
@@ -91,9 +134,10 @@ class LPopCmd : public Cmd {
   void DoInitial() override;
 };
 
-class LPushCmd : public Cmd {
+
+class LPushCmd : public BlockingBaseCmd {
  public:
-  LPushCmd(const std::string& name, int arity, uint16_t flag) : Cmd(name, arity, flag){};
+  LPushCmd(const std::string& name, int arity, uint16_t flag) : BlockingBaseCmd(name, arity, flag){};
   std::vector<std::string> current_key() const override {
     std::vector<std::string> res;
     res.push_back(key_);
@@ -210,6 +254,25 @@ class LTrimCmd : public Cmd {
   void DoInitial() override;
 };
 
+class BRPopCmd final : public BlockingBaseCmd {
+ public:
+  BRPopCmd(const std::string& name, int arity, uint16_t flag) : BlockingBaseCmd(name, arity, flag){};
+  virtual std::vector<std::string> current_key() const {
+    return { keys_ };
+  }
+  virtual void Do(std::shared_ptr<Slot> slot = nullptr);
+  virtual void Split(std::shared_ptr<Slot> slot, const HintKeys& hint_keys){};
+  virtual void Merge(){};
+  virtual Cmd* Clone() override { return new BRPopCmd(*this); }
+  void DoInitial() override;
+  void DoBinlog(const std::shared_ptr<SyncMasterSlot>& slot) override;
+ private:
+  std::vector<std::string> keys_;
+  int64_t expire_time_{0};
+  WriteBinlogOfPopArgs binlog_args_;
+  bool is_binlog_deferred_{false};
+};
+
 class RPopCmd : public Cmd {
  public:
   RPopCmd(const std::string& name, int arity, uint16_t flag) : Cmd(name, arity, flag){};
@@ -229,14 +292,14 @@ class RPopCmd : public Cmd {
   void DoInitial() override;
 };
 
-class RPopLPushCmd : public Cmd {
+class RPopLPushCmd : public BlockingBaseCmd {
  public:
-  RPopLPushCmd(const std::string& name, int arity, uint16_t flag) : Cmd(name, arity, flag) {
+  RPopLPushCmd(const std::string& name, int arity, uint16_t flag) : BlockingBaseCmd(name, arity, flag) {
     rpop_cmd_ = std::make_shared<RPopCmd>(kCmdNameRPop, 2, kCmdFlagsWrite | kCmdFlagsSingleSlot | kCmdFlagsList);
     lpush_cmd_ = std::make_shared<LPushCmd>(kCmdNameLPush, -3, kCmdFlagsWrite | kCmdFlagsSingleSlot | kCmdFlagsList);
   };
   RPopLPushCmd(const RPopLPushCmd& other)
-      : Cmd(other),
+      : BlockingBaseCmd(other),
         source_(other.source_),
         receiver_(other.receiver_),
         value_poped_from_source_(other.value_poped_from_source_),
@@ -267,9 +330,9 @@ class RPopLPushCmd : public Cmd {
   void DoInitial() override;
 };
 
-class RPushCmd : public Cmd {
+class RPushCmd : public BlockingBaseCmd {
  public:
-  RPushCmd(const std::string& name, int arity, uint16_t flag) : Cmd(name, arity, flag){};
+  RPushCmd(const std::string& name, int arity, uint16_t flag) : BlockingBaseCmd(name, arity, flag){};
   std::vector<std::string> current_key() const override {
     std::vector<std::string> res;
     res.push_back(key_);

src/net/include/redis_conn.h

@@ -39,6 +39,7 @@ class RedisConn : public NetConn {
   void NotifyEpoll(bool success);
 
   virtual int DealMessage(const RedisCmdArgsType& argv, std::string* response) = 0;
+  virtual const std::string& GetCurrentTable() = 0;
 
  private:
   static int ParserDealMessageCb(RedisParser* parser, const RedisCmdArgsType& argv);

src/net/include/server_thread.h

@@ -20,6 +20,7 @@
 #include "net/include/net_define.h"
 #include "net/include/net_thread.h"
 #include "net/src/net_multiplexer.h"
+#include "pstd/include/env.h"
 #include "pstd/include/pstd_mutex.h"
 #include "pstd/include/pstd_status.h"
 

src/net/src/dispatch_thread.cc

@@ -19,22 +19,23 @@ DispatchThread::DispatchThread(int port, int work_num, ConnFactory* conn_factory
     : ServerThread::ServerThread(port, cron_interval, handle),
       last_thread_(0),
       work_num_(work_num),
-      queue_limit_(queue_limit) {
+      queue_limit_(queue_limit){
   for (int i = 0; i < work_num_; i++) {
     worker_thread_.emplace_back(std::make_unique<WorkerThread>(conn_factory, this, queue_limit, cron_interval));
   }
+  timed_scan_thread.SetTimedTask(0.3, [this]{this->ScanExpiredBlockedConnsOfBlrpop();});
 }
 
 DispatchThread::DispatchThread(const std::string& ip, int port, int work_num, ConnFactory* conn_factory,
                                int cron_interval, int queue_limit, const ServerHandle* handle)
     : ServerThread::ServerThread(ip, port, cron_interval, handle),
       last_thread_(0),
       work_num_(work_num),
-      queue_limit_(queue_limit) {
-
+      queue_limit_(queue_limit){
   for (int i = 0; i < work_num_; i++) {
     worker_thread_.emplace_back(std::make_unique<WorkerThread>(conn_factory, this, queue_limit, cron_interval));
   }
+  timed_scan_thread.SetTimedTask(0.3, [this]{this->ScanExpiredBlockedConnsOfBlrpop();});
 }
 
 DispatchThread::DispatchThread(const std::set<std::string>& ips, int port, int work_num, ConnFactory* conn_factory,
@@ -43,11 +44,10 @@ DispatchThread::DispatchThread(const std::set<std::string>& ips, int port, int w
       last_thread_(0),
       work_num_(work_num),
       queue_limit_(queue_limit) {
-
   for (int i = 0; i < work_num_; i++) {
     worker_thread_.emplace_back(std::make_unique<WorkerThread>(conn_factory, this, queue_limit, cron_interval));
-
   }
+  timed_scan_thread.SetTimedTask(0.3, [this]{this->ScanExpiredBlockedConnsOfBlrpop();});
 }
 
 DispatchThread::~DispatchThread() = default;
@@ -67,6 +67,7 @@ int DispatchThread::StartThread() {
       return ret;
     }
   }
+  timed_scan_thread.StartThread();
   return ServerThread::StartThread();
 }
 
@@ -87,6 +88,7 @@ int DispatchThread::StopThread() {
       worker_thread_[i]->private_data_ = nullptr;
     }
   }
+  timed_scan_thread.StopThread();
   return ServerThread::StopThread();
 }
 
@@ -168,6 +170,92 @@ void DispatchThread::HandleNewConn(const int connfd, const std::string& ip_port)
   }
 }
 
+bool BlockedConnNode::IsExpired() {
+  if (expire_time_ == 0) {
+    return false;
+  }
+  auto now = std::chrono::system_clock::now();
+  int64_t now_in_ms = std::chrono::time_point_cast<std::chrono::milliseconds>(now).time_since_epoch().count();
+  if (expire_time_ <= now_in_ms) {
+    return true;
+  }
+  return false;
+}
+
+std::shared_ptr<RedisConn>& BlockedConnNode::GetConnBlocked() { return conn_blocked_; }
+BlockKeyType BlockedConnNode::GetBlockType() const { return block_type_; }
+
+void DispatchThread::CleanWaitNodeOfUnBlockedBlrConn(std::shared_ptr<net::RedisConn> conn_unblocked) {
+  // removed all the waiting info of this conn/ doing cleaning work
+  auto pair = blocked_conn_to_keys_.find(conn_unblocked->fd());
+  if (pair == blocked_conn_to_keys_.end()) {
+    LOG(WARNING) << "blocking info of blpop/brpop went wrong, blpop/brpop can't working correctly";
+    return;
+  }
+  auto& blpop_keys_list = pair->second;
+  for (auto& blpop_key : *blpop_keys_list) {
+    auto& wait_list_of_this_key = key_to_blocked_conns_.find(blpop_key)->second;
+    for (auto conn = wait_list_of_this_key->begin(); conn != wait_list_of_this_key->end();) {
+      if (conn->GetConnBlocked()->fd() == conn_unblocked->fd()) {
+        conn = wait_list_of_this_key->erase(conn);
+        break;
+      }
+      conn++;
+    }
+  }
+  blocked_conn_to_keys_.erase(conn_unblocked->fd());
+}
+
+void DispatchThread::CleanKeysAfterWaitNodeCleaned() {
+  // after wait info of a conn is cleaned, some wait list of keys might be empty, must erase them from the map
+  std::vector<BlockKey> keys_to_erase;
+  for (auto& pair : key_to_blocked_conns_) {
+    if (pair.second->empty()) {
+      // wait list of this key is empty, just erase this key
+      keys_to_erase.emplace_back(pair.first);
+    }
+  }
+  for (auto& blrpop_key : keys_to_erase) {
+    key_to_blocked_conns_.erase(blrpop_key);
+  }
+}
+
+void DispatchThread::ClosingConnCheckForBlrPop(std::shared_ptr<net::RedisConn> conn_to_close) {
+  if (!conn_to_close) {
+    // dynamic pointer cast failed, it's not an instance of RedisConn, no need of the process below
+    return;
+  }
+  {
+    std::shared_lock l(block_mtx_);
+    if (blocked_conn_to_keys_.find(conn_to_close->fd()) == blocked_conn_to_keys_.end()) {
+      // this conn_to_close is not disconnected from blocking state cause by "blpop/brpop"
+      return;
+    }
+  }
+  std::lock_guard l(block_mtx_);
+  CleanWaitNodeOfUnBlockedBlrConn(conn_to_close);
+  CleanKeysAfterWaitNodeCleaned();
+}
+
+void DispatchThread::ScanExpiredBlockedConnsOfBlrpop() {
+  std::unique_lock latch(block_mtx_);
+  for (auto& pair : key_to_blocked_conns_) {
+    auto& conns_list = pair.second;
+    for (auto conn_node = conns_list->begin(); conn_node != conns_list->end();) {
+      if (conn_node->IsExpired()) {
+        std::shared_ptr conn_ptr = conn_node->GetConnBlocked();
+        conn_ptr->WriteResp("$-1\r\n");
+        conn_ptr->NotifyEpoll(true);
+        conn_node = conns_list->erase(conn_node);
+        CleanWaitNodeOfUnBlockedBlrConn(conn_ptr);
+      } else {
+        conn_node++;
+      }
+    }
+  }
+  CleanKeysAfterWaitNodeCleaned();
+}
+
 void DispatchThread::SetQueueLimit(int queue_limit) { queue_limit_ = queue_limit; }
 
 extern ServerThread* NewDispatchThread(int port, int work_num, ConnFactory* conn_factory, int cron_interval,