Merge pull request #1772 from ClickHouse/small_fixes

[client-v2] Connection reuse strategy parameter
ClickHouse · Aug 16, 2024 · 5609c18 · 5609c18
2 parents 82f8b21 + 00b6d04
commit 5609c18
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 4 deletions.
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java b/client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java
@@ -5,6 +5,7 @@
 import com.clickhouse.client.api.Client;
 import com.clickhouse.client.api.ClientException;
 import com.clickhouse.client.api.ClientMisconfigurationException;
+import com.clickhouse.client.api.ConnectionReuseStrategy;
 import com.clickhouse.client.api.ServerException;
 import com.clickhouse.client.api.enums.ProxyType;
 import com.clickhouse.client.config.ClickHouseClientOption;
@@ -158,8 +159,21 @@ private HttpClientConnectionManager basicConnectionManager(SSLContext sslContext
 
     private HttpClientConnectionManager poolConnectionManager(SSLContext sslContext, SocketConfig socketConfig) {
         PoolingHttpClientConnectionManagerBuilder connMgrBuilder = PoolingHttpClientConnectionManagerBuilder.create()
-                .setConnPoolPolicy(PoolReusePolicy.LIFO)
-                .setPoolConcurrencyPolicy(PoolConcurrencyPolicy.STRICT);
+                .setPoolConcurrencyPolicy(PoolConcurrencyPolicy.LAX);
+
+        ConnectionReuseStrategy connectionReuseStrategy =
+                ConnectionReuseStrategy.valueOf(chConfiguration.get("connection_reuse_strategy"));
+        switch (connectionReuseStrategy) {
+            case LIFO:
+                connMgrBuilder.setConnPoolPolicy(PoolReusePolicy.LIFO);
+                break;
+            case FIFO:
+                connMgrBuilder.setConnPoolPolicy(PoolReusePolicy.FIFO);
+                break;
+            default:
+                throw new ClientMisconfigurationException("Unknown connection reuse strategy: " + connectionReuseStrategy);
+        }
+        LOG.info("Connection reuse strategy: {}", connectionReuseStrategy);
 
         connMgrBuilder.setDefaultConnectionConfig(createConnectionConfig());
         connMgrBuilder.setMaxConnTotal(Integer.MAX_VALUE); // as we do not know how many routes we will have

diff --git a/client-v2/src/test/java/com/clickhouse/client/ConnectionManagementTests.java b/client-v2/src/test/java/com/clickhouse/client/ConnectionManagementTests.java
@@ -2,6 +2,7 @@
 
 import com.clickhouse.client.api.Client;
 import com.clickhouse.client.api.ConnectionInitiationException;
+import com.clickhouse.client.api.ConnectionReuseStrategy;
 import com.clickhouse.client.api.enums.ProxyType;
 import com.clickhouse.client.api.query.GenericRecord;
 import com.clickhouse.client.api.query.QueryResponse;
@@ -150,17 +151,39 @@ public void testConnectionRequestTimeout() {
 
         try (Client client = clientBuilder.build()) {
             CompletableFuture<QueryResponse> f1 = client.query("select 1");
+            Thread.sleep(500L);
             CompletableFuture<QueryResponse> f2 = client.query("select 1");
             f2.get();
         } catch (ExecutionException e) {
             e.printStackTrace();
-            Assert.assertTrue(e.getCause() instanceof ConnectionInitiationException);
-            Assert.assertTrue(e.getCause().getCause() instanceof ConnectionRequestTimeoutException);
+            Assert.assertEquals(e.getCause().getClass(), ConnectionInitiationException.class);
+            Assert.assertEquals(e.getCause().getCause().getClass(), ConnectionRequestTimeoutException.class);
         } catch (Exception e) {
             e.printStackTrace();
             Assert.fail("Unexpected exception", e);
         } finally {
             proxy.stop();
         }
     }
+
+    @Test
+    public void testConnectionReuseStrategy() {
+        ClickHouseNode server = getServer(ClickHouseProtocol.HTTP);
+
+        try (Client client = new Client.Builder()
+                .addEndpoint(server.getBaseUri())
+                .setUsername("default")
+                .setPassword(getPassword())
+                .useNewImplementation(true)
+                .setConnectionReuseStrategy(ConnectionReuseStrategy.LIFO)
+                .build()) {
+
+            List<GenericRecord> records = client.queryAll("SELECT timezone()");
+            Assert.assertTrue(records.size() > 0);
+            Assert.assertEquals(records.get(0).getString(1), "UTC");
+        } catch (Exception e) {
+            e.printStackTrace();
+            Assert.fail(e.getMessage());
+        }
+    }
 }