feat(replication): implement connection health checks and task retry logic#65
feat(replication): implement connection health checks and task retry logic#65
Conversation
…logic
This PR enhances the ConnectionPool and TransferService to robustly handle stale TCP connections and transient network failures, ensuring
high reliability in distributed environments where connections may be closed by peers or intermediate firewalls due to idleness.
1. Proactive Health Checks in ConnectionPool
* Zero-Cost Detection: Added IsConnectionAlive() using MSG_PEEK | MSG_DONTWAIT to verify socket health before handing connections to
worker threads. This detects remote closures (FIN) or resets (RST) without consuming protocol data.
* Self-Healing Logic: Updated GetConnection() to automatically discard dead sockets and replenish the pool with fresh connections. The
method respects the original user timeout even if multiple stale connections are encountered.
* Automatic Replenishment: Improved ReleaseConnection() to ensure the pool maintains its configured size by opening new connections when a
socket is discarded as unusable.
2. Resilient Task Execution in TransferService
* Transparent Retries: Implemented a dual-attempt retry mechanism for PutTask, GetTask, and RespondToGetTask. Tasks now automatically
recover from "race condition" connection failures that occur between the health check and the first protocol write.
* Prevention of Pool Poisoning: Added ScopedConnection::SetUnusable(). The TransferService now explicitly flags connections as bad upon
encountering any socket-level error (SendAll/RecvHeader failures). This ensures broken sockets are closed permanently and not returned
to the pool for reuse by other tasks.
* Exhaustive Protocol Handling: Replaced default cases with exhaustive switch statements for all MessageType handling, ensuring strict
protocol compliance and easier debugging of unexpected message types.
3. Verification & Testing
* Simulated Stale Connections: Added DiscardDeadConnection unit test to verify the pool's ability to recover when the remote side closes a
socket while it is idle in the queue.
* Invalidation Logic: Added SetUnusablePreventsReuse to confirm that connections flagged as unusable are correctly disposed of.
* Regression Testing: Verified that all 21 peer-to-peer integration tests in TransferServiceP2PTest pass with the new recovery logic.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly improves the robustness and reliability of the replication transfer service by addressing issues related to stale TCP connections and transient network failures. It introduces proactive health checks within the connection pool to identify and replace dead connections, implements a retry mechanism for transfer tasks to handle intermittent communication issues, and prevents the reuse of broken sockets, ensuring a more stable and self-healing system. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request significantly enhances the robustness of the connection pool and transfer service by adding proactive connection health checks and task retry logic. The changes are well-structured and the added features are crucial for handling transient network failures and stale connections. My review identifies a couple of high-severity issues related to concurrency and data correctness, and a medium-severity style guide violation in the new tests. Addressing these points will further improve the quality and reliability of the implementation.
| if (new_fd >= 0) { | ||
| available_connections_.push(new_fd); | ||
| // The loop will continue and pick up this or another connection. | ||
| } | ||
| } |
There was a problem hiding this comment.
Calling CreateConnection() while holding the pool's mutex (mtx_) can cause significant contention, as it's a blocking network operation. This will block any other thread trying to get or release a connection. Consider releasing the lock before this call and re-acquiring it after to push the new connection to the pool. Remember to re-check stopping_ after re-acquiring the lock.
| 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"; | ||
| } |
There was a problem hiding this comment.
As in GetConnection, CreateConnection() is called here while holding the pool's mutex. This blocking call can harm concurrency. The lock should be released before creating the connection and re-acquired after. Don't forget to handle the stopping_ case after re-acquiring the lock.
| } | ||
| ScopedConnection conn = std::move(conn_opt.value()); | ||
| metric_container.connection_acquired_time = absl::Now(); | ||
|
|
There was a problem hiding this comment.
The header object is not initialized here. This means that any fields not explicitly set will contain garbage values, which are then sent over the network. This is inconsistent with ExecutePutTask and ExecuteRespondToGetTask where the header is zero-initialized. Please initialize the header to prevent sending indeterminate data.
| ObjInfoHeader header{}; |
| TEST_F(ConnectionPoolTest, DiscardDeadConnection) { | ||
| ConnectionPool pool("127.0.0.1", port_, 1); | ||
| EXPECT_TRUE(pool.Initialize()); | ||
|
|
||
| // Wait for the connection to be accepted by the server. | ||
| std::this_thread::sleep_for(std::chrono::milliseconds(100)); | ||
|
|
||
| // Close the connection on the server side to make it stale in the pool. | ||
| { | ||
| std::lock_guard<std::mutex> lock(accepted_fds_mutex_); | ||
| ASSERT_EQ(accepted_fds_.size(), 1); | ||
| close(accepted_fds_[0]); | ||
| accepted_fds_.clear(); | ||
| } | ||
|
|
||
| // GetConnection should detect it's dead, discard it, and create a new one. | ||
| auto conn = pool.GetConnection(); | ||
| EXPECT_TRUE(conn.has_value()); | ||
| EXPECT_TRUE(conn->IsValid()); | ||
|
|
||
| // Verify that a new connection was indeed established. | ||
| std::this_thread::sleep_for(std::chrono::milliseconds(100)); | ||
| std::lock_guard<std::mutex> lock(accepted_fds_mutex_); | ||
| EXPECT_EQ(accepted_fds_.size(), 1); | ||
| } |
There was a problem hiding this comment.
According to the repository's testing style guide (rule #1 under 'Testing Rules'), tests should be structured with # Given, # When, and # Then comments to separate setup, action, and assertion phases. This improves readability. This applies to both DiscardDeadConnection and SetUnusablePreventsReuse tests.
TEST_F(ConnectionPoolTest, DiscardDeadConnection) {
// # Given
ConnectionPool pool("127.0.0.1", port_, 1);
EXPECT_TRUE(pool.Initialize());
// Wait for the connection to be accepted by the server.
std::this_thread::sleep_for(std::chrono::milliseconds(100));
// Close the connection on the server side to make it stale in the pool.
{
std::lock_guard<std::mutex> lock(accepted_fds_mutex_);
ASSERT_EQ(accepted_fds_.size(), 1);
close(accepted_fds_[0]);
accepted_fds_.clear();
}
// # When
// GetConnection should detect it's dead, discard it, and create a new one.
auto conn = pool.GetConnection();
// # Then
EXPECT_TRUE(conn.has_value());
EXPECT_TRUE(conn->IsValid());
// Verify that a new connection was indeed established.
std::this_thread::sleep_for(std::chrono::milliseconds(100));
std::lock_guard<std::mutex> lock(accepted_fds_mutex_);
EXPECT_EQ(accepted_fds_.size(), 1);
}References
- Format tests such that the setup, action and assertion parts are separated by an empty line when possible. Add the comments "# Given", "# When" and "# Then" before each section respectively. (link)
Python Code Coverage Summary
Minimum allowed line rate is |
C++ Code Coverage Summary
Minimum allowed line rate is |
1 participant
This PR enhances the ConnectionPool and TransferService to robustly handle stale TCP connections and transient network failures, ensuring high reliability in scenarios where connections may be closed by peers or intermediate firewalls due to idleness.