fix: prevent statement hang in case close() called when query is in p…

…rogress

Note: Statement#close() is still not thread-safe, however it is much more robust with the fix

fixes #1022

pgjdbc/src/main/java/org/postgresql/jdbc/PgPreparedStatement.java

@@ -182,24 +182,10 @@ protected boolean isOneShotQuery(CachedQuery cachedQuery) {
   }
 
   @Override
-  public void close() throws SQLException {
-    if (isClosed) {
-      return;
-    }
-
+  public void closeImpl() throws SQLException {
     if (preparedQuery != null) {
-      // See #368. We need to prevent closing the same statement twice
-      // Otherwise we might "release" a query that someone else is already using
-      // In other words, client does .close() as usual, however cleanup thread might fail to observe
-      // isClosed=true
-      synchronized (preparedQuery) {
-        if (!isClosed) {
-          ((PgConnection) connection).releaseQuery(preparedQuery);
-        }
-      }
+      ((PgConnection) connection).releaseQuery(preparedQuery);
     }
-
-    super.close();
   }
 
   public void setNull(int parameterIndex, int sqlType) throws SQLException {

pgjdbc/src/main/java/org/postgresql/jdbc/PgStatement.java

@@ -441,9 +441,7 @@ public void setCursorName(String name) throws SQLException {
     // No-op.
   }
 
-  // This is intentionally non-volatile to avoid performance hit in isClosed checks
-  // see #close()
-  protected boolean isClosed = false;
+  private volatile boolean isClosed = false;
 
   public int getUpdateCount() throws SQLException {
     checkClosed();
@@ -595,17 +593,27 @@ public void clearWarnings() throws SQLException {
    *
    * {@inheritDoc}
    */
-  public void close() throws SQLException {
+  public final void close() throws SQLException {
     // closing an already closed Statement is a no-op.
-    if (isClosed) {
-      return;
+    synchronized (this) {
+      if (isClosed) {
+        return;
+      }
+      isClosed = true;
     }
 
-    cleanupTimer();
+    cancel();
 
     closeForNextExecution();
 
-    isClosed = true;
+    closeImpl();
+  }
+
+  /**
+   * This is guaranteed to be called exactly once even in case of concurrent {@link #close()} calls.
+   * @throws SQLException in case of error
+   */
+  protected void closeImpl() throws SQLException {
   }
 
   /*
@@ -646,7 +654,7 @@ public boolean isUseServerPrepare() {
   }
 
   protected void checkClosed() throws SQLException {
-    if (isClosed) {
+    if (isClosed()) {
       throw new PSQLException(GT.tr("This statement has been closed."),
           PSQLState.OBJECT_NOT_IN_STATE);
     }
@@ -805,18 +813,20 @@ protected BatchResultHandler createBatchHandler(Query[] queries,
   }
 
   public void cancel() throws SQLException {
-    if (!STATE_UPDATER.compareAndSet(this, StatementCancelState.IN_QUERY, StatementCancelState.CANCELING)) {
+    if (statementState == StatementCancelState.IDLE) {
+      return;
+    }
+    if (!STATE_UPDATER.compareAndSet(this, StatementCancelState.IN_QUERY,
+        StatementCancelState.CANCELING)) {
       // Not in query, there's nothing to cancel
       return;
     }
-    try {
-      // Synchronize on connection to avoid spinning in killTimerTask
-      synchronized (connection) {
+    // Synchronize on connection to avoid spinning in killTimerTask
+    synchronized (connection) {
+      try {
         connection.cancelQuery();
-      }
-    } finally {
-      STATE_UPDATER.set(this, StatementCancelState.CANCELLED);
-      synchronized (connection) {
+      } finally {
+        STATE_UPDATER.set(this, StatementCancelState.CANCELLED);
         connection.notifyAll(); // wake-up killTimerTask
       }
     }
@@ -925,8 +935,10 @@ private void killTimerTask() {
     // "timeout error"
     // We wait till state becomes "cancelled"
     boolean interrupted = false;
-    while (!STATE_UPDATER.compareAndSet(this, StatementCancelState.CANCELLED, StatementCancelState.IDLE)) {
-      synchronized (connection) {
+    synchronized (connection) {
+      // state check is performed under synchronized so it detects "cancelled" state faster
+      // In other words, it prevents unnecessary ".wait()" call
+      while (!STATE_UPDATER.compareAndSet(this, StatementCancelState.CANCELLED, StatementCancelState.IDLE)) {
         try {
           // Note: wait timeout here is irrelevant since synchronized(connection) would block until
           // .cancel finishes

pgjdbc/src/main/java/org/postgresql/util/PSQLState.java

@@ -87,6 +87,7 @@
   OBJECT_NOT_IN_STATE("55000"),
   OBJECT_IN_USE("55006"),
 
+  QUERY_CANCELED("57014"),
 
   SYSTEM_ERROR("60000"),
   IO_ERROR("58030"),

pgjdbc/src/test/java/org/postgresql/test/jdbc2/StatementTest.java

@@ -18,6 +18,7 @@
 import org.postgresql.util.PSQLState;
 
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -27,7 +28,9 @@
 import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
-
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -792,6 +795,118 @@ public void testMultipleCancels() throws Exception {
     assertEquals(0, sharedTimer.getRefCount());
   }
 
+  @Test(timeout = 10000)
+  public void testCloseInProgressStatement() throws Exception {
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    final Connection outerLockCon = TestUtil.openDB();
+    outerLockCon.setAutoCommit(false);
+    //Acquire an exclusive lock so we can block the notice generating statement
+    outerLockCon.createStatement().execute("LOCK TABLE test_lock IN ACCESS EXCLUSIVE MODE;");
+
+    try {
+      con.createStatement().execute("SET SESSION client_min_messages = 'NOTICE'");
+      con.createStatement()
+          .execute("CREATE OR REPLACE FUNCTION notify_then_sleep() RETURNS VOID AS "
+              + "$BODY$ "
+              + "BEGIN "
+              + "RAISE NOTICE 'start';"
+              + "LOCK TABLE test_lock IN ACCESS EXCLUSIVE MODE;"
+              + "END "
+              + "$BODY$ "
+              + "LANGUAGE plpgsql;");
+      int cancels = 0;
+      for (int i = 0; i < 100; i++) {
+        final Statement st = con.createStatement();
+        executor.submit(new Callable<Void>() {
+          @Override
+          public Void call() throws Exception {
+            long start = System.currentTimeMillis();
+            while (st.getWarnings() == null) {
+              long dt = System.currentTimeMillis() - start;
+              if (dt > 10000) {
+                throw new IllegalStateException("Expected to receive a notice within 10 seconds");
+              }
+            }
+            st.close();
+            return null;
+          }
+        });
+        st.setQueryTimeout(120);
+        try {
+          st.execute("select notify_then_sleep()");
+        } catch (SQLException e) {
+          Assert.assertEquals(
+              "Query is expected to be cancelled via st.close(), got " + e.getMessage(),
+              PSQLState.QUERY_CANCELED.getState(),
+              e.getSQLState()
+          );
+          cancels++;
+        } finally {
+          TestUtil.closeQuietly(st);
+        }
+      }
+      Assert.assertNotEquals("At least one QUERY_CANCELED state is expected", 0, cancels);
+    } finally {
+      executor.shutdown();
+      TestUtil.closeQuietly(outerLockCon);
+    }
+  }
+
+  @Test(timeout = 20000)
+  public void testFastCloses() throws SQLException {
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    con.createStatement().execute("SET SESSION client_min_messages = 'NOTICE'");
+    con.createStatement()
+        .execute("CREATE OR REPLACE FUNCTION notify_then_sleep() RETURNS VOID AS "
+            + "$BODY$ "
+            + "BEGIN "
+            + "RAISE NOTICE 'start';"
+            + "EXECUTE pg_sleep(1);" // Note: timeout value does not matter here, we just test if test crashes or locks somehow
+            + "END "
+            + "$BODY$ "
+            + "LANGUAGE plpgsql;");
+    Map<String, Integer> cnt = new HashMap<String, Integer>();
+    final Random rnd = new Random();
+    for (int i = 0; i < 1000; i++) {
+      final Statement st = con.createStatement();
+      executor.submit(new Callable<Void>() {
+        @Override
+        public Void call() throws Exception {
+          int s = rnd.nextInt(10);
+          if (s > 8) {
+            Thread.sleep(s - 9);
+          }
+          st.close();
+          return null;
+        }
+      });
+      ResultSet rs = null;
+      String sqlState = "0";
+      try {
+        rs = st.executeQuery("select 1");
+        // Acceptable
+      } catch (SQLException e) {
+        sqlState = e.getSQLState();
+        if (!PSQLState.OBJECT_NOT_IN_STATE.getState().equals(sqlState)
+            && !PSQLState.QUERY_CANCELED.getState().equals(sqlState)) {
+          Assert.assertEquals(
+              "Query is expected to be cancelled via st.close(), got " + e.getMessage(),
+              PSQLState.QUERY_CANCELED.getState(),
+              e.getSQLState()
+          );
+        }
+      } finally {
+        TestUtil.closeQuietly(rs);
+        TestUtil.closeQuietly(st);
+      }
+      Integer val = cnt.get(sqlState);
+      val = (val == null ? 0 : val) + 1;
+      cnt.put(sqlState, val);
+    }
+    System.out.println("[testFastCloses] total counts for each sql state: " + cnt);
+    executor.shutdown();
+  }
+
   /**
    * Tests that calling {@code java.sql.Statement#close()} from a concurrent thread does not result
    * in {@link java.util.ConcurrentModificationException}