From f3b1a9e2461006790182e534dbf8a98572bd710d Mon Sep 17 00:00:00 2001 From: Alex Kasko Date: Tue, 3 Jun 2025 11:25:08 +0100 Subject: [PATCH] Do not allow to cancel other running stmts (1.3) This is a backport of the PR #245 to `v1.3-ossivalis` stable branch. `Statement#cancel()` call works on connection level. So currently it is well possible to run query in one statement, and then interrupt this query calling cancel on another statement. This change adds a check that cancellation can only be performed if the query on current statement is still running. Otherwise `stmt.cancel()` call is a no-op. Testing: new test added that checks that other statement cannot be cancelled. --- .../java/org/duckdb/DuckDBConnection.java | 13 +++--- .../org/duckdb/DuckDBPreparedStatement.java | 41 ++++++++++++++++--- src/test/java/org/duckdb/TestClosure.java | 34 +++++++++++++++ 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/src/main/java/org/duckdb/DuckDBConnection.java b/src/main/java/org/duckdb/DuckDBConnection.java index c5177f1fb..9efbe9dfa 100644 --- a/src/main/java/org/duckdb/DuckDBConnection.java +++ b/src/main/java/org/duckdb/DuckDBConnection.java @@ -36,7 +36,7 @@ public final class DuckDBConnection implements java.sql.Connection { public static final String DEFAULT_SCHEMA = "main"; ByteBuffer connRef; - final Lock connRefLock = new ReentrantLock(); + final ReentrantLock connRefLock = new ReentrantLock(); final LinkedHashSet preparedStatements = new LinkedHashSet<>(); volatile boolean closing = false; @@ -488,14 +488,11 @@ void checkOpen() throws SQLException { * This function calls the underlying C++ interrupt function which aborts the query running on this connection. */ void interrupt() throws SQLException { - checkOpen(); - connRefLock.lock(); - try { - checkOpen(); - DuckDBNative.duckdb_jdbc_interrupt(connRef); - } finally { - connRefLock.unlock(); + if (!connRefLock.isHeldByCurrentThread()) { + throw new SQLException("Connection lock state error"); } + checkOpen(); + DuckDBNative.duckdb_jdbc_interrupt(connRef); } QueryProgress queryProgress() throws SQLException { diff --git a/src/main/java/org/duckdb/DuckDBPreparedStatement.java b/src/main/java/org/duckdb/DuckDBPreparedStatement.java index 44da9156c..eff4ca980 100644 --- a/src/main/java/org/duckdb/DuckDBPreparedStatement.java +++ b/src/main/java/org/duckdb/DuckDBPreparedStatement.java @@ -44,7 +44,7 @@ public class DuckDBPreparedStatement implements PreparedStatement { private DuckDBConnection conn; private ByteBuffer stmtRef = null; - final Lock stmtRefLock = new ReentrantLock(); + final ReentrantLock stmtRefLock = new ReentrantLock(); volatile boolean closeOnCompletion = false; private DuckDBResultSet selectResult = null; @@ -159,6 +159,11 @@ private boolean execute(boolean startTransaction) throws SQLException { checkOpen(); checkPrepared(); + // Wait with dispatching a new query if connection is locked by cancel() call + Lock connLock = getConnRefLock(); + connLock.lock(); + connLock.unlock(); + ByteBuffer resultRef = null; stmtRefLock.lock(); @@ -442,12 +447,27 @@ public void setQueryTimeout(int seconds) throws SQLException { @Override public void cancel() throws SQLException { checkOpen(); + // Only proceed to interrupt call after ensuring that the query on + // this statement is still running. + if (!stmtRefLock.isLocked()) { + return; + } + // Cancel is intended to be called concurrently with execute, + // thus we cannot take the statement lock that is held while + // query is running. NPE may be thrown if connection is closed + // concurrently. try { - // Cancel is intended to be called concurrently with execute, - // thus we cannot take the statement lock that is held while - // query is running. NPE may be thrown if connection is closed - // concurrently. - conn.interrupt(); + // Taking connection lock will prevent new queries to be executed + Lock connLock = getConnRefLock(); + connLock.lock(); + try { + if (!stmtRefLock.isLocked()) { + return; + } + conn.interrupt(); + } finally { + connLock.unlock(); + } } catch (NullPointerException e) { throw new SQLException(e); } @@ -1215,4 +1235,13 @@ private int[] intArrayFromLong(long[] arr) { } return res; } + + private Lock getConnRefLock() throws SQLException { + // NPE can be thrown if statement is closed concurrently. + try { + return conn.connRefLock; + } catch (NullPointerException e) { + throw new SQLException(e); + } + } } diff --git a/src/test/java/org/duckdb/TestClosure.java b/src/test/java/org/duckdb/TestClosure.java index db2fe53c4..fa25b414b 100644 --- a/src/test/java/org/duckdb/TestClosure.java +++ b/src/test/java/org/duckdb/TestClosure.java @@ -251,4 +251,38 @@ public static void test_results_fetch_no_hang() throws Exception { } } } + + public static void test_stmt_can_only_cancel_self() throws Exception { + try (Connection conn = DriverManager.getConnection(JDBC_URL); Statement stmt1 = conn.createStatement(); + Statement stmt2 = conn.createStatement()) { + stmt1.execute("DROP TABLE IF EXISTS test_fib1"); + stmt1.execute("CREATE TABLE test_fib1(i bigint, p double, f double)"); + stmt1.execute("INSERT INTO test_fib1 values(1, 0, 1)"); + long start = System.currentTimeMillis(); + Thread th = new Thread(() -> { + try { + Thread.sleep(200); + stmt1.cancel(); + } catch (Exception e) { + e.printStackTrace(); + } + }); + th.start(); + try ( + ResultSet rs = stmt2.executeQuery( + "WITH RECURSIVE cte AS (" + + + "SELECT * from test_fib1 UNION ALL SELECT cte.i + 1, cte.f, cte.p + cte.f from cte WHERE cte.i < 40000) " + + "SELECT avg(f) FROM cte")) { + rs.next(); + assertTrue(rs.getDouble(1) > 0); + } + th.join(); + long elapsed = System.currentTimeMillis() - start; + assertTrue(elapsed > 1000); + assertFalse(conn.isClosed()); + assertFalse(stmt1.isClosed()); + assertFalse(stmt2.isClosed()); + } + } }