From 0205bfdf4d2620353a30113c5c39bb537a55bd8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Thu, 26 Dec 2024 18:51:04 +0100
Subject: [PATCH] fix: clear interrupted flag after cancel

Clear the interrupted flag after cancelling a statement when using
a direct executor.

Fixes #1879
---
 .../spanner/jdbc/AbstractJdbcStatement.java     | 17 +++++++++++++++++
 .../spanner/jdbc/JdbcStatementTimeoutTest.java  |  1 -
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
index 5c15f336..3b98591a 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
@@ -34,6 +35,7 @@
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
@@ -47,6 +49,7 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
   final AbstractStatementParser parser;
   private final Lock executingLock;
   private volatile Thread executingThread;
+  private final AtomicBoolean cancelled = new AtomicBoolean();
   private boolean closed;
   private boolean closeOnCompletion;
   private boolean poolable;
@@ -259,10 +262,18 @@ private <T> T doWithStatementTimeout(
       connection.recordClientLibLatencyMetric(executionDuration.toMillis());
       return result;
     } catch (SpannerException spannerException) {
+      if (this.cancelled.get()
+          && spannerException.getErrorCode() == ErrorCode.CANCELLED
+          && this.executingLock != null) {
+        // Clear the interrupted flag of the thread.
+        //noinspection ResultOfMethodCallIgnored
+        Thread.interrupted();
+      }
       throw JdbcSqlExceptionFactory.of(spannerException);
     } finally {
       if (this.executingLock != null) {
         this.executingThread = null;
+        this.cancelled.set(false);
         this.executingLock.unlock();
       }
       if (shouldResetTimeout.apply(result)) {
@@ -374,8 +385,14 @@ public void cancel() throws SQLException {
       // This is a best-effort operation. It could be that the executing thread is set to null
       // between the if-check and the actual execution. Just ignore if that happens.
       try {
+        this.cancelled.set(true);
         this.executingThread.interrupt();
       } catch (NullPointerException ignore) {
+        // ignore, this just means that the execution finished before we got to the point where we
+        // could interrupt the thread.
+      } catch (SecurityException securityException) {
+        throw JdbcSqlExceptionFactory.of(
+            securityException.getMessage(), Code.PERMISSION_DENIED, securityException);
       }
     } else {
       connection.getSpannerConnection().cancel();
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java
index 2c8c43ca..7ce51d36 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java
@@ -161,7 +161,6 @@ public void testCancel() throws Exception {
                         message instanceof ExecuteSqlRequest
                             && ((ExecuteSqlRequest) message).getSql().equals(sql),
                     5000L);
-                System.out.println("Cancelling statement");
                 statement.cancel();
                 return null;
               });