Preserve current endpoint across failed reconnects

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 value 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 {
@@ -264,18 +255,11 @@ class Client::Impl {
 
     void InitializeStreams(std::unique_ptr<SocketBase>&& socket);
 
-    inline size_t GetConnectionAttempts() const
-    {
-        return options_.endpoints.size() * options_.send_retries;
-    }
-
 private:
     /// In case of network errors tries to reconnect to server and
     /// call fuc several times.
     void RetryGuard(std::function<void()> func);
 
-    void RetryConnectToTheEndpoint(std::function<void()>& func);
-
 private:
     enum class State : uint8_t {
         Idle = 0,
@@ -317,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 remains wrapped as optional for backwards compatibility.
     std::optional<Endpoint> current_endpoint_;
 
     ServerInfo server_info_;
@@ -342,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();
 
@@ -619,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;
             }
@@ -1232,32 +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.
-    size_t connection_attempts_count = GetConnectionAttempts();
-    for (size_t i = 0; i < connection_attempts_count;)
+    // Trying to establish with another one from the list.
+    size_t connection_attempts_count = options_.endpoints.size() * options_.send_retries;
+    std::optional<Endpoint> last_endpoint = current_endpoint_;
+    for (size_t i = 1; ; ++i)
     {
         try
         {
@@ -1267,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