Added autoclosing of YdbContext

ydb-platform · Apr 29, 2024 · d094e58 · d094e58
1 parent dc75bfb
commit d094e58
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 25 deletions.
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/YdbDriver.java b/jdbc/src/main/java/tech/ydb/jdbc/YdbDriver.java
@@ -52,17 +52,25 @@ public YdbConnection connect(String url, Properties info) throws SQLException {
             config.getSafeProps()
         });
 
-        if (config.isCacheConnectionsInDriver()) {
-            return new YdbConnectionImpl(getCachedContext(config));
+        if (!config.isCacheConnectionsInDriver()) {
+            final YdbContext context = YdbContext.createContext(config);
+            return new YdbConnectionImpl(context) {
+                @Override
+                public void close() throws SQLException {
+                    super.close();
+                    context.close();
+                }
+            };
         }
 
-        // findOrCreateJdbcParams new context
-        final YdbContext context = YdbContext.createContext(config);
-        return new YdbConnectionImpl(context) {
+        YdbContext cached = getCachedContext(config);
+        return new YdbConnectionImpl(cached) {
             @Override
             public void close() throws SQLException {
                 super.close();
-                context.close();
+                if (!cached.hasConnections() && cache.remove(config, cached)) {
+                    cached.close();
+                }
             }
         };
     }

diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java
@@ -161,6 +161,10 @@ public void close() {
         }
     }
 
+    public boolean hasConnections() {
+        return connectionsCount.get() > 0;
+    }
+
     public void register() {
         int actual = connectionsCount.incrementAndGet();
         int maxSize = tableClient.sessionPoolStats().getMaxSize();

diff --git a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java b/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java
@@ -43,15 +43,13 @@ public void connect() throws SQLException {
 
     @Test
     public void testContextCache() throws SQLException {
-        YdbContext ctx;
-        try (Connection conn1 = DriverManager.getConnection(jdbcURL.build())) {
-            Assertions.assertTrue(conn1.isValid(5000));
+        Connection firstConnection = DriverManager.getConnection(jdbcURL.build());
+        Assertions.assertTrue(firstConnection.isValid(5000));
 
-            YdbConnection unwrapped = conn1.unwrap(YdbConnection.class);
-            Assertions.assertNotNull(unwrapped.getCtx());
+        YdbConnection first = firstConnection.unwrap(YdbConnection.class);
+        Assertions.assertNotNull(first.getCtx());
 
-            ctx = unwrapped.getCtx();
-        }
+        YdbContext ctx = first.getCtx();
 
         try (Connection conn2 = DriverManager.getConnection(jdbcURL.build())) {
             Assertions.assertTrue(conn2.isValid(5000));
@@ -87,10 +85,7 @@ public void testContextCache() throws SQLException {
             Assertions.assertNotSame(ctx, unwrapped.getCtx());
         }
 
-        if (YdbDriver.isRegistered()) {
-            YdbDriver.deregister();
-            YdbDriver.register();
-        }
+        firstConnection.close();
 
         try (Connection conn6 = DriverManager.getConnection(jdbcURL.build())) {
             Assertions.assertTrue(conn6.isValid(5000));
@@ -103,24 +98,27 @@ public void testContextCache() throws SQLException {
 
     @Test
     public void testContextCacheConncurrent() throws SQLException {
-        List<CompletableFuture<YdbContext>> list = new ArrayList<>();
+        List<CompletableFuture<YdbConnection>> list = new ArrayList<>();
 
         for (int idx = 0; idx < 20; idx++) {
             list.add(CompletableFuture.supplyAsync(() -> {
-                try (Connection conn1 = DriverManager.getConnection(jdbcURL.build())) {
-                    YdbConnection unwrapped = conn1.unwrap(YdbConnection.class);
-                    Assertions.assertNotNull(unwrapped.getCtx());
-                    return unwrapped.getCtx();
+                try {
+                    Connection connection = DriverManager.getConnection(jdbcURL.build());
+                    return connection.unwrap(YdbConnection.class);
                 } catch (SQLException ex) {
                     throw new RuntimeException("Cannot connect", ex);
                 }
             }));
         }
 
-        YdbContext first = list.get(0).join();
+        YdbContext first = list.get(0).join().getCtx();
+
+        for (CompletableFuture<YdbConnection> future: list) {
+            Assertions.assertEquals(first, future.join().getCtx());
+        }
 
-        for (CompletableFuture<YdbContext> future: list) {
-            Assertions.assertEquals(first, future.join());
+        for (CompletableFuture<YdbConnection> future: list) {
+            future.join().close();
         }
     }