chore: close socket on destructor for outgoing migrations

Signed-off-by: kostas <[email protected]>

Author: kostasrim

src/server/cluster/outgoing_slot_migration.cc

@@ -87,8 +87,8 @@ class OutgoingMigration::SliceSlotMigration : private ProtocolClient {
   }
 
   void Cancel() {
-    // Close socket for clean disconnect.
-    CloseSocket();
+    // Shutdown socket and allow IO loops to return.
+    ShutdownSocket();
     streamer_.Cancel();
   }
 
@@ -100,6 +100,8 @@ class OutgoingMigration::SliceSlotMigration : private ProtocolClient {
     return exec_st_.GetError();
   }
 
+  using ProtocolClient::CloseSocket;
+
  private:
   RestoreStreamer streamer_;
 };
@@ -120,7 +122,10 @@ OutgoingMigration::~OutgoingMigration() {
   exec_st_.JoinErrorHandler();
   // Destroy each flow in its dedicated thread, because we could be the last
   // owner of the db tables
-  OnAllShards([](auto& migration) { migration.reset(); });
+  OnAllShards([](auto& migration) {
+    migration->CloseSocket();
+    migration.reset();
+  });
 }
 
 bool OutgoingMigration::ChangeState(MigrationState new_state) {

src/server/protocol_client.cc

@@ -220,21 +220,30 @@ error_code ProtocolClient::ConnectAndAuth(std::chrono::milliseconds connect_time
   return error_code{};
 }
 
-void ProtocolClient::CloseSocket() {
+void ProtocolClient::ShutdownSocketImpl(bool should_close) {
   unique_lock lk(sock_mu_);
   if (sock_) {
-    sock_->proactor()->Await([this] {
+    sock_->proactor()->Await([this, should_close] {
       if (sock_->IsOpen()) {
         auto ec = sock_->Shutdown(SHUT_RDWR);
         LOG_IF(ERROR, ec) << "Could not shutdown socket " << ec;
       }
-
-      auto ec = sock_->Close();  // Quietly close.
-      LOG_IF(WARNING, ec) << "Error closing socket " << ec << "/" << ec.message();
+      if (should_close) {
+        auto ec = sock_->Close();  // Quietly close.
+        LOG_IF(WARNING, ec) << "Error closing socket " << ec << "/" << ec.message();
+      }
     });
   }
 }
 
+void ProtocolClient::CloseSocket() {
+  return ShutdownSocketImpl(true);
+}
+
+void ProtocolClient::ShutdownSocket() {
+  return ShutdownSocketImpl(false);
+}
+
 void ProtocolClient::DefaultErrorHandler(const GenericError& err) {
   LOG(WARNING) << "Socket error: " << err.Format() << " in " << server_context_.Description()
                << ", socket info: " << GetSocketInfo(sock_ ? sock_->native_handle() : -1);

src/server/protocol_client.h

@@ -40,7 +40,15 @@ class ProtocolClient {
   ProtocolClient(std::string master_host, uint16_t port);
   virtual ~ProtocolClient();
 
-  void CloseSocket();  // Close replica sockets.
+  // First Shutdown() the socket and immediately Close() it.
+  // Any attempt for IO in the socket after Close() will crash with CHECK fail.
+  void CloseSocket();
+
+  // Shutdown the underline socket but do not Close() it. By decoupling this, api
+  // callers can shutdown the socket, wait for the relevant flows to gracefully exit
+  // (by observing during an IO operation that the socket was shut down) and then finally
+  // Close() the socket.
+  void ShutdownSocket();
 
   uint64_t LastIoTime() const;
   void TouchIoTime();
@@ -117,6 +125,8 @@ class ProtocolClient {
  private:
   std::error_code Recv(util::FiberSocketBase* input, base::IoBuf* dest);
 
+  void ShutdownSocketImpl(bool should_close);
+
   ServerContext server_context_;
 
   std::unique_ptr<facade::RedisParser> parser_;