Replace barriers by resultset keeping

alex268 · alex268 · commit 79db3eb9db2c · 2024-09-25T18:08:41.000+01:00
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java
@@ -2,7 +2,6 @@
 
 import java.sql.SQLException;
 import java.util.Collections;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 
 import tech.ydb.jdbc.YdbConst;
@@ -20,28 +19,31 @@
  */
 public abstract class BaseYdbExecutor implements YdbExecutor {
     private final SessionRetryContext retryCtx;
-    private final AtomicReference<Barrier> barrier;
+    private final AtomicReference<YdbQueryResult> currResult;
 
     public BaseYdbExecutor(YdbContext ctx) {
         this.retryCtx = ctx.getRetryCtx();
-        this.barrier = new AtomicReference<>(new Barrier());
-        this.barrier.get().open();
+        this.currResult = new AtomicReference<>();
     }
 
-    protected void checkBarrier() {
-        barrier.get().waitOpening();
+    protected void closeCurrentResult() throws SQLException {
+        YdbQueryResult rs = currResult.get();
+        if (rs != null) {
+            rs.close();
+        }
     }
 
-    protected Barrier createBarrier() {
-        Barrier newBarrier = new Barrier();
-        Barrier prev = barrier.getAndSet(newBarrier);
-        prev.waitOpening();
-        return newBarrier;
+    protected YdbQueryResult updateCurrentResult(YdbQueryResult result) throws SQLException {
+        YdbQueryResult old = currResult.getAndSet(result);
+        if (old != null) {
+            old.close();
+        }
+        return result;
     }
 
     @Override
     public void ensureOpened() throws SQLException {
-        checkBarrier();
+        closeCurrentResult();
         if (isClosed()) {
             throw new SQLException(YdbConst.CLOSED_CONNECTION);
         }
@@ -61,7 +63,7 @@ public YdbQueryResult executeSchemeQuery(YdbStatement statement, YdbQuery query)
                 () -> retryCtx.supplyStatus(session -> session.executeSchemeQuery(yql, settings))
         );
 
-        return new StaticQueryResult(query, Collections.emptyList());
+        return updateCurrentResult(new StaticQueryResult(query, Collections.emptyList()));
     }
 
     @Override
@@ -75,18 +77,6 @@ public YdbQueryResult executeBulkUpsert(YdbStatement statement, YdbQuery query,
                 () -> retryCtx.supplyStatus(session -> session.executeBulkUpsert(tablePath, rows))
         );
 
-        return new StaticQueryResult(query, Collections.emptyList());
-    }
-
-    public class Barrier {
-        private final CompletableFuture<Void> future = new CompletableFuture<>();
-
-        public void open() {
-            future.complete(null);
-        }
-
-        public void waitOpening() {
-            future.join();
-        }
+        return updateCurrentResult(new StaticQueryResult(query, Collections.emptyList()));
     }
 }
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java
@@ -81,8 +81,8 @@ protected QuerySession createNewQuerySession(YdbValidator validator) throws SQLE
     }
 
     @Override
-    public void close() {
-        checkBarrier();
+    public void close() throws SQLException {
+        closeCurrentResult();
         cleanTx();
         isClosed = true;
     }
@@ -144,14 +144,14 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
     }
 
     @Override
