fix: client kill preempts in atomic section on shutdown

src/server/server_family.cc

@@ -499,7 +499,7 @@ void ClientId(CmdArgList args, SinkReplyBuilder* builder, ConnectionContext* cnt
 }
 
 void ClientKill(CmdArgList args, absl::Span<facade::Listener*> listeners, SinkReplyBuilder* builder,
-                ConnectionContext* cntx) {
+                ConnectionContext* cntx, util::ProactorPool* pp) {
   std::function<bool(facade::Connection * conn)> evaluator;
 
   if (args.size() == 1) {
@@ -535,14 +535,39 @@ void ClientKill(CmdArgList args, absl::Span<facade::Listener*> listeners, SinkRe
 
   const bool is_admin_request = cntx->conn()->IsPrivileged();
 
+  std::vector<util::fb2::Fiber> fibers(pp->size() * listeners.size());
+  std::vector<std::deque<facade::Connection::WeakRef>> consumers(fibers.size());
+
   atomic<uint32_t> killed_connections = 0;
   atomic<uint32_t> kill_errors = 0;
-  auto cb = [&](unsigned thread_index, util::Connection* conn) {
+
+  size_t round = 0;
+
+  auto client_kill_fb_body = [&consumers, &killed_connections](unsigned idx) {
+    auto& kill_list = consumers[idx];
+    while (!kill_list.empty()) {
+      facade::Connection::WeakRef ref = std::move(kill_list.front());
+      kill_list.pop_front();
+      facade::Connection* conn = ref.Get();
+      // TODO think how to handle migration for eval. See RequestAsyncMigration
+      // DCHECK(socket_->proactor()->GetPoolIndex() == conn.thread())
+      if (conn) {
+        conn->ShutdownSelf();
+        killed_connections.fetch_add(1);
+      }
+    }
+  };
+
+  auto cb = [&](unsigned thread_index, util::Connection* conn) mutable {
     facade::Connection* dconn = static_cast<facade::Connection*>(conn);
+    size_t idx = round + thread_index;
+    auto& fb = fibers[idx];
     if (evaluator(dconn)) {
       if (is_admin_request || !dconn->IsPrivileged()) {
-        dconn->ShutdownSelf();
-        killed_connections.fetch_add(1);
+        if (!fb.IsJoinable()) {
+          fb = [client_kill_fb_body, idx]() mutable { client_kill_fb_body(idx); };
+        }
+        consumers[idx].push_back(dconn->Borrow());
       } else {
         kill_errors.fetch_add(1);
       }
@@ -551,6 +576,11 @@ void ClientKill(CmdArgList args, absl::Span<facade::Listener*> listeners, SinkRe
 
   for (auto* listener : listeners) {
     listener->TraverseConnections(cb);
+    round += pp->size();
+  }
+
+  for (auto& fb : fibers) {
+    fb.JoinIfNeeded();
   }
 
   if (kill_errors.load() == 0) {
@@ -2035,7 +2065,8 @@ void ServerFamily::Client(CmdArgList args, const CommandContext& cmd_cntx) {
   } else if (sub_cmd == "TRACKING") {
     return ClientTracking(sub_args, builder, cntx);
   } else if (sub_cmd == "KILL") {
-    return ClientKill(sub_args, absl::MakeSpan(listeners_), builder, cntx);
+    return ClientKill(sub_args, absl::MakeSpan(listeners_), builder, cntx,
+                      &service_.proactor_pool());
   } else if (sub_cmd == "CACHING") {
     return ClientCaching(sub_args, builder, cmd_cntx.tx, cntx);
   } else if (sub_cmd == "SETINFO") {

tests/dragonfly/connection_test.py

@@ -1331,3 +1331,42 @@ async def test_client_detached_crash(df_factory):
     async_client = server.client()
     await async_client.client_pause(2, all=False)
     server.stop()
+
+
+async def test_tls_client_kill_preemption(
+    with_ca_tls_server_args, with_ca_tls_client_args, df_factory
+):
+    server = df_factory.create(proactor_threads=4, port=BASE_PORT, **with_ca_tls_server_args)
+    server.start()
+
+    client = aioredis.Redis(port=server.port, **with_ca_tls_client_args)
+    assert await client.dbsize() == 0
+
+    # Get the list of clients
+    clients_info = await client.client_list()
+    assert len(clients_info) == 1
+
+    kill_id = clients_info[0]["id"]
+
+    async def seed():
+        with pytest.raises(aioredis.ConnectionError) as roe:
+            while True:
+                p = client.pipeline(transaction=True)
+                expected = []
+                for i in range(100):
+                    p.lpush(str(i), "V")
+                    expected.append(f"LPUSH {i} V")
+
+                await p.execute()
+
+    task = asyncio.create_task(seed())
+
+    await asyncio.sleep(0.1)
+
+    cl = aioredis.Redis(port=server.port, **with_ca_tls_client_args)
+    await cl.execute_command(f"CLIENT KILL ID {kill_id}")
+
+    await task
+    server.stop()
+    lines = server.find_in_logs("Preempting inside of atomic section, fiber")
+    assert len(lines) == 0