feat(replication): implement connection health checks and task retry logic

src/ml_flashpoint/replication/transfer_service/connection_pool.cpp

@@ -51,7 +51,7 @@ ScopedConnection::~ScopedConnection() { Release(); }
 
 void ScopedConnection::Release() {
   if (pool_ != nullptr && sockfd_ >= 0) {
-    pool_->ReleaseConnection(sockfd_, true);
+    pool_->ReleaseConnection(sockfd_, reuse_);
   } else if (sockfd_ >= 0) {
     close(sockfd_);
   }
@@ -154,27 +154,93 @@ int ConnectionPool::CreateConnection() {
   return -1;
 }
 
+bool ConnectionPool::IsConnectionAlive(int sockfd) {
+  if (sockfd < 0) return false;
+
+  char buf;
+  // We use MSG_PEEK | MSG_DONTWAIT to check the status of the TCP connection
+  // without actually consuming any data from the socket's receive buffer.
+  // This is a fast, zero-copy way to ask the kernel if the connection has
+  // been closed or has encountered an error while it was idle in the pool.
+  ssize_t r = recv(sockfd, &buf, 1, MSG_PEEK | MSG_DONTWAIT);
+
+  if (r == 0) {
+    // If recv returns 0, it means the remote peer has performed an orderly
+    // shutdown (sent a FIN packet). The connection is no longer usable for
+    // sending new requests.
+    return false;
+  } else if (r < 0) {
+    // If recv returns -1, we check errno to distinguish between "no data"
+    // and a real network error.
+    if (errno == EAGAIN || errno == EWOULDBLOCK) {
+      // EAGAIN/EWOULDBLOCK means the connection is still alive and healthy,
+      // but there is currently no data waiting to be read. This is the
+      // expected state for an idle connection in the pool.
+      return true;
+    }
+    // Any other error (like ECONNRESET or EPIPE) indicates that the
+    // connection has been broken or timed out.
+    return false;
+  }
+  // If r > 0, there is actual data waiting in the buffer. While unusual for
+  // an idle pool connection, it indicates the connection is definitely alive.
+  return true;
+}
+
 std::optional<ScopedConnection> ConnectionPool::GetConnection(int timeout_ms) {
   CHECK_GT(timeout_ms, 0) << "timeout_ms must be positive";
-  std::unique_lock<std::mutex> lock(mtx_);
-  if (!cv_.wait_for(lock, std::chrono::milliseconds(timeout_ms), [this] {
-        return !available_connections_.empty() || stopping_;
-      })) {
-    LOG(WARNING) << "ConnectionPool::GetConnection: timeout";
-    return std::nullopt;
-  }
-  if (stopping_) {
-    LOG(WARNING) << "ConnectionPool::GetConnection: stopping";
-    return std::nullopt;
-  }
-  if (available_connections_.empty()) {
-    // TODO: Handle the case when we run out of connections
-    LOG(WARNING) << "ConnectionPool::GetConnection: no available connections";
-    return std::nullopt;
+
+  // Calculate the absolute deadline to ensure we respect the user-provided
+  // timeout even if we have to loop through several dead connections.
+  auto start_time = std::chrono::steady_clock::now();
+  auto end_time = start_time + std::chrono::milliseconds(timeout_ms);
+
+  while (true) {
+    std::unique_lock<std::mutex> lock(mtx_);
+
+    // Re-calculate the remaining wait time for each iteration of the loop.
+    auto remaining = std::chrono::duration_cast<std::chrono::milliseconds>(
+        end_time - std::chrono::steady_clock::now());
+
+    if (remaining.count() <= 0 || !cv_.wait_for(lock, remaining, [this] {
+          return !available_connections_.empty() || stopping_;
+        })) {
+      LOG(WARNING) << "ConnectionPool::GetConnection: timeout reached while "
+                      "searching for a healthy connection";
+      return std::nullopt;
+    }
+
+    if (stopping_) {
+      LOG(WARNING) << "ConnectionPool::GetConnection: pool is shutting down";
+      return std::nullopt;
+    }
+
+    // Pop the oldest connection from the FIFO queue.
+    int fd = available_connections_.front();
+    available_connections_.pop();
+
+    // Verify the connection's health before handing it to the caller.
+    // This protects against "stale" connections that were closed by the
+    // peer or a firewall while sitting idle in the pool.
+    if (IsConnectionAlive(fd)) {
+      return ScopedConnection(fd, this);
+    }
+
+    // The connection is dead. We close it and attempt to retrieve another
+    // one from the queue.
+    LOG(INFO) << "ConnectionPool::GetConnection: discarded dead connection; "
+                 "retrying with next available connection";
+    close(fd);
+
+    // To maintain the desired pool size, we immediately attempt to open a
+    // replacement connection. This ensures the pool doesn't slowly drain
+    // if many connections go stale at once.
+    int new_fd = CreateConnection();
+    if (new_fd >= 0) {
+      available_connections_.push(new_fd);
+      // The loop will continue and pick up this or another connection.
+    }
   }
-  int fd = available_connections_.front();
-  available_connections_.pop();
-  return ScopedConnection(fd, this);
 }
 
 // Returns a connection to the pool, allowing it to be reused.
@@ -186,18 +252,18 @@ void ConnectionPool::ReleaseConnection(int sockfd, bool reuse) {
     LOG(WARNING) << "ConnectionPool::ReleaseConnection: invalid sockfd";
     return;
   }
+
   std::unique_lock<std::mutex> lock(mtx_);
   if (stopping_) {
     LOG(WARNING)
         << "ConnectionPool::ReleaseConnection: stopping, close connection";
     close(sockfd);
     return;
   }
+
   if (reuse) {
     if (available_connections_.size() < max_size_) {
       LOG(INFO) << "ConnectionPool::ReleaseConnection: reuse connection";
-      // TODO: Check if we need cleanup for the connection before return it to
-      // the pool
       available_connections_.push(sockfd);
       cv_.notify_one();
     } else {
@@ -207,8 +273,22 @@ void ConnectionPool::ReleaseConnection(int sockfd, bool reuse) {
     }
   } else {
     LOG(INFO)
-        << "ConnectionPool::ReleaseConnection: do not reuse, close connection";
+        << "ConnectionPool::ReleaseConnection: connection marked as unusable, "
+           "closing and replenishing pool";
     close(sockfd);
+
+    // Since we are discarding a connection that was previously part of the
+    // pool's "active" set, we create a new one to maintain the fixed pool size.
+    // This prevents the pool from permanently shrinking when network errors
+    // occur.
+    int new_fd = CreateConnection();
+    if (new_fd >= 0) {
+      available_connections_.push(new_fd);
+      cv_.notify_one();
+    } else {
+      LOG(ERROR) << "ConnectionPool::ReleaseConnection: failed to replenish "
+                    "pool after discarding unusable connection";
+    }
   }
 }
 }  // namespace ml_flashpoint::replication::transfer_service

src/ml_flashpoint/replication/transfer_service/connection_pool.h

@@ -57,11 +57,17 @@ class ScopedConnection {
 
   int fd() const { return sockfd_; }
   bool IsValid() const { return sockfd_ >= 0; }
+
+  // Marks the connection as unusable (e.g., after a socket error).
+  // This prevents it from being returned to the pool for reuse.
+  void SetUnusable() { reuse_ = false; }
+
   void Release();
 
  private:
   int sockfd_;
   ConnectionPool* pool_;
+  bool reuse_ = true;
 };
 
 // Manages a thread-safe pool of TCP connections to a single peer.
@@ -108,6 +114,9 @@ class ConnectionPool {
   // Releases a connection back to the pool or closes it.
   void ReleaseConnection(int sockfd, bool reuse = true);
 
+  // Checks if a connection is still alive by performing a non-blocking peek.
+  bool IsConnectionAlive(int sockfd);
+
   // Creates a new connection to the peer.
   // Returns the socket file descriptor on a successful connection, or -1 on
   // failure.