Fixed YdbExecutor concurrency problems (#68)

alex268 · web-flow · commit ed0e695ba275 · 2024-09-25T19:28:09.000+01:00
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -0,0 +1,93 @@
+name: YDB JDBC Driver CI
+
+on:
+  push:
+    branches:
+      - master
+      - release*
+  pull_request:
+      type: [opened, reopened, edited, synchronize]
+
+jobs:
+  prepare:
+    name: Prepare Maven cache
+    runs-on: ubuntu-latest
+
+    env:
+        MAVEN_ARGS: --batch-mode -Dstyle.color=always
+
+    steps:
+      - name: Checkout YDB JDBC Driver
+        uses: actions/checkout@v4
+        with:
+          path: jdbc
+
+      - name: Checkout YDB Java Examples
+        uses: actions/checkout@v4
+        with:
+          repository: ydb-platform/ydb-java-examples
+          path: examples
+
+      - name: Set up Java
+        uses: actions/setup-java@v4
+        with:
+          java-version: 17
+          distribution: 'temurin'
+          cache: 'maven'
+          cache-dependency-path: |
+              jdbc/pom.xml
+
+      - name: Download YDB JDBC Driver dependencies
+        working-directory: ./jdbc
+        run: mvn $MAVEN_ARGS dependency:go-offline
+
+      - name: Download YDB Java Examples dependencies
+        working-directory: ./examples
+        run: mvn $MAVEN_ARGS dependency:go-offline
+
+  build:
+    name: YDB JDBC Driver CI on JDK
+    runs-on: ubuntu-latest
+    needs: prepare
+
+    strategy:
+      matrix:
+        java: [ '8', '11', '17']
+
+    env:
+        MAVEN_ARGS: --batch-mode -Dstyle.color=always -DYDB_DOCKER_ISOLATION=true
+
+    steps:
+      - name: Checkout YDB JDBC Driver
+        uses: actions/checkout@v4
+        with:
+          path: jdbc
+
+      - name: Set up Java
+        uses: actions/setup-java@v4
+        with:
+          java-version: ${{ matrix.java }}
+          distribution: 'temurin'
+          cache: 'maven'
+          cache-dependency-path: |
+              jdbc/pom.xml
+
+      - name: Extract YDB JDBC Driver version
+        working-directory: ./jdbc
+        run: |
+          VERSION=$(mvn help:evaluate -Dexpression=project.version -q -DforceStdout)
+          echo "JDBC_VERSION=$VERSION" >> "$GITHUB_ENV"
+
+      - name: Build YDB JDBC Driver
+        working-directory: ./jdbc
+        run: mvn $MAVEN_ARGS install
+
+      - name: Checkout YDB Java Examples
+        uses: actions/checkout@v4
+        with:
+          repository: ydb-platform/ydb-java-examples
+          path: examples
+
+      - name: Test examples with Maven
+        working-directory: ./examples
+        run: mvn $MAVEN_ARGS -Dydb.jdbc.version=$JDBC_VERSION test
diff --git a/README.md b/README.md
@@ -1,6 +1,7 @@
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/ydb-platform/ydb-jdbc-driver/blob/master/LICENSE)
 [![Maven metadata URL](https://img.shields.io/maven-metadata/v?metadataUrl=https%3A%2F%2Frepo1.maven.org%2Fmaven2%2Ftech%2Fydb%2Fjdbc%2Fydb-jdbc-driver%2Fmaven-metadata.xml)](https://mvnrepository.com/artifact/tech.ydb.jdbc/ydb-jdbc-driver)
 [![Build](https://img.shields.io/github/actions/workflow/status/ydb-platform/ydb-jdbc-driver/build.yaml)](https://github.com/ydb-platform/ydb-jdbc-driver/actions/workflows/build.yaml)
+[![CI](https://img.shields.io/github/actions/workflow/status/ydb-platform/ydb-jdbc-driver/ci.yaml?label=CI)](https://github.com/ydb-platform/ydb-jdbc-driver/actions/workflows/ci.yaml)
 [![Codecov](https://img.shields.io/codecov/c/github/ydb-platform/ydb-jdbc-driver)](https://app.codecov.io/gh/ydb-platform/ydb-jdbc-driver)
 
 ## JDBC Driver for YDB
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,9 @@
 
 import java.sql.SQLException;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicReference;
 
+import tech.ydb.jdbc.YdbConst;
 import tech.ydb.jdbc.YdbStatement;
 import tech.ydb.jdbc.impl.YdbQueryResult;
 import tech.ydb.jdbc.query.QueryType;
@@ -17,9 +19,34 @@
  */
 public abstract class BaseYdbExecutor implements YdbExecutor {
     private final SessionRetryContext retryCtx;
+    private final AtomicReference<YdbQueryResult> currResult;
 
     public BaseYdbExecutor(YdbContext ctx) {
         this.retryCtx = ctx.getRetryCtx();
+        this.currResult = new AtomicReference<>();
+    }
+
+    protected void closeCurrentResult() throws SQLException {
+        YdbQueryResult rs = currResult.get();
+        if (rs != null) {
+            rs.close();
+        }
+    }
+
+    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 {
+        closeCurrentResult();
+        if (isClosed()) {
+            throw new SQLException(YdbConst.CLOSED_CONNECTION);
+        }
     }
 
     @Override
@@ -36,20 +63,20 @@ 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
     public YdbQueryResult executeBulkUpsert(YdbStatement statement, YdbQuery query, String tablePath, ListValue rows)
             throws SQLException {
         ensureOpened();
+
         String yql = query.getPreparedYql();
         YdbValidator validator = statement.getValidator();
-
         validator.execute(QueryType.BULK_QUERY + " >>\n" + yql,
                 () -> retryCtx.supplyStatus(session -> session.executeBulkUpsert(tablePath, rows))
         );
 
-        return new StaticQueryResult(query, Collections.emptyList());
+        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,7 +81,8 @@ protected QuerySession createNewQuerySession(YdbValidator validator) throws SQLE
     }
 
     @Override
-    public void close() {
+    public void close() throws SQLException {
+        closeCurrentResult();
         cleanTx();
         isClosed = true;
     }
@@ -96,6 +97,8 @@ private void cleanTx() {
 
     @Override
     public void setTransactionLevel(int level) throws SQLException {
+        ensureOpened();
+
         if (level == transactionLevel) {
             return;
         }
@@ -111,6 +114,8 @@ public void setTransactionLevel(int level) throws SQLException {
 
     @Override
     public void setReadOnly(boolean readOnly) throws SQLException {
+        ensureOpened();
+
         if (readOnly == isReadOnly) {
             return;
         }
@@ -125,6 +130,8 @@ public void setReadOnly(boolean readOnly) throws SQLException {
 
     @Override
     public void setAutoCommit(boolean autoCommit) throws SQLException {
+        ensureOpened();
+
         if (autoCommit == isAutoCommit) {
             return;
         }
@@ -137,12 +144,14 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
     }
 
     @Override
-    public boolean isClosed() {
+    public boolean isClosed() throws SQLException {
+        closeCurrentResult();
         return isClosed;
     }
 
     @Override
-    public String txID() {
+    public String txID() throws SQLException {
+        closeCurrentResult();
         return tx != null ? tx.getId() : null;
     }
 
@@ -234,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();
@@ -257,13 +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, () -> {
-                session.close();
-            });
+            return result.execute(stream, session::close);
         });
+
+        return updateCurrentResult(lazy);
     }
 
     @Override
@@ -283,7 +292,7 @@ public YdbQueryResult executeSchemeQuery(YdbStatement statement, YdbQuery query)
             );
         }
 
-        return new StaticQueryResult(query, Collections.emptyList());
+        return updateCurrentResult(new StaticQueryResult(query, Collections.emptyList()));
     }
 
     @Override
@@ -309,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/StreamQueryResult.java b/jdbc/src/main/java/tech/ydb/jdbc/context/StreamQueryResult.java
@@ -85,9 +85,12 @@ public CompletableFuture<Result<StreamQueryResult>> execute(QueryStream stream,
         return startFuture;
     }
 
-    public CompletableFuture<Result<StreamQueryResult>> execute(GrpcReadStream<ResultSetReader> stream) {
+    public CompletableFuture<Result<StreamQueryResult>> execute(
+            GrpcReadStream<ResultSetReader> stream, Runnable finish
+    ) {
         stream.start(rsr -> onResultSet(0, rsr))
-                .whenComplete(this::onStreamFinished);
+                .whenComplete(this::onStreamFinished)
+                .thenRun(finish);
         return startFuture;
     }
 
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,7 +51,8 @@ public TableServiceExecutor(YdbContext ctx, int transactionLevel, boolean autoCo
     }
 
     @Override
-    public void close() {
+    public void close() throws SQLException {
+        closeCurrentResult();
         tx = null;
     }
 
@@ -74,26 +75,31 @@ protected Session createNewTableSession(YdbValidator validator) throws SQLExcept
 
     @Override
     public void setTransactionLevel(int level) throws SQLException {
+        ensureOpened();
         updateState(tx.withTransactionLevel(level));
     }
 
     @Override
     public void setReadOnly(boolean readOnly) throws SQLException {
+        ensureOpened();
         updateState(tx.withReadOnly(readOnly));
     }
 
     @Override
     public void setAutoCommit(boolean autoCommit) throws SQLException {
+        ensureOpened();
         updateState(tx.withAutoCommit(autoCommit));
     }
 
     @Override
-    public boolean isClosed() {
+    public boolean isClosed() throws SQLException {
+        closeCurrentResult();
         return tx == null;
     }
 
     @Override
-    public String txID() {
+    public String txID() throws SQLException {
+        closeCurrentResult();
         return tx != null ? tx.txID() : null;
     }
 
@@ -191,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()));
         }
     }
 
@@ -220,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;
@@ -239,14 +245,16 @@ public YdbQueryResult executeScanQuery(YdbStatement statement, YdbQuery query, S
                 .withRequestTimeout(scanQueryTimeout)
                 .build();
 
-        final Session session = tx.getSession(validator);
+        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);
+            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
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java

Original file line number	Diff line number	Diff line change
`@@ -85,9 +85,12 @@ public CompletableFuture<Result<StreamQueryResult>> execute(QueryStream stream,`
`85`	`85`	`return startFuture;`
`86`	`86`	`}`
`87`	`87`
`88`		`- public CompletableFuture<Result<StreamQueryResult>> execute(GrpcReadStream<ResultSetReader> stream) {`
	`88`	`+ public CompletableFuture<Result<StreamQueryResult>> execute(`
	`89`	`+ GrpcReadStream<ResultSetReader> stream, Runnable finish`
	`90`	`+ ) {`
`89`	`91`	`stream.start(rsr -> onResultSet(0, rsr))`
`90`		`- .whenComplete(this::onStreamFinished);`
	`92`	`+ .whenComplete(this::onStreamFinished)`
	`93`	`+ .thenRun(finish);`
`91`	`94`	`return startFuture;`
`92`	`95`	`}`
`93`	`96`