-    public boolean isClosed() {
-        checkBarrier();
+    public boolean isClosed() throws SQLException {
+        closeCurrentResult();
         return isClosed;
     }
 
     @Override
-    public String txID() {
-        checkBarrier();
+    public String txID() throws SQLException {
+        closeCurrentResult();
         return tx != null ? tx.getId() : null;
     }
 
@@ -243,7 +243,7 @@ public YdbQueryResult executeDataQuery(
             for (ResultSetReader rst: result) {
                 readers.add(new YdbStaticResultSet(statement, rst));
             }
-            return new StaticQueryResult(query, readers);
+            return updateCurrentResult(new StaticQueryResult(query, readers));
         } finally {
             if (!tx.isActive()) {
                 cleanTx();
@@ -256,8 +256,6 @@ public YdbQueryResult executeScanQuery(YdbStatement statement, YdbQuery query, S
             throws SQLException {
         ensureOpened();
 
-        Barrier barrier = createBarrier();
-
         YdbContext ctx = statement.getConnection().getCtx();
         YdbValidator validator = statement.getValidator();
 
@@ -268,11 +266,13 @@ public YdbQueryResult executeScanQuery(YdbStatement statement, YdbQuery query, S
 
         final QuerySession session = createNewQuerySession(validator);
         String msg = "STREAM_QUERY >>\n" + yql;
-        return validator.call(msg, () -> {
+        StreamQueryResult lazy = validator.call(msg, () -> {
             QueryStream stream = session.createQuery(yql, TxMode.SNAPSHOT_RO, params, settings);
             StreamQueryResult result = new StreamQueryResult(msg, statement, query, stream::cancel);
-            return result.execute(stream, barrier::open);
+            return result.execute(stream, session::close);
         });
+
+        return updateCurrentResult(lazy);
     }
 
     @Override
@@ -292,7 +292,7 @@ public YdbQueryResult executeSchemeQuery(YdbStatement statement, YdbQuery query)
             );
         }
 
-        return new StaticQueryResult(query, Collections.emptyList());
+        return updateCurrentResult(new StaticQueryResult(query, Collections.emptyList()));
     }
 
     @Override
@@ -318,7 +318,9 @@ public YdbQueryResult executeExplainQuery(YdbStatement statement, YdbQuery query
                 throw new SQLException("No explain data");
             }
 
-            return new StaticQueryResult(statement, res.getStats().getQueryAst(), res.getStats().getQueryPlan());
+            return updateCurrentResult(
+                    new StaticQueryResult(statement, res.getStats().getQueryAst(), res.getStats().getQueryPlan())
+            );
         }
     }
 
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java
@@ -51,8 +51,8 @@ public TableServiceExecutor(YdbContext ctx, int transactionLevel, boolean autoCo
     }
 
     @Override
-    public void close() {
-        checkBarrier();
+    public void close() throws SQLException {
+        closeCurrentResult();
         tx = null;
     }
 
@@ -92,14 +92,14 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
     }
 
     @Override
-    public boolean isClosed() {
-        checkBarrier();
+    public boolean isClosed() throws SQLException {
+        closeCurrentResult();
         return tx == null;
     }
 
     @Override
-    public String txID() {
-        checkBarrier();
+    public String txID() throws SQLException {
+        closeCurrentResult();
         return tx != null ? tx.txID() : null;
     }
 
@@ -197,7 +197,7 @@ public YdbQueryResult executeExplainQuery(YdbStatement statement, YdbQuery query
         try (Session session = createNewTableSession(validator)) {
             String msg = QueryType.EXPLAIN_QUERY + " >>\n" + yql;
             ExplainDataQueryResult res = validator.call(msg, () -> session.explainDataQuery(yql, settings));
-            return new StaticQueryResult(statement, res.getQueryAst(), res.getQueryPlan());
+            return updateCurrentResult(new StaticQueryResult(statement, res.getQueryAst(), res.getQueryPlan()));
         }
     }
 
@@ -226,7 +226,7 @@ public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, S
                 readers.add(new YdbStaticResultSet(statement, rs));
             }
 
-            return new StaticQueryResult(query, readers);
+            return updateCurrentResult(new StaticQueryResult(query, readers));
         } catch (SQLException | RuntimeException ex) {
             updateState(tx.withRollback(session));
             throw ex;
@@ -238,8 +238,6 @@ public YdbQueryResult executeScanQuery(YdbStatement statement, YdbQuery query, S
             throws SQLException {
         ensureOpened();
 
-        final Barrier barrier = createBarrier();
-
         YdbContext ctx = statement.getConnection().getCtx();
         YdbValidator validator = statement.getValidator();
         Duration scanQueryTimeout = ctx.getOperationProperties().getScanQueryTimeout();
@@ -250,14 +248,13 @@ public YdbQueryResult executeScanQuery(YdbStatement statement, YdbQuery query, S
         final Session session = createNewTableSession(validator);
 
         String msg = QueryType.SCAN_QUERY + " >>\n" + yql;
-        return validator.call(msg, () -> {
+        StreamQueryResult lazy = validator.call(msg, () -> {
             GrpcReadStream<ResultSetReader> stream = session.executeScanQuery(yql, params, settings);
             StreamQueryResult result = new StreamQueryResult(msg, statement, query, stream::cancel);
-            return result.execute(stream, () -> {
-                session.close();
-                barrier.open();
-            });
+            return result.execute(stream, session::close);
         });
+
+        return updateCurrentResult(lazy);
     }
 
     @Override
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbExecutor.java
@@ -13,11 +13,11 @@
  * @author Aleksandr Gorshenin
  */
 public interface YdbExecutor {
-    void close();
-    boolean isClosed();
+    void close() throws SQLException;
+    boolean isClosed() throws SQLException;
     void ensureOpened() throws SQLException;
 
-    String txID();
+    String txID() throws SQLException;
     int transactionLevel() throws SQLException;
 
     boolean isInsideTransaction() throws SQLException;
