Fix deadlock when fetching next chunk (1.3)

This is a backport of the PR duckdb#242 to `v1.3-ossivalis` stable branch.

When connection is being closed, it closes all statements and all result
sets on these statements. All this is done while holding connection
lock.

When fetching next chunk, result set needs to take both result set and
connection locks. Because of the wrong order of taking these locks,
concurrent `close()` call on result set, initiated from
`Connection#close()` was causing a deadlock.

This change reorders locks in `ResultSet#fetchChunk()` to fix this.

Testing: new test added that reproduces the deadlock reliably (may need
to raise the number of iterations).

Fixes: duckdb#241

Author: staticlibs

src/main/java/org/duckdb/DuckDBConnection.java

@@ -129,7 +129,7 @@ public void close() throws SQLException {
                 return;
             }
 
-            // Mark this instance as 'closing' to skip untrack call in
+            // Mark this instance as 'closing' to skip untrack logic in
             // prepared statements, that requires connection lock and can
             // cause a deadlock when the statement closure is caused by the
             // connection interrupt called by us.

src/main/java/org/duckdb/DuckDBPreparedStatement.java

@@ -361,7 +361,7 @@ public void close() throws SQLException {
 
                 // Untrack prepared statement from parent connection,
                 // if 'closing' flag is set it means that the parent connection itself
-                // is being closed and we don't need to call untrack from the statement.
+                // is being closed and we don't need to untrack this instance from the statement.
                 if (!conn.closing) {
                     conn.connRefLock.lock();
                     try {

src/main/java/org/duckdb/DuckDBResultSet.java

@@ -1399,22 +1399,25 @@ private void checkOpen() throws SQLException {
     }
 
     private DuckDBVector[] fetchChunk() throws SQLException {
-        // Take both result set and connection locks for fetching
-        resultRefLock.lock();
+        // Take both result set and connection locks for fetching,
+        // connection lock must be taken first because concurrent
+        // rs#close() call can be initiated from conn#close()
+        // that holds connection lock.
+        conn.connRefLock.lock();
         try {
-            checkOpen();
-            conn.connRefLock.lock();
+            conn.checkOpen();
+            resultRefLock.lock();
             try {
-                conn.checkOpen();
+                checkOpen();
                 return DuckDBNative.duckdb_jdbc_fetch(resultRef, conn.connRef);
             } finally {
-                conn.connRefLock.unlock();
+                resultRefLock.unlock();
             }
         } catch (SQLException e) {
             close();
             throw e;
         } finally {
-            resultRefLock.unlock();
+            conn.connRefLock.unlock();
         }
     }
 }

src/test/java/org/duckdb/TestClosure.java

@@ -5,6 +5,7 @@
 
 import java.io.File;
 import java.sql.*;
+import java.util.Properties;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -220,4 +221,34 @@ public static void test_results_close_prepared_stmt_no_crash() throws Exception
             }
         }
     }
+
+    public static void test_results_fetch_no_hang() throws Exception {
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        Properties config = new Properties();
+        config.put(DuckDBDriver.JDBC_STREAM_RESULTS, true);
+        long rowsCount = 1 << 24;
+        int iterations = 1;
+        for (int i = 0; i < iterations; i++) {
+            try (Connection conn = DriverManager.getConnection(JDBC_URL, config);
+                 Statement stmt = conn.createStatement();
+                 ResultSet rs = stmt.executeQuery("SELECT i, i::VARCHAR FROM range(0, " + rowsCount + ") AS t(i)")) {
+                executor.submit(() -> {
+                    try {
+                        Thread.sleep(100);
+                        conn.close();
+                    } catch (Exception e) {
+                        e.printStackTrace();
+                    }
+                });
+                long[] resultsCount = new long[1];
+                assertThrows(() -> {
+                    while (rs.next()) {
+                        resultsCount[0]++;
+                    }
+                }, SQLException.class);
+                assertTrue(resultsCount[0] > 0);
+                assertTrue(resultsCount[0] < rowsCount);
+            }
+        }
+    }
 }

src/test/java/org/duckdb/TestDuckDBJDBC.java

@@ -3455,7 +3455,7 @@ public static void test_query_progress() throws Exception {
                 @Override
                 public QueryProgress call() throws Exception {
                     try {
-                        Thread.sleep(1000);
+                        Thread.sleep(1500);
                         QueryProgress qp = stmt.getQueryProgress();
                         stmt.cancel();
                         return qp;