From 838408ff0a5f16deeeb8ca1f856ea6de30ccb004 Mon Sep 17 00:00:00 2001
From: lzydmxy <13126752315@163.com>
Date: Thu, 31 Oct 2024 16:09:25 +0800
Subject: [PATCH 1/2] Force_connected for retry resume connection establisher

---
 src/Client/ConnectionEstablisher.cpp    | 3 ++-
 src/Client/ConnectionEstablisher.h      | 3 +++
 src/Client/HedgedConnectionsFactory.cpp | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/Client/ConnectionEstablisher.cpp b/src/Client/ConnectionEstablisher.cpp
index 178e26b124c..d1960e51e03 100644
--- a/src/Client/ConnectionEstablisher.cpp
+++ b/src/Client/ConnectionEstablisher.cpp
@@ -151,7 +151,8 @@ void ConnectionEstablisherAsync::Task::run(AsyncCallback async_callback, Suspend
 {
     connection_establisher_async.reset();
     connection_establisher_async.connection_establisher.setAsyncCallback(async_callback);
-    connection_establisher_async.connection_establisher.run(connection_establisher_async.result, connection_establisher_async.fail_message);
+    connection_establisher_async.connection_establisher.run(connection_establisher_async.result,
+        connection_establisher_async.fail_message, connection_establisher_async.force_connected);
     connection_establisher_async.is_finished = true;
 }
 
diff --git a/src/Client/ConnectionEstablisher.h b/src/Client/ConnectionEstablisher.h
index ff071e59aea..e7833237600 100644
--- a/src/Client/ConnectionEstablisher.h
+++ b/src/Client/ConnectionEstablisher.h
@@ -76,6 +76,8 @@ public:
 
     const std::string & getFailMessage() const { return fail_message; }
 
+    void resumeConnectionWithForceOption(bool force_connected_) {force_connected = force_connected_; resume();}
+
 private:
     bool checkBeforeTaskResume() override;
 
@@ -125,6 +127,7 @@ private:
 
     bool is_finished = false;
     bool restarted = false;
+    bool force_connected = false;
 };
 
 #endif
diff --git a/src/Client/HedgedConnectionsFactory.cpp b/src/Client/HedgedConnectionsFactory.cpp
index df63a124539..0900f1c0293 100644
--- a/src/Client/HedgedConnectionsFactory.cpp
+++ b/src/Client/HedgedConnectionsFactory.cpp
@@ -281,7 +281,7 @@ int HedgedConnectionsFactory::getReadyFileDescriptor(bool blocking, AsyncCallbac
 
 HedgedConnectionsFactory::State HedgedConnectionsFactory::resumeConnectionEstablisher(int index, Connection *& connection_out)
 {
-    replicas[index].connection_establisher->resume();
+    replicas[index].connection_establisher->resumeConnectionWithForceOption(/*force_connected_*/ shuffled_pools[index].error_count != 0);
 
     if (replicas[index].connection_establisher->isCancelled())
         return State::CANNOT_CHOOSE;

From 30c33aa1a7c5c05097f7f8b32c66950c84a68669 Mon Sep 17 00:00:00 2001
From: lzydmxy <13126752315@163.com>
Date: Tue, 12 Nov 2024 10:18:37 +0800
Subject: [PATCH 2/2] Fix test_distributed_respect_user_timeouts expected

---
 tests/integration/test_distributed_respect_user_timeouts/test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/test_distributed_respect_user_timeouts/test.py b/tests/integration/test_distributed_respect_user_timeouts/test.py
index dd0953c47dd..8ea4efde701 100644
--- a/tests/integration/test_distributed_respect_user_timeouts/test.py
+++ b/tests/integration/test_distributed_respect_user_timeouts/test.py
@@ -80,6 +80,7 @@ def _check_exception(exception, expected_tries=3):
     for i, line in enumerate(lines[3 : 3 + expected_tries]):
         expected_lines = (
             "Code: 209. " + EXCEPTION_NETWORK + EXCEPTION_TIMEOUT,
+            "Code: 209. " + EXCEPTION_NETWORK + "Timeout: connect timed out",
             EXCEPTION_CONNECT_TIMEOUT,
             EXCEPTION_TIMEOUT,
         )