Better fix + updated tests

ClickHouse · Sep 30, 2023 · ece5e9b · ece5e9b
1 parent b4fe02f
commit ece5e9b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 21 deletions.
diff --git a/clickhouse/client.cpp b/clickhouse/client.cpp
@@ -385,15 +385,17 @@ void Client::Impl::ResetConnectionEndpoint() {
 }
 
 void Client::Impl::CreateConnection() {
-    // i <= 1 to try at least once
-    for (size_t i = 0; i <= 1 || i < options_.send_retries;)
+    // 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;)
     {
         try
         {
+            // Try to connect to each endpoint before throwing exception.
             ResetConnectionEndpoint();
             return;
         } catch (const std::system_error&) {
-            if (++i >= options_.send_retries)
+            if (++i >= max_attempts)
             {
                 throw;
             }

diff --git a/ut/client_ut.cpp b/ut/client_ut.cpp
@@ -1369,58 +1369,79 @@ INSTANTIATE_TEST_SUITE_P(ResetConnectionClientTest, ResetConnectionTestCase,
 ));
 
 struct CountingSocketFactoryAdapter : public SocketFactory {
-    struct Counters
-    {
-        size_t connect_count = 0;
 
-        void Reset() {
-            *this = Counters();
-        }
-    };
+    using ConnectRequests = std::vector<std::pair<ClientOptions, Endpoint>>;
 
     SocketFactory & socket_factory;
-    Counters& counters;
+    ConnectRequests & connect_requests;
 
-    CountingSocketFactoryAdapter(SocketFactory & socket_factory, Counters& counters)
+    CountingSocketFactoryAdapter(SocketFactory & socket_factory, ConnectRequests& connect_requests)
         : socket_factory(socket_factory)
-        , counters(counters)
+        , connect_requests(connect_requests)
     {}
 
     std::unique_ptr<SocketBase> connect(const ClientOptions& opts, const Endpoint& endpoint) {
-        ++counters.connect_count;
+        connect_requests.emplace_back(opts, endpoint);
+
         return socket_factory.connect(opts, endpoint);
     }
 
     void sleepFor(const std::chrono::milliseconds& duration) {
         return socket_factory.sleepFor(duration);
     }
 
+    size_t GetConnectRequestsCount() const {
+        return connect_requests.size();
+    }
+
 };
 
 TEST(SimpleClientTest, issue_335) {
+    // Make sure Client connects to server even with ClientOptions.SetSendRetries(0)
     auto vals = MakeStrings();
     auto col = std::make_shared<ColumnString>(vals);
 
-    CountingSocketFactoryAdapter::Counters counters;
+    CountingSocketFactoryAdapter::ConnectRequests connect_requests;
     std::unique_ptr<SocketFactory> wrapped_socket_factory = std::make_unique<NonSecureSocketFactory>();
-    std::unique_ptr<SocketFactory> socket_factory = std::make_unique<CountingSocketFactoryAdapter>(*wrapped_socket_factory, counters);
+    std::unique_ptr<SocketFactory> socket_factory = std::make_unique<CountingSocketFactoryAdapter>(*wrapped_socket_factory, connect_requests);
 
     Client client(ClientOptions(LocalHostEndpoint)
                       .SetSendRetries(0), // <<=== crucial for reproducing https://github.com/ClickHouse/clickhouse-cpp/issues/335
                   std::move(socket_factory));
 
-    EXPECT_EQ(1u, counters.connect_count);
+    EXPECT_EQ(1u, connect_requests.size());
     EXPECT_TRUE(CompareRecursive(vals, *RoundtripColumnValuesTyped(client, col)));
 
-    counters.Reset();
+    connect_requests.clear();
 
     client.ResetConnection();
-    EXPECT_EQ(1u, counters.connect_count);
+    EXPECT_EQ(1u, connect_requests.size());
     EXPECT_TRUE(CompareRecursive(vals, *RoundtripColumnValuesTyped(client, col)));
 
-    counters.Reset();
+    connect_requests.clear();
 
     client.ResetConnectionEndpoint();
-    EXPECT_EQ(1u, counters.connect_count);
+    EXPECT_EQ(1u, connect_requests.size());
     EXPECT_TRUE(CompareRecursive(vals, *RoundtripColumnValuesTyped(client, col)));
 }
+
+TEST(SimpleClientTest, issue_335_reconnects_count) {
+    // Make sure that client attempts to connect to each endpoint at least once.
+    CountingSocketFactoryAdapter::ConnectRequests connect_requests;
+    std::unique_ptr<SocketFactory> wrapped_socket_factory = std::make_unique<NonSecureSocketFactory>();
+    std::unique_ptr<SocketFactory> socket_factory = std::make_unique<CountingSocketFactoryAdapter>(*wrapped_socket_factory, connect_requests);
+
+    const auto endpoints = {
+        Endpoint{"foo-invalid-hostname", 1234},
+        Endpoint{"bar-invalid-hostname", 4567},
+    };
+
+    EXPECT_ANY_THROW(
+        Client(ClientOptions()
+                          .SetEndpoints(endpoints)
+                          .SetSendRetries(0), // <<=== crucial for reproducing https://github.com/ClickHouse/clickhouse-cpp/issues/335
+                      std::move(socket_factory));
+    );
+
+    EXPECT_EQ(endpoints.size(), connect_requests.size());
+}