Graceful shutdown the workers when reducing worker threads

src/commands/blocking_commander.h

@@ -44,7 +44,6 @@ class BlockingCommander : public Commander,
   // in other words, returning true indicates ending the blocking
   virtual bool OnBlockingWrite() = 0;
 
-  bool IsBlocking() const override { return true; }
   // to start the blocking process
   // usually put to the end of the Execute method
   Status StartBlocking(int64_t timeout, std::string *output) {

src/commands/cmd_pubsub.cc

@@ -82,7 +82,6 @@ void SubscribeCommandReply(std::string *output, const std::string &name, const s
 
 class CommandSubscribe : public Commander {
  public:
-  bool IsBlocking() const override { return true; }
   Status Execute(Server *srv, Connection *conn, std::string *output) override {
     for (unsigned i = 1; i < args_.size(); i++) {
       conn->SubscribeChannel(args_[i]);
@@ -112,7 +111,6 @@ class CommandUnSubscribe : public Commander {
 
 class CommandPSubscribe : public Commander {
  public:
-  bool IsBlocking() const override { return true; }
   Status Execute(Server *srv, Connection *conn, std::string *output) override {
     for (size_t i = 1; i < args_.size(); i++) {
       conn->PSubscribeChannel(args_[i]);

src/commands/cmd_stream.cc

@@ -664,7 +664,6 @@ class CommandXRead : public Commander,
                      private EvbufCallbackBase<CommandXRead, false>,
                      private EventCallbackBase<CommandXRead> {
  public:
-  bool IsBlocking() const override { return true; }
   Status Parse(const std::vector<std::string> &args) override {
     size_t streams_word_idx = 0;
 

src/commands/commander.h

@@ -70,7 +70,6 @@ class Commander {
   void SetAttributes(const CommandAttributes *attributes) { attributes_ = attributes; }
   const CommandAttributes *GetAttributes() const { return attributes_; }
   void SetArgs(const std::vector<std::string> &args) { args_ = args; }
-  virtual bool IsBlocking() const { return false; }
   virtual Status Parse() { return Parse(args_); }
   virtual Status Parse(const std::vector<std::string> &args) { return Status::OK(); }
   virtual Status Execute(Server *srv, Connection *conn, std::string *output) {

src/server/redis_connection.cc

@@ -438,6 +438,7 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
     if (!reply.empty()) Reply(reply);
     reply.clear();
   }
+  if (current_cmd != nullptr) current_cmd.reset();
 }
 
 void Connection::ResetMultiExec() {

src/server/redis_connection.h

@@ -111,6 +111,7 @@ class Connection : public EvbufCallbackBase<Connection> {
   Worker *Owner() { return owner_; }
   void SetOwner(Worker *new_owner) { owner_ = new_owner; };
   int GetFD() { return bufferevent_getfd(bev_); }
+  bool HasRunningCommand() const { return current_cmd != nullptr; }
   evbuffer *Input() { return bufferevent_get_input(bev_); }
   evbuffer *Output() { return bufferevent_get_output(bev_); }
   bufferevent *GetBufferEvent() { return bev_; }
@@ -127,8 +128,8 @@ class Connection : public EvbufCallbackBase<Connection> {
   void ResetMultiExec();
   std::deque<redis::CommandTokens> *GetMultiExecCommands() { return &multi_cmds_; }
 
-  std::unique_ptr<Commander> current_cmd;
   std::function<void(int)> close_cb = nullptr;
+  std::unique_ptr<Commander> current_cmd;
 
   std::set<std::string> watched_keys;
   std::atomic<bool> watched_keys_modified = false;

src/server/server.cc

@@ -115,7 +115,7 @@ Server::~Server() {
   for (auto &worker_thread : worker_threads_) {
     worker_thread.reset();
   }
-  cleanupExitedWorkerThreads();
+  cleanupExitedWorkerThreads(true /* force */);
   CleanupExitedSlaves();
 
   lua::DestroyState(lua_);
@@ -226,7 +226,7 @@ void Server::Stop() {
   slaveof_mu_.unlock();
 
   for (const auto &worker : worker_threads_) {
-    worker->Stop();
+    worker->Stop(0 /* immediately terminate  */);
   }
 
   rocksdb::CancelAllBackgroundWork(storage->GetDB(), true);
@@ -739,8 +739,9 @@ void Server::cron() {
       storage->SetDBInRetryableIOError(false);
     }
 
-    if (counter != 0 && counter % 10 == 0) {
-      cleanupExitedWorkerThreads();
+    // check if we need to clean up exited worker threads every 5s
+    if (counter != 0 && counter % 50 == 0) {
+      cleanupExitedWorkerThreads(false);
     }
 
     CleanupExitedSlaves();
@@ -1685,12 +1686,12 @@ void Server::AdjustWorkerThreads() {
   if (new_worker_threads > worker_threads_.size()) {
     delta = new_worker_threads - worker_threads_.size();
     increaseWorkerThreads(delta);
-    LOG(INFO) << "[server] Increase worker threads to " << new_worker_threads;
+    LOG(INFO) << "[server] Increase worker threads from " << worker_threads_.size() << " to " << new_worker_threads;
     return;
   }
 
   delta = worker_threads_.size() - new_worker_threads;
-  LOG(INFO) << "[server] Decrease worker threads to " << new_worker_threads;
+  LOG(INFO) << "[server] Decrease worker threads from " << worker_threads_.size() << " to " << new_worker_threads;
   decreaseWorkerThreads(delta);
 }
 
@@ -1721,17 +1722,26 @@ void Server::decreaseWorkerThreads(size_t delta) {
       auto target_worker = worker_threads_[iter.first % remain_worker_threads]->GetWorker();
       worker_thread->GetWorker()->MigrateConnection(target_worker, iter.second);
     }
-    worker_thread->Stop();
+    worker_thread->Stop(10 /* graceful timeout */);
     // Don't join the worker thread here, because it may join itself.
     recycle_worker_threads_.push(std::move(worker_thread));
   }
 }
 
-void Server::cleanupExitedWorkerThreads() {
+void Server::cleanupExitedWorkerThreads(bool force) {
   std::unique_ptr<WorkerThread> worker_thread = nullptr;
-  while (recycle_worker_threads_.try_pop(worker_thread)) {
-    worker_thread->Join();
-    worker_thread.reset();
+  auto total = recycle_worker_threads_.unsafe_size();
+  for (size_t i = 0; i < total; i++) {
+    if (!recycle_worker_threads_.try_pop(worker_thread)) {
+      break;
+    }
+    if (worker_thread->IsTerminated() || force) {
+      worker_thread->Join();
+      worker_thread.reset();
+    } else {
+      // Push the worker thread back to the queue if it's still running.
+      recycle_worker_threads_.push(std::move(worker_thread));
+    }
   }
 }
 

src/server/server.h

@@ -305,7 +305,7 @@ class Server {
   void updateAllWatchedKeys();
   void increaseWorkerThreads(size_t delta);
   void decreaseWorkerThreads(size_t delta);
-  void cleanupExitedWorkerThreads();
+  void cleanupExitedWorkerThreads(bool force);
 
   std::atomic<bool> stop_ = false;
   std::atomic<bool> is_loading_ = false;

src/server/worker.cc

@@ -252,7 +252,8 @@ Status Worker::listenTCP(const std::string &host, uint32_t port, int backlog) {
     }
 
     evutil_make_socket_nonblocking(fd);
-    auto lev = NewEvconnlistener<&Worker::newTCPConnection>(base_, LEV_OPT_CLOSE_ON_FREE, backlog, fd);
+    auto lev =
+        NewEvconnlistener<&Worker::newTCPConnection>(base_, LEV_OPT_THREADSAFE | LEV_OPT_CLOSE_ON_FREE, backlog, fd);
     listen_events_.emplace_back(lev);
   }
 
@@ -292,14 +293,22 @@ void Worker::Run(std::thread::id tid) {
   if (event_base_dispatch(base_) != 0) {
     LOG(ERROR) << "[worker] Failed to run server, err: " << strerror(errno);
   }
+  is_terminated_ = true;
 }
 
-void Worker::Stop() {
-  event_base_loopbreak(base_);
+void Worker::Stop(uint32_t wait_seconds) {
   for (const auto &lev : listen_events_) {
     // It's unnecessary to close the listener fd since we have set the LEV_OPT_CLOSE_ON_FREE flag
     evconnlistener_free(lev);
   }
+  // wait_seconds == 0 means stop immediately, or it will wait N seconds
+  // for the worker to process the remaining requests before stopping.
+  if (wait_seconds > 0) {
+    timeval tv = {wait_seconds, 0};
+    event_base_loopexit(base_, &tv);
+  } else {
+    event_base_loopbreak(base_);
+  }
 }
 
 Status Worker::AddConnection(redis::Connection *c) {
@@ -351,17 +360,19 @@ redis::Connection *Worker::removeConnection(int fd) {
 // blocked on a key or stream.
 void Worker::MigrateConnection(Worker *target, redis::Connection *conn) {
   if (!target || !conn) return;
-  if (conn->current_cmd != nullptr && conn->current_cmd->IsBlocking()) {
+  // We cannot migrate the connection if it has a running command
+  // since it will cause data race since the old worker may still process the command.
+  if (conn->HasRunningCommand()) {
     // don't need to close the connection since destroy worker thread will close it
     return;
   }
 
+  // remove the connection from current worker
+  DetachConnection(conn);
   if (!target->AddConnection(conn).IsOK()) {
-    // destroy worker thread will close the connection
+    conn->Close();
     return;
   }
-  // remove the connection from current worker
-  DetachConnection(conn);
   auto bev = conn->GetBufferEvent();
   bufferevent_base_set(target->base_, bev);
   conn->SetCB(bev);
@@ -540,7 +551,7 @@ void WorkerThread::Start() {
   LOG(INFO) << "[worker] Thread #" << t_.get_id() << " started";
 }
 
-void WorkerThread::Stop() { worker_->Stop(); }
+void WorkerThread::Stop(uint32_t wait_seconds) { worker_->Stop(wait_seconds); }
 
 void WorkerThread::Join() {
   if (auto s = util::ThreadJoin(t_); !s) {

src/server/worker.h

@@ -50,8 +50,9 @@ class Worker : EventCallbackBase<Worker>, EvconnlistenerBase<Worker> {
   Worker(Worker &&) = delete;
   Worker &operator=(const Worker &) = delete;
 
-  void Stop();
+  void Stop(uint32_t wait_seconds);
   void Run(std::thread::id tid);
+  bool IsTerminated() const { return is_terminated_; }
 
   void MigrateConnection(Worker *target, redis::Connection *conn);
   void DetachConnection(redis::Connection *conn);
@@ -94,6 +95,7 @@ class Worker : EventCallbackBase<Worker>, EvconnlistenerBase<Worker> {
   struct bufferevent_rate_limit_group *rate_limit_group_ = nullptr;
   struct ev_token_bucket_cfg *rate_limit_group_cfg_ = nullptr;
   lua_State *lua_;
+  std::atomic<bool> is_terminated_ = false;
 };
 
 class WorkerThread {
@@ -106,8 +108,9 @@ class WorkerThread {
 
   Worker *GetWorker() { return worker_.get(); }
   void Start();
-  void Stop();
+  void Stop(uint32_t wait_seconds);
   void Join();
+  bool IsTerminated() const { return worker_->IsTerminated(); }
 
  private:
   std::thread t_;

tests/gocase/unit/config/config_test.go

@@ -151,10 +151,7 @@ func TestDynamicChangeWorkerThread(t *testing.T) {
 	defer srv.Close()
 
 	ctx := context.Background()
-	rdb := srv.NewClientWithOption(&redis.Options{
-		MaxIdleConns: 20,
-		MaxRetries:   -1, // Disable retry to check connections are alive after config change
-	})
+	rdb := srv.NewClient()
 	defer func() { require.NoError(t, rdb.Close()) }()
 
 	t.Run("Test dynamic change worker thread", func(t *testing.T) {