Enhance asynchronous performance optimization

Author: wuxianrong

src/pika_kv.cc

@@ -12,8 +12,11 @@
 #include "include/pika_cache.h"
 #include "include/pika_conf.h"
 #include "include/pika_slot_command.h"
+#include "include/pika_server.h"
+#include "praft/praft.h"
 
 extern std::unique_ptr<PikaConf> g_pika_conf;
+extern PikaServer* g_pika_server;
 /* SET key value [NX] [XX] [EX <seconds>] [PX <milliseconds>] */
 void SetCmd::DoInitial() {
   if (!CheckArg(argv_.size())) {
@@ -67,24 +70,74 @@ void SetCmd::DoInitial() {
 void SetCmd::Do() {
   int32_t res = 1;
   STAGE_TIMER_GUARD(storage_duration_ms, true);
+  
+  bool is_raft_leader = false;
+  if (g_pika_server && g_pika_server->GetRaftManager()) {
+    auto node = g_pika_server->GetRaftManager()->GetRaftNode(db_->GetDBName());
+    is_raft_leader = (node && node->IsLeader());
+  }
+  
+  storage::CommitCallback callback = nullptr;
+  
+  if (is_raft_leader && db_->storage() && db_->storage()->IsRaftEnabled()) {
+    auto self = std::static_pointer_cast<SetCmd>(shared_from_this());
+    auto resp_ptr = std::make_shared<std::string>();
+    
+    auto pika_conn = std::dynamic_pointer_cast<PikaClientConn>(GetConn());
+    if (!pika_conn) {
+      res_.SetRes(CmdRes::kErrOther, "Invalid connection");
+      return;
+    }
+    
+    callback = [self, resp_ptr, pika_conn](rocksdb::Status status) {
+      int32_t result = (status.ok() || status.IsNotFound()) ? 1 : 0;
+      
+      if (status.ok() || status.IsNotFound()) {
+        if (self->condition_ == SetCmd::kVX) {
+          self->res_.AppendInteger(self->success_);
+        } else {
+          if (result == 1) {
+            self->res_.SetRes(CmdRes::kOk);
+            AddSlotKey("k", self->key_, self->db_);
+          } else {
+            self->res_.AppendStringLen(-1);
+          }
+        }
+      } else {
+        self->res_.SetRes(CmdRes::kErrOther, status.ToString());
+      }
+      
+      *resp_ptr = std::move(self->res_.message());
+      pika_conn->WriteResp(*resp_ptr);
+      pika_conn->NotifyEpoll(true);
+    };
+  }
+  
+  // Call storage layer with optional callback
   switch (condition_) {
     case SetCmd::kXX:
-      s_ = db_->storage()->Setxx(key_, value_, &res, static_cast<int32_t>(sec_));
+      s_ = db_->storage()->Setxx(key_, value_, &res, static_cast<int32_t>(sec_), callback);
       break;
     case SetCmd::kNX:
-      s_ = db_->storage()->Setnx(key_, value_, &res, static_cast<int32_t>(sec_));
+      s_ = db_->storage()->Setnx(key_, value_, &res, static_cast<int32_t>(sec_), callback);
       break;
     case SetCmd::kVX:
       s_ = db_->storage()->Setvx(key_, target_, value_, &success_, static_cast<int32_t>(sec_));
       break;
     case SetCmd::kEXORPX:
-      s_ = db_->storage()->Setex(key_, value_, static_cast<int32_t>(sec_));
+      s_ = db_->storage()->Setex(key_, value_, static_cast<int32_t>(sec_), callback);
       break;
     default:
-      s_ = db_->storage()->Set(key_, value_);
+      s_ = db_->storage()->Set(key_, value_, callback);
       break;
   }
 
+  // For async mode (Leader), response is handled by callback
+  if (is_raft_leader && db_->storage() && db_->storage()->IsRaftEnabled()) {
+    return;
+  }
+  
+  // For sync mode (non-Leader or non-Raft), set response immediately
   if (s_.ok() || s_.IsNotFound()) {
     if (condition_ == SetCmd::kVX) {
       res_.AppendInteger(success_);

src/pika_server.cc

@@ -117,10 +117,11 @@ PikaServer::PikaServer()
 
     std::lock_guard rwl(storage_options_rw_);
     storage_options_.append_log_function = 
-      [this](const ::pikiwidb::Binlog& binlog, std::promise<rocksdb::Status>&& promise) {
+      [this](const ::pikiwidb::Binlog& binlog, std::promise<rocksdb::Status>&& promise,
+             storage::CommitCallback callback) {
         std::string db_name = "db0";
 
-        raft_manager_->AppendLog(db_name, binlog, std::move(promise));
+        raft_manager_->AppendLog(db_name, binlog, std::move(promise), callback);
       };
     LOG(INFO) << "Raft append_log_function registered in storage_options";
   }

src/praft/include/praft/praft.h

@@ -21,8 +21,10 @@
 #include "pstd/include/pstd_mutex.h"
 #include "pstd/include/pstd_status.h"
 #include "rocksdb/status.h"
+#include "storage/batch.h"
 
 class PikaServer;
+class Cmd;
 
 // Forward declarations
 namespace storage {
@@ -33,6 +35,10 @@ namespace pikiwidb {
 class Binlog;
 }
 
+namespace net {
+class NetConn;
+}
+
 namespace pika_raft {
 
 // Raft log entry data structure
@@ -44,13 +50,22 @@ class WriteDoneClosure : public braft::Closure {
 
   void Run() override;
 
-  // Set promise for synchronous Raft apply
+  // Set promise for synchronous Raft apply (used in Follower on_apply)
   void SetPromise(std::shared_ptr<std::promise<rocksdb::Status>> p) {
     promise_ = p;
   }
+  
+  // Set callback for async response (used in Leader)
+  void SetCallback(storage::CommitCallback callback) {
+    callback_ = callback;
+  }
 
  private:
+  // For synchronous mode (Follower)
   std::shared_ptr<std::promise<rocksdb::Status>> promise_;
+  
+  // For asynchronous mode (Leader)
+  storage::CommitCallback callback_;
 };
 
 // Pika state machine implementation
@@ -162,10 +177,12 @@ class RaftManager {
   // Get cluster information
   pstd::Status GetClusterInfo(const std::string& db_name, std::string* info);
 
-  // Append binlog
+  // Append binlog (supports both sync and async modes)
+  // Sync mode: pass promise, async mode: pass callback
   void AppendLog(const std::string& db_name, 
                  const ::pikiwidb::Binlog& log, 
-                 std::promise<rocksdb::Status>&& promise);
+                 std::promise<rocksdb::Status>&& promise,
+                 storage::CommitCallback callback = nullptr);
 
   // Get Raft node for a specific DB
   std::shared_ptr<PikaRaftNode> GetRaftNode(const std::string& db_name);

src/praft/src/praft.cc

@@ -16,6 +16,8 @@
 #include "brpc/closure_guard.h"
 #include "include/pika_conf.h"
 #include "include/pika_server.h"
+#include "include/pika_command.h"
+#include "include/pika_client_conn.h"
 #include "binlog.pb.h"
 #include "storage/storage.h"
 #include "storage/batch.h"
@@ -568,7 +570,7 @@ braft::PeerId RaftManager::ParsePeerId(const std::string& peer_str) {
 void WriteDoneClosure::Run() {
   std::unique_ptr<WriteDoneClosure> self_guard(this);
 
-  // If promise is set, notify the waiting thread
+  // If promise is set, notify the waiting thread (synchronous mode)
   if (promise_) {
     if (status().ok()) {
       promise_->set_value(rocksdb::Status::OK());
@@ -578,35 +580,60 @@ void WriteDoneClosure::Run() {
     return;
   }
 
-  // Legacy path for non-promise closures
-  if (!status().ok()) {
+  // If callback is set, call it (asynchronous mode for Leader)
+  if (callback_) {
+    rocksdb::Status s;
+    if (status().ok()) {
+      s = rocksdb::Status::OK();
+    } else {
+      s = rocksdb::Status::IOError(status().error_str());
+    }
+    // Call callback with status only (result is captured in lambda)
+    callback_(s);
     return;
   }
+  
+  // Legacy path for closures without promise or callback
+  if (!status().ok()) {
+    LOG(WARNING) << "Raft operation failed: " << status().error_str();
+  }
 }
 
 void RaftManager::AppendLog(const std::string& db_name, 
                             const ::pikiwidb::Binlog& log, 
-                            std::promise<rocksdb::Status>&& promise) {
+                            std::promise<rocksdb::Status>&& promise,
+                            storage::CommitCallback callback) {
   auto node = GetRaftNode(db_name);
   if (!node) {
     LOG(ERROR) << "Raft node not found for DB: " << db_name;
-    promise.set_value(rocksdb::Status::NotFound("Raft node not found"));
+    if (callback) {
+      callback(rocksdb::Status::NotFound("Raft node not found"));
+    } else {
+      promise.set_value(rocksdb::Status::NotFound("Raft node not found"));
+    }
     return;
   }
 
   if (!node->IsLeader()) {
     braft::PeerId leader = node->GetLeaderId();
     LOG(WARNING) << "Current node is not leader for DB: " << db_name 
                  << ", leader: " << leader.to_string();
-    promise.set_value(rocksdb::Status::Incomplete("Not leader"));
+    if (callback) {
+      callback(rocksdb::Status::Incomplete("Not leader, leader is: " + leader.to_string()));
+    } else {
+      promise.set_value(rocksdb::Status::Incomplete("Not leader"));
+    }
     return;
   }
 
-  // 创建 WriteDoneClosure 并传递 promise
   auto* done = new WriteDoneClosure();
-  done->SetPromise(std::make_shared<std::promise<rocksdb::Status>>(std::move(promise)));
 
-  // 序列化 binlog
+  if (callback) {
+    done->SetCallback(callback);
+  } else {
+    done->SetPromise(std::make_shared<std::promise<rocksdb::Status>>(std::move(promise)));
+  }
+  
   butil::IOBuf data;
   butil::IOBufAsZeroCopyOutputStream wrapper(&data);
   if (!log.SerializeToZeroCopyStream(&wrapper)) {
@@ -615,12 +642,10 @@ void RaftManager::AppendLog(const std::string& db_name,
     return;
   }
 
-  // 创建 Raft 任务
   braft::Task task;
   task.data = &data;
   task.done = done;
 
-  // 提交到 Raft
   node->GetRaftNode()->apply(task);
 }
 
@@ -649,7 +674,6 @@ rocksdb::Status RaftManager::ApplyBinlogEntry(const std::string& binlog_data) {
   }
 
   // 调用 Storage::OnBinlogWrite() 应用 binlog
-  // 注意：log_index 暂时传 0，后续可以从外部传入
   auto status = storage->OnBinlogWrite(binlog, 0);
 
   if (!status.ok()) {

src/storage/include/storage/batch.h

@@ -17,12 +17,24 @@ namespace pikiwidb {
 class Binlog;
 }
 
+class Cmd;
+
+namespace net {
+class NetConn;
+}
+
 namespace storage {
 
 class Storage;
 class Redis;
 
-using AppendLogFunction = std::function<void(const ::pikiwidb::Binlog&, std::promise<rocksdb::Status>&&)>;
+// Callback for async commit result
+// Parameters: status (result should be captured in the lambda from storage layer)
+// Note: Connection and result should be captured in the lambda, not passed as parameters
+using CommitCallback = std::function<void(rocksdb::Status)>;
+
+using AppendLogFunction = std::function<void(const ::pikiwidb::Binlog&, std::promise<rocksdb::Status>&&, 
+                                             CommitCallback)>;
 
 using ColumnFamilyIndex = uint32_t;
 
@@ -34,7 +46,7 @@ class Batch {
 
   virtual void Delete(ColumnFamilyIndex cf_idx, const rocksdb::Slice& key) = 0;
 
-  virtual rocksdb::Status Commit() = 0;
+  virtual rocksdb::Status Commit(CommitCallback callback = nullptr) = 0;
 
   int32_t Count() const { return count_; }
 
@@ -52,7 +64,7 @@ class RocksBatch : public Batch {
 
   void Put(ColumnFamilyIndex cf_idx, const rocksdb::Slice& key, const rocksdb::Slice& value) override;
   void Delete(ColumnFamilyIndex cf_idx, const rocksdb::Slice& key) override;
-  rocksdb::Status Commit() override;
+  rocksdb::Status Commit(CommitCallback callback = nullptr) override;
 
 private:
   rocksdb::WriteBatch batch_;
@@ -69,8 +81,8 @@ class BinlogBatch : public Batch {
   void Put(ColumnFamilyIndex cf_idx, const rocksdb::Slice& key, const rocksdb::Slice& value) override;
   void Delete(ColumnFamilyIndex cf_idx, const rocksdb::Slice& key) override;
 
-  // 同步等待 Raft 应用完成（使用 promise/future）
-  rocksdb::Status Commit() override;
+  // Commit to Raft (sync mode if callback is null, async mode otherwise)
+  rocksdb::Status Commit(CommitCallback callback = nullptr) override;
 
 private:
   AppendLogFunction append_log_func_;