From 5a77a64031f10f1aff599f539747c565eef28699 Mon Sep 17 00:00:00 2001
From: Alexandr Gorshenin <alexandr268@ydb.tech>
Date: Mon, 11 Nov 2024 17:56:32 +0000
Subject: [PATCH] Added option cachedTransportsCount to shard JDBC connections

---
 .../tech/ydb/jdbc/settings/YdbConfig.java     | 18 +++++++++++++--
 .../ydb/jdbc/YdbDriverIntegrationTest.java    | 23 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java b/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java
index 4bdea69..5643233 100644
--- a/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java
+++ b/jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java
@@ -10,6 +10,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Properties;
+import java.util.concurrent.ThreadLocalRandom;
 
 import tech.ydb.core.utils.URITools;
 import tech.ydb.jdbc.YdbConst;
@@ -50,6 +51,9 @@ public class YdbConfig {
     static final YdbProperty<Boolean> TRANSACTION_TRACER = YdbProperty.bool(
             "enableTxTracer", "Enable collecting of transaction execution traces", false
     );
+    static final YdbProperty<Integer> CACHED_TRANSPORT_COUNT = YdbProperty.integer(
+            "cachedTransportsCount", "Use specified count of YDB transports in context cache", 1
+    );
 
     private final String url;
     private final String username;
@@ -67,6 +71,7 @@ public class YdbConfig {
 
     private final boolean fullScanDetectorEnabled;
     private final boolean txTracerEnabled;
+    private final int transportIndex;
 
     private YdbConfig(
             String url, String safeUrl, String connectionString, String username, String password, Properties props
@@ -85,6 +90,13 @@ private YdbConfig(
 
         this.fullScanDetectorEnabled = FULLSCAN_DETECTOR_ENABLED.readValue(props).getValue();
         this.txTracerEnabled = TRANSACTION_TRACER.readValue(props).getValue();
+
+        int transportsCount = CACHED_TRANSPORT_COUNT.readValue(props).getValue();
+        if (transportsCount > 1) {
+            this.transportIndex = ThreadLocalRandom.current().nextInt(transportsCount);
+        } else {
+            this.transportIndex = 0;
+        }
     }
 
     public Properties getSafeProps() {
@@ -144,12 +156,14 @@ public boolean equals(Object o) {
             return false;
         }
         YdbConfig that = (YdbConfig) o;
-        return Objects.equals(url, that.url) && Objects.equals(properties, that.properties);
+        return Objects.equals(url, that.url)
+                && Objects.equals(properties, that.properties)
+                && transportIndex == that.transportIndex;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(url, properties);
+        return Objects.hash(url, properties, transportIndex);
     }
 
     public String getUrl() {
diff --git a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java b/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java
index 58aebeb..dc54116 100644
--- a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java
+++ b/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java
@@ -7,8 +7,10 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Deque;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Properties;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
 import org.junit.jupiter.api.Assertions;
@@ -195,6 +197,27 @@ public void testContextCacheConncurrent() throws SQLException {
         }
     }
 
+    @Test
+    public void testCachedTransportCount() throws SQLException {
+        Properties props = new Properties();
+        props.put("cachedTransportsCount", "3");
+
+        List<Connection> list = new ArrayList<>();
+
+        Set<YdbContext> contexts = new HashSet<>();
+        for (int idx = 0; idx < 20; idx++) {
+            Connection connection = DriverManager.getConnection(jdbcURL.build(), props);
+            contexts.add(connection.unwrap(YdbConnection.class).getCtx());
+            list.add(connection);
+        }
+
+        Assertions.assertEquals(3, contexts.size());
+
+        for (Connection connection: list) {
+            connection.close();
+        }
+    }
+
     @Test
     public void testResizeSessionPool() throws SQLException {
         String url = jdbcURL.build();