Preserve current endpoint across failed reconnects

Keep the current endpoint set to the last successfully connected endpoint. ResetConnection retries only that endpoint, while ResetConnectionEndpoint tries it first before failing over to the remaining endpoints. Restore the previous endpoint if all reconnection attempts fail.

Validate non-empty endpoint lists in RoundRobinEndpointsIterator and add coverage for reconnect and failover ordering.

Author: slabko

clickhouse/base/endpoints_iterator.cpp

@@ -3,8 +3,21 @@
 
 namespace clickhouse {
 
+namespace {
+
+const std::vector<Endpoint> & ValidateEndpoints(const std::vector<Endpoint>& endpoints)
+{
+    if (endpoints.empty()) {
+        throw ValidationError("The list of endpoints is empty");
+    }
+    return endpoints;
+}
+
+} // anonymous namespace
+
 RoundRobinEndpointsIterator::RoundRobinEndpointsIterator(const std::vector<Endpoint>& _endpoints)
-   :  endpoints (_endpoints)
+   :  endpoints (ValidateEndpoints(_endpoints))
+   // set `current_index` to the values such that `Next` returns an element at index 0
    , current_index (endpoints.size() - 1ull)
 {
 }

clickhouse/client.cpp

@@ -179,15 +179,6 @@ std::unique_ptr<SocketFactory> GetSocketFactory(const ClientOptions& opts) {
         return std::make_unique<NonSecureSocketFactory>();
 }
 
-std::unique_ptr<EndpointsIteratorBase> GetEndpointsIterator(const ClientOptions& opts) {
-    if (opts.endpoints.empty())
-    {
-        throw ValidationError("The list of endpoints is empty");
-    }
-
-    return std::make_unique<RoundRobinEndpointsIterator>(opts.endpoints);
-}
-
 } // anonymous namespace
 
 class Client::Impl {
@@ -269,8 +260,6 @@ class Client::Impl {
     /// call fuc several times.
     void RetryGuard(std::function<void()> func);
 
-    void RetryConnectToTheEndpoint(std::function<void()>& func);
-
 private:
     enum class State : uint8_t {
         Idle = 0,
@@ -312,6 +301,8 @@ class Client::Impl {
     std::unique_ptr<SocketBase> socket_;
     std::unique_ptr<EndpointsIteratorBase> endpoints_iterator;
 
+    // current_endpoint_ points to the last successfully connected endpoint, and always
+    // holds a value. The variable remain wrapped as optional for backwards compatibility.
     std::optional<Endpoint> current_endpoint_;
 
     ServerInfo server_info_;
@@ -337,7 +328,8 @@ Client::Impl::Impl(const ClientOptions& opts,
     : options_(modifyClientOptions(opts))
     , events_(nullptr)
     , socket_factory_(std::move(socket_factory))
-    , endpoints_iterator(GetEndpointsIterator(options_))
+    , endpoints_iterator(std::make_unique<RoundRobinEndpointsIterator>(options_.endpoints))
+    , current_endpoint_(endpoints_iterator->Next())
 {
     CreateConnection();
 
@@ -614,36 +606,39 @@ void Client::Impl::ResetConnection() {
 }
 
 void Client::Impl::ResetConnectionEndpoint() {
-    current_endpoint_.reset();
-    for (size_t i = 0; i < options_.endpoints.size();)
+    std::optional<Endpoint> last_endpoint = current_endpoint_;
+    for (size_t i = 1; ; ++i)
     {
         try
         {
-            current_endpoint_ = endpoints_iterator->Next();
             ResetConnection();
             return;
         } catch (const std::system_error&) {
-            if (++i == options_.endpoints.size())
+            current_endpoint_ = endpoints_iterator->Next();
+            if (i >= options_.endpoints.size())
             {
-                current_endpoint_.reset();
+                current_endpoint_ = last_endpoint;
                 throw;
             }
+        } catch (...) {
+            current_endpoint_ = last_endpoint;
+            throw;
         }
     }
 }
 
 void Client::Impl::CreateConnection() {
     // make sure to try to connect to each endpoint at least once even if `options_.send_retries` is 0
     const size_t max_attempts = (options_.send_retries ? options_.send_retries : 1);
-    for (size_t i = 0; i < max_attempts;)
+    for (size_t i = 1; ; ++i)
     {
         try
         {
             // Try to connect to each endpoint before throwing exception.
             ResetConnectionEndpoint();
             return;
         } catch (const std::system_error&) {
-            if (++i >= max_attempts)
+            if (i >= max_attempts)
             {
                 throw;
             }
@@ -1227,33 +1222,36 @@ bool Client::Impl::ReceiveHello() {
 
 void Client::Impl::RetryGuard(std::function<void()> func) {
 
-    if (current_endpoint_)
-    {
-        for (unsigned int i = 0; ; ++i) {
-            try {
-                func();
-                return;
-            } catch (const std::system_error&) {
-                bool ok = true;
+    for (unsigned int i = 1; ; ++i) {
+        try {
+            func();
+            return;
+        } catch (const std::system_error&) {
+            // if send_retries == 0 do not try anymore, throw right away
+            if (options_.send_retries == 0) {
+                throw;
+            }
 
-                try {
-                    socket_factory_->sleepFor(options_.retry_timeout);
-                    ResetConnection();
-                } catch (...) {
-                    ok = false;
-                }
+            // If `send_retries` attempts failed, try other endpoints
+            if (i >= options_.send_retries) {
+                break;
+            }
 
-                if (!ok && i == options_.send_retries) {
-                    break;
-                }
+            // otherwise sleep and try again
+            try {
+                socket_factory_->sleepFor(options_.retry_timeout);
+                ResetConnection();
+            } catch (const std::system_error&) {
             }
+
         }
     }
+
     // Connections with current_endpoint_ are broken.
-    // Trying to establish  with the another one from the list.
+    // Trying to establish with another one from the list.
     size_t connection_attempts_count = options_.endpoints.size() * options_.send_retries;
-
-    for (size_t i = 0; i < connection_attempts_count;)
+    std::optional<Endpoint> last_endpoint = current_endpoint_;
+    for (size_t i = 1; ; ++i)
     {
         try
         {
@@ -1263,11 +1261,14 @@ void Client::Impl::RetryGuard(std::function<void()> func) {
             func();
             return;
         } catch (const std::system_error&) {
-            if (++i == connection_attempts_count)
+            if (i >= connection_attempts_count)
             {
-                current_endpoint_.reset();
+                current_endpoint_ = last_endpoint;
                 throw;
             }
+        } catch (...) {
+            current_endpoint_ = last_endpoint;
+            throw;
         }
     }
 }

clickhouse/client.h

@@ -338,11 +338,12 @@ class Client {
 
     const ServerInfo& GetServerInfo() const;
 
-    /// Get current connected endpoint.
-    /// In case when client is not connected to any endpoint, nullopt will returned.
+    /// Get current endpoint, i.e. the last successfully connected endpoint.
+    /// It remains optional for backward compatibility, but now always contains a value.
     const std::optional<Endpoint>& GetCurrentEndpoint() const;
 
-    // Try to connect to different endpoints one by one only one time. If it doesn't work, throw an exception.
+    /// Try to reconnect to different endpoints one by one only one time. If it doesn't work, throw
+    /// an exception. The function starts with the last successfully connected endpoint.
     void ResetConnectionEndpoint();
 
     struct Version

ut/BUILD.bazel

@@ -110,6 +110,7 @@ cc_test(
         "readonly_client_test.cpp",
         "readonly_client_test.h",
         "roundtrip_tests.cpp",
+        "test_socket_factory_adapters.h",
         # Test entry point and shared support code.
         "main.cpp",
         "roundtrip_column.cpp",

ut/client_ut.cpp

@@ -8,6 +8,7 @@
 
 #include "readonly_client_test.h"
 #include "connection_failed_client_test.h"
+#include "test_socket_factory_adapters.h"
 #include "ut/utils_comparison.h"
 #include "utils.h"
 #include "ut/roundtrip_column.h"
@@ -1849,31 +1850,44 @@ INSTANTIATE_TEST_SUITE_P(MultipleEndpointsFailed, ConnectionFailedClientTest,
 
 class ResetConnectionTestCase : public testing::TestWithParam<ClientOptions> {};
 
-TEST_P(ResetConnectionTestCase, ResetConnectionEndpointTest) {
-    const auto & client_options = GetParam();
-    std::unique_ptr<Client> client;
+TEST(ResetConnectionEndpointTest, ReconnectsCurrentBeforeFailover) {
+    const Endpoint primary{"primary", 9000};
+    const Endpoint secondary{"secondary", 9000};
+    const Endpoint actual_endpoint{LocalHostEndpoint.host, LocalHostEndpoint.port};
 
-    try {
-        client = std::make_unique<Client>(client_options);
-        auto endpoint = client->GetCurrentEndpoint().value();
-        ASSERT_EQ("localhost", endpoint.host);
-        ASSERT_EQ(9000u, endpoint.port);
+    ClientOptions options(LocalHostEndpoint);
+    options.SetHost("");
+    options.SetEndpoints({primary, secondary});
 
-        client->ResetConnectionEndpoint();
-        endpoint = client->GetCurrentEndpoint().value();
-        ASSERT_EQ("127.0.0.1", endpoint.host);
-        ASSERT_EQ(9000u, endpoint.port);
+    // Redirect both logical endpoints to the same reachable test server.
+    auto base_socket_factory = std::make_unique<NonSecureSocketFactory>();
+    auto socket_factory = std::make_unique<FailOnceSocketFactoryAdapter>(*base_socket_factory, actual_endpoint);
+    auto * const adapter = socket_factory.get();
 
-        client->ResetConnectionEndpoint();
+    // The initial connection selects the first endpoint.
+    Client client(options, std::move(socket_factory));
+    ASSERT_EQ(primary, client.GetCurrentEndpoint().value());
 
-        endpoint = client->GetCurrentEndpoint().value();
-        ASSERT_EQ("localhost", endpoint.host);
-        ASSERT_EQ(9000u, endpoint.port);
+    // A healthy current endpoint is retried without advancing.
+    adapter->SetFailEndpoint(std::nullopt);
+    adapter->ClearConnectRequests();
+    client.ResetConnectionEndpoint();
+    EXPECT_EQ(primary, client.GetCurrentEndpoint().value());
+    EXPECT_EQ(std::vector<Endpoint>{primary}, adapter->ConnectRequests());
 
-        SUCCEED();
-    } catch (const std::exception & e) {
-        FAIL() << "Got an unexpected exception : " << e.what();
-    }
+    // Failure of the current endpoint advances to the next endpoint.
+    adapter->SetFailEndpoint(primary);
+    adapter->ClearConnectRequests();
+    client.ResetConnectionEndpoint();
+    EXPECT_EQ(secondary, client.GetCurrentEndpoint().value());
+    EXPECT_EQ((std::vector<Endpoint>{primary, secondary}), adapter->ConnectRequests());
+
+    // Failure of the last endpoint wraps around to the first endpoint.
+    adapter->SetFailEndpoint(secondary);
+    adapter->ClearConnectRequests();
+    client.ResetConnectionEndpoint();
+    EXPECT_EQ(primary, client.GetCurrentEndpoint().value());
+    EXPECT_EQ((std::vector<Endpoint>{secondary, primary}), adapter->ConnectRequests());
 }
 
 TEST_P(ResetConnectionTestCase, ResetConnectionTest) {

ut/test_socket_factory_adapters.h

@@ -0,0 +1,62 @@
+#pragma once
+
+#include "clickhouse/base/socket.h"
+
+#include <chrono>
+#include <memory>
+#include <optional>
+#include <system_error>
+#include <utility>
+#include <vector>
+
+namespace clickhouse {
+
+/** Records requested endpoints and optionally fails one matching connection attempt.
+ *
+ * Successful attempts are redirected to actual_endpoint, allowing tests to exercise
+ * failover between distinct logical endpoints using a single reachable server. Setting
+ * fail_endpoint makes the next matching attempt throw connection_refused and then
+ * clears the value. The wrapped factory must outlive the adapter.
+ */
+struct FailOnceSocketFactoryAdapter : public SocketFactory {
+    SocketFactory & socket_factory;
+    Endpoint actual_endpoint;
+    std::vector<Endpoint> connect_requests{};
+    std::optional<Endpoint> fail_endpoint{};
+
+    FailOnceSocketFactoryAdapter(SocketFactory & socket_factory,
+                                 Endpoint actual_endpoint)
+        : socket_factory(socket_factory)
+        , actual_endpoint(std::move(actual_endpoint))
+    {}
+
+    std::unique_ptr<SocketBase> connect(const ClientOptions& opts,
+                                        const Endpoint& endpoint) override {
+        connect_requests.push_back(endpoint);
+
+        if (fail_endpoint && fail_endpoint.value() == endpoint) {
+            fail_endpoint.reset();
+            throw std::system_error(std::make_error_code(std::errc::connection_refused));
+        }
+
+        return socket_factory.connect(opts, actual_endpoint);
+    }
+
+    void SetFailEndpoint(std::optional<Endpoint> endpoint) {
+        fail_endpoint = std::move(endpoint);
+    }
+
+    const std::vector<Endpoint> & ConnectRequests() const {
+        return connect_requests;
+    }
+
+    void ClearConnectRequests() {
+        connect_requests.clear();
+    }
+
+    void sleepFor(const std::chrono::milliseconds& duration) override {
+        socket_factory.sleepFor(duration);
+    }
+};
+
+} // namespace clickhouse