fix(jdbc): cancel running statements when queries stop

crates/dbx-plugin-runtime/src/plugins.rs

@@ -431,7 +431,29 @@ impl PluginDriverSession {
         if let Some(logical) = &self.logical {
             return logical.invoke(method, params, timeout_duration).await;
         }
-        self.sidecar.invoke_with_timeout(method, params, Some(&self.driver_id), timeout_duration).await
+        let mut params = params;
+        let mut cancel = self.cancel_statement_on_drop(&mut params);
+        let result = self.sidecar.invoke_with_timeout(method, params, Some(&self.driver_id), timeout_duration).await;
+        if let Some(cancel) = cancel.as_mut() {
+            cancel.token = None;
+        }
+        result
+    }
+
+    /// The JDBC plugin runs one statement at a time on its shared connection, so a request the host
+    /// abandons (the user pressed stop) would otherwise keep the next request waiting until the
+    /// database finishes it. Tagging the request lets the plugin cancel exactly that statement.
+    fn cancel_statement_on_drop(&self, params: &mut serde_json::Value) -> Option<CancelStatementOnDrop> {
+        if self.driver_id != "jdbc" {
+            return None;
+        }
+        let token = uuid::Uuid::new_v4().to_string();
+        params.as_object_mut()?.insert(JDBC_CANCEL_TOKEN_PARAM.into(), token.clone().into());
+        Some(CancelStatementOnDrop {
+            sidecar: self.sidecar.clone(),
+            driver_id: self.driver_id.clone(),
+            token: Some(token),
+        })
     }
 
     pub async fn shutdown(&self) -> Result<(), String> {
@@ -473,6 +495,40 @@ impl PluginDriverSession {
     }
 }
 
+const JDBC_CANCEL_STATEMENT_METHOD: &str = "cancelStatement";
+const JDBC_CANCEL_TOKEN_PARAM: &str = "cancelToken";
+const JDBC_CANCEL_STATEMENT_TIMEOUT: Duration = Duration::from_secs(5);
+
+struct CancelStatementOnDrop {
+    sidecar: Arc<PluginSidecarSession>,
+    driver_id: String,
+    token: Option<String>,
+}
+
+impl Drop for CancelStatementOnDrop {
+    fn drop(&mut self) {
+        let Some(token) = self.token.take() else {
+            return;
+        };
+        let Ok(handle) = tokio::runtime::Handle::try_current() else {
+            return;
+        };
+        let sidecar = self.sidecar.clone();
+        let driver_id = std::mem::take(&mut self.driver_id);
+        handle.spawn(async move {
+            // Plugins without `cancelStatement` reply with an error once the abandoned statement ends.
+            let _: Result<serde_json::Value, _> = sidecar
+                .invoke_with_timeout(
+                    JDBC_CANCEL_STATEMENT_METHOD,
+                    serde_json::json!({ JDBC_CANCEL_TOKEN_PARAM: token }),
+                    Some(&driver_id),
+                    Some(JDBC_CANCEL_STATEMENT_TIMEOUT),
+                )
+                .await;
+        });
+    }
+}
+
 impl Drop for PluginDriverSession {
     fn drop(&mut self) {
         if let Some(logical) = &self.logical {
@@ -642,6 +698,89 @@ sleep 30
         let _ = std::fs::remove_dir_all(dir);
     }
 
+    #[cfg(unix)]
+    #[tokio::test]
+    async fn dropped_jdbc_request_cancels_its_statement_by_token() {
+        use std::os::unix::fs::PermissionsExt;
+
+        let dir = tempfile::tempdir().unwrap();
+        let executable = dir.path().join("plugin.py");
+        std::fs::write(
+            &executable,
+            r#"#!/usr/bin/python3
+import json, os, sys
+root = os.path.dirname(__file__)
+for line in sys.stdin:
+    request = json.loads(line)
+    params = request.get('params', {})
+    if request['method'] == 'cancelStatement':
+        with open(os.path.join(root, 'cancelled'), 'w') as out:
+            out.write(params.get('cancelToken', ''))
+        print(json.dumps({'id': request['id'], 'result': {'cancelled': True}}), flush=True)
+    elif request['method'] == 'executeQuery' and params.get('sql') == 'slow':
+        with open(os.path.join(root, 'running'), 'w') as out:
+            out.write(params.get('cancelToken', ''))
+    else:
+        print(json.dumps({'id': request['id'], 'result': {'token': params.get('cancelToken')}}), flush=True)
+"#,
+        )
+        .unwrap();
+        std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o755)).unwrap();
+        let plugin = InstalledPlugin::new(
+            PluginManifest {
+                id: "jdbc".to_string(),
+                name: "JDBC".to_string(),
+                version: "test".to_string(),
+                protocol_version: 1,
+                executable: Some("plugin.py".to_string()),
+                drivers: vec![PluginDriverManifest {
+                    id: "jdbc".to_string(),
+                    label: "JDBC".to_string(),
+                    kind: "external".to_string(),
+                    database_type: Some("jdbc".to_string()),
+                }],
+                ..PluginManifest::default()
+            },
+            dir.path().to_path_buf(),
+            env!("CARGO_PKG_VERSION"),
+        );
+        let session = std::sync::Arc::new(
+            PluginDriverSession::start_for_test(plugin, "jdbc".to_string(), PluginRuntimeEnv::default()).await.unwrap(),
+        );
+
+        let completed: serde_json::Value = session.invoke("listTables", serde_json::json!({})).await.unwrap();
+        assert!(completed["token"].is_string(), "every JDBC request carries a cancel token");
+
+        let running = dir.path().join("running");
+        let slow_session = session.clone();
+        let slow = tokio::spawn(async move {
+            slow_session.invoke::<serde_json::Value>("executeQuery", serde_json::json!({ "sql": "slow" })).await
+        });
+        tokio::time::timeout(std::time::Duration::from_secs(5), async {
+            while !running.exists() {
+                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .unwrap();
+        slow.abort();
+        assert!(slow.await.unwrap_err().is_cancelled());
+
+        let cancelled = dir.path().join("cancelled");
+        tokio::time::timeout(std::time::Duration::from_secs(5), async {
+            while !cancelled.exists() {
+                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .expect("dropping the request must send cancelStatement");
+        let slow_token = std::fs::read_to_string(&running).unwrap();
+        assert!(!slow_token.is_empty());
+        assert_eq!(std::fs::read_to_string(&cancelled).unwrap(), slow_token);
+        assert_ne!(completed["token"], slow_token.as_str(), "completed requests are never cancelled");
+        session.shutdown().await.unwrap();
+    }
+
     #[cfg(unix)]
     fn process_exists(pid: u32) -> bool {
         std::process::Command::new("kill")

plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java

@@ -64,6 +64,8 @@ public final class DbxJdbcPlugin {
     private static final ObjectMapper MAPPER =
         new ObjectMapper().configure(SerializationFeature.WRITE_BIGDECIMAL_AS_PLAIN, true);
     private static final int MAX_ROWS = 10_000;
+    static final String CANCEL_STATEMENT_METHOD = "cancelStatement";
+    static final String CANCEL_TOKEN_PARAM = "cancelToken";
     private static final String JDBCX_URL_PREFIX = "jdbcx:";
     private static final String JDBCX_EXTENSION_WHITELIST_PROPERTY = "jdbcx.extension.whitelist";
     private static final String JDBCX_HIGH_PRIVILEGE_EXTENSIONS_OPT_IN = "-Ddbx.jdbcx.allowHighPrivilegeExtensions=";
@@ -180,6 +182,7 @@ public final class DbxJdbcPlugin {
     private static final ThreadLocal<JdbcConnectionState> CONNECTION_STATE =
         ThreadLocal.withInitial(() -> DEFAULT_CONNECTION_STATE);
     private static final Map<String, JdbcConnectionState> LOGICAL_SESSIONS = new java.util.concurrent.ConcurrentHashMap<>();
+    private static final ThreadLocal<CancellationRequest> CANCELLATION_REQUEST = new ThreadLocal<>();
     private static volatile boolean shuttingDown;
 
     private static final class JdbcConnectionState {
@@ -187,11 +190,24 @@ public final class DbxJdbcPlugin {
         Connection sharedConnection;
         boolean manualTransactionActive;
         final Map<String, QuerySession> querySessions = new HashMap<>();
+        final Map<String, CancellationRequest> cancellationRequests = new HashMap<>();
         volatile Statement activeStatement;
         volatile boolean closing;
         final java.util.concurrent.locks.ReentrantLock operation = new java.util.concurrent.locks.ReentrantLock();
     }
 
+    static final class CancellationRequest {
+        final String token;
+        boolean cancelled;
+        Statement statement;
+
+        CancellationRequest(String token) {
+            this.token = token;
+        }
+    }
+
+    record StatementCancellation(boolean matched, Statement statement) {}
+
     private static JdbcConnectionState connectionState() {
         return CONNECTION_STATE.get();
     }
@@ -311,17 +327,35 @@ public final class DbxJdbcPlugin {
     }
 
     public static void main(String[] args) throws Exception {
-        java.util.concurrent.ExecutorService requests = java.util.concurrent.Executors.newCachedThreadPool();
         try (
             BufferedReader reader = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
             BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(System.out, StandardCharsets.UTF_8))
         ) {
+            serve(reader, writer);
+        }
+    }
+
+    static void serve(BufferedReader reader, BufferedWriter writer) throws Exception {
+        // Requests without a logical session run in order on one worker so the reader stays free to
+        // handle `cancelStatement` while a statement is still executing on the shared connection.
+        java.util.concurrent.ExecutorService defaultRequests = java.util.concurrent.Executors.newSingleThreadExecutor();
+        serve(reader, writer, defaultRequests);
+    }
+
+    static void serve(BufferedReader reader, BufferedWriter writer,
+                      java.util.concurrent.ExecutorService defaultRequests) throws Exception {
+        java.util.concurrent.ExecutorService requests = java.util.concurrent.Executors.newCachedThreadPool();
+        try {
             String line;
             while ((line = reader.readLine()) != null) {
                 if (line.isBlank()) {
                     continue;
                 }
                 JsonNode incoming = MAPPER.readTree(line);
+                if (CANCEL_STATEMENT_METHOD.equals(incoming.path("method").asText())) {
+                    writeResponse(writer, cancelStatement(incoming));
+                    continue;
+                }
                 if (incoming.path("params").hasNonNull("jdbcSessionId")) {
                     JdbcConnectionState retainedState = retainLogicalState(incoming);
                     requests.execute(() -> {
@@ -333,24 +367,153 @@ public final class DbxJdbcPlugin {
                     });
                     continue;
                 }
-                if ("close".equals(incoming.path("method").asText())) shuttingDown = true;
-                ObjectNode response = handleLine(line);
-                writeResponse(writer, response);
-                if (response.path("_dbx_close").asBoolean(false)) {
-                    break;
+                CancellationRequest cancellationRequest = registerCancellationRequest(incoming);
+                if ("close".equals(incoming.path("method").asText())) {
+                    shuttingDown = true;
+                    String closeLine = line;
+                    ObjectNode response = defaultRequests.submit(
+                        () -> handleDefaultRequest(closeLine, incoming, cancellationRequest)
+                    ).get();
+                    writeResponse(writer, response);
+                    if (response.path("_dbx_close").asBoolean(false)) {
+                        break;
+                    }
+                    continue;
                 }
+                String requestLine = line;
+                defaultRequests.execute(() -> {
+                    try {
+                        writeResponse(writer, handleDefaultRequest(requestLine, incoming, cancellationRequest));
+                    } catch (Exception error) {
+                        System.err.println("JDBC response failed: " + error.getClass().getSimpleName());
+                    }
+                });
             }
         } finally {
             shuttingDown = true;
+            defaultRequests.shutdown();
             for (JdbcConnectionState state : LOGICAL_SESSIONS.values()) {
                 closeLogicalSession(state);
             }
             LOGICAL_SESSIONS.clear();
             requests.shutdown();
+            if (!defaultRequests.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)) {
+                cancelActiveStatement(DEFAULT_CONNECTION_STATE);
+                defaultRequests.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS);
+            }
             closeSharedConnection();
         }
     }
 
+    private static ObjectNode handleDefaultRequest(String line, JsonNode request,
+                                                   CancellationRequest cancellationRequest) throws Exception {
+        CANCELLATION_REQUEST.set(cancellationRequest);
+        try {
+            if (isCancellationRequested(cancellationRequest)) return cancelledResponse(request);
+            return handleLine(line);
+        } finally {
+            finishCancellationRequest(cancellationRequest);
+            CANCELLATION_REQUEST.remove();
+        }
+    }
+
+    static CancellationRequest registerCancellationRequest(JsonNode request) {
+        String token = optionalText(request.path("params"), CANCEL_TOKEN_PARAM);
+        if (token == null) return null;
+        CancellationRequest cancellationRequest = new CancellationRequest(token);
+        JdbcConnectionState state = DEFAULT_CONNECTION_STATE;
+        synchronized (state) {
+            state.cancellationRequests.put(token, cancellationRequest);
+        }
+        return cancellationRequest;
+    }
+
+    private static boolean isCancellationRequested(CancellationRequest cancellationRequest) {
+        if (cancellationRequest == null) return false;
+        synchronized (DEFAULT_CONNECTION_STATE) {
+            return cancellationRequest.cancelled;
+        }
+    }
+
+    private static ObjectNode cancelledResponse(JsonNode request) {
+        JsonNode id = request.path("id");
+        ObjectNode response = MAPPER.createObjectNode();
+        response.set("id", id.isMissingNode() ? MAPPER.getNodeFactory().numberNode(1) : id);
+        response.set("error", MAPPER.createObjectNode().put("message", "JDBC request cancelled"));
+        return response;
+    }
+
+    static void finishCancellationRequest(CancellationRequest cancellationRequest) {
+        JdbcConnectionState state = DEFAULT_CONNECTION_STATE;
+        synchronized (state) {
+            Statement statement = cancellationRequest == null ? null : cancellationRequest.statement;
+            if (cancellationRequest != null) {
+                state.cancellationRequests.remove(cancellationRequest.token, cancellationRequest);
+                cancellationRequest.statement = null;
+            }
+            if (cancellationRequest == null || state.activeStatement == statement) {
+                state.activeStatement = null;
+            }
+        }
+    }
+
+    /**
+     * Cancels the statement of the request that carried the same cancel token. The host sends this
+     * when it abandons a request (the user pressed stop); a token that no longer matches belongs to a
+     * request that already finished, so a later request on the same connection is never cancelled.
+     */
+    private static ObjectNode cancelStatement(JsonNode request) {
+        JsonNode id = request.path("id");
+        ObjectNode response = MAPPER.createObjectNode();
+        response.set("id", id.isMissingNode() ? MAPPER.getNodeFactory().numberNode(1) : id);
+        String token = optionalText(request.path("params"), CANCEL_TOKEN_PARAM);
+        StatementCancellation cancellation = markCancellationRequested(token);
+        cancelCapturedStatement(cancellation.statement());
+        ObjectNode result = MAPPER.createObjectNode();
+        result.put("cancelled", cancellation.matched());
+        response.set("result", result);
+        return response;
+    }
+
+    static StatementCancellation markCancellationRequested(String token) {
+        JdbcConnectionState state = DEFAULT_CONNECTION_STATE;
+        synchronized (state) {
+            CancellationRequest cancellationRequest = token == null ? null : state.cancellationRequests.get(token);
+            if (cancellationRequest == null) return new StatementCancellation(false, null);
+            cancellationRequest.cancelled = true;
+            return new StatementCancellation(true, cancellationRequest.statement);
+        }
+    }
+
+    static void cancelCapturedStatement(Statement statement) {
+        if (statement == null) return;
+        try {
+            statement.cancel();
+        } catch (SQLException ignored) {
+            // The statement finished or was closed while the cancel was in flight.
+        }
+    }
+
+    private static void cancelActiveStatement(JdbcConnectionState state) {
+        cancelCapturedStatement(state.activeStatement);
+    }
+
+    private static void trackActiveStatement(Statement statement) {
+        trackActiveStatement(statement, CANCELLATION_REQUEST.get());
+    }
+
+    static void trackActiveStatement(Statement statement, CancellationRequest cancellationRequest) {
+        JdbcConnectionState state = connectionState();
+        boolean cancelRequested;
+        synchronized (state) {
+            state.activeStatement = statement;
+            if (cancellationRequest != null) cancellationRequest.statement = statement;
+            cancelRequested = cancellationRequest != null && cancellationRequest.cancelled;
+        }
+        // A cancel can arrive before the request created its statement.
+        if (cancelRequested) cancelCapturedStatement(statement);
+    }
+
     private static void writeResponse(BufferedWriter writer, ObjectNode response) throws IOException {
         synchronized (writer) {
             writer.write(MAPPER.writeValueAsString(response));
@@ -1247,7 +1410,7 @@ public final class DbxJdbcPlugin {
         boolean bitStringColumns = usesBitStringColumns(connection);
         boolean fullPrecisionIntegers = usesFullPrecisionIntegers(connection);
         try (Statement statement = conn.createStatement()) {
-            connectionState().activeStatement = statement;
+            trackActiveStatement(statement);
             if (connectionState().closing) throw new SQLException("JDBC logical session is closed");
             applyStatementOptions(statement, maxRows, fetchSize, timeoutSecs, quirks);
             String trimmedSql = trimStatementSql(sql);
@@ -1429,7 +1592,7 @@ public final class DbxJdbcPlugin {
         Statement statement;
         try {
             statement = createPagedQueryStatement(conn);
-            connectionState().activeStatement = statement;
+            trackActiveStatement(statement);
         } catch (Exception | LinkageError error) {
             restorePagedQueryTransaction(conn, restoreAutoCommit);
             throw error;
@@ -1535,7 +1698,7 @@ public final class DbxJdbcPlugin {
     }
 
     private static JsonNode readQuerySessionPage(QuerySession session, int pageSize) throws SQLException {
-        connectionState().activeStatement = session.statement;
+        trackActiveStatement(session.statement);
         if (connectionState().closing) throw new SQLException("JDBC logical session is closed");
         int effectivePageSize = Math.max(1, pageSize);
         ArrayNode rows = MAPPER.createArrayNode();

plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java

@@ -0,0 +1,253 @@
+package app.dbx.jdbc;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.BufferedReader;
+import java.io.BufferedWriter;
+import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
+import java.io.PipedInputStream;
+import java.io.PipedOutputStream;
+import java.lang.reflect.Proxy;
+import java.nio.charset.StandardCharsets;
+import java.sql.Statement;
+import java.util.Map;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+final class StatementCancelTest {
+    private static final ObjectMapper MAPPER = new ObjectMapper();
+    private static final String SLOW_SQL =
+        "SELECT COUNT(*) FROM SYSTEM_RANGE(1, 100000) A, SYSTEM_RANGE(1, 100000) B WHERE A.X + B.X > 0";
+
+    private final Map<Integer, BlockingQueue<JsonNode>> responses = new ConcurrentHashMap<>();
+    private BufferedWriter toPlugin;
+    private Thread server;
+
+    @AfterEach
+    void stopServer() throws Exception {
+        if (server == null) return;
+        toPlugin.close();
+        server.join(TimeUnit.SECONDS.toMillis(15));
+        var shuttingDown = DbxJdbcPlugin.class.getDeclaredField("shuttingDown");
+        shuttingDown.setAccessible(true);
+        shuttingDown.set(null, false);
+    }
+
+    @Test
+    void cancelStatementStopsTheAbandonedQuerySoTheNextOneRunsImmediately() throws Exception {
+        startServer();
+        assertNull(send(1, "connect", params()).path("error").textValue());
+
+        ObjectNode slow = params().put("sql", SLOW_SQL).put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "slow-token");
+        write(2, "executeQuery", slow);
+        Thread.sleep(500);
+        assertNull(responses.get(2).peek(), "the slow query must still be running when it is cancelled");
+
+        long started = System.nanoTime();
+        ObjectNode staleCancel = MAPPER.createObjectNode().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "finished-token");
+        assertFalse(send(3, DbxJdbcPlugin.CANCEL_STATEMENT_METHOD, staleCancel).path("result").path("cancelled").asBoolean());
+        ObjectNode cancel = MAPPER.createObjectNode().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "slow-token");
+        assertTrue(send(4, DbxJdbcPlugin.CANCEL_STATEMENT_METHOD, cancel).path("result").path("cancelled").asBoolean());
+
+        JsonNode next = send(5, "executeQuery", params().put("sql", "SELECT 1"));
+        assertEquals(1, next.path("result").path("rows").path(0).path(0).asInt(), next.toString());
+        assertTrue(TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - started) < 5);
+        assertTrue(await(2).has("error"), "the cancelled query reports an error");
+    }
+
+    @Test
+    void cancelStatementRemembersARequestWhileItIsQueued() throws Exception {
+        ExecutorService defaultRequests = Executors.newSingleThreadExecutor();
+        startServer(defaultRequests);
+        assertNull(send(1, "connect", params()).path("error").textValue());
+
+        CountDownLatch workerBlocked = new CountDownLatch(1);
+        CountDownLatch releaseWorker = new CountDownLatch(1);
+        defaultRequests.execute(() -> {
+            workerBlocked.countDown();
+            try {
+                releaseWorker.await();
+            } catch (InterruptedException error) {
+                Thread.currentThread().interrupt();
+            }
+        });
+        assertTrue(workerBlocked.await(5, TimeUnit.SECONDS));
+
+        try {
+            ObjectNode queued = params()
+                .put("sql", "CREATE TABLE CANCELLED_BEFORE_START(ID INT)")
+                .put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "queued-token");
+            write(2, "executeQuery", queued);
+            ObjectNode cancel = MAPPER.createObjectNode().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "queued-token");
+            assertTrue(send(3, DbxJdbcPlugin.CANCEL_STATEMENT_METHOD, cancel)
+                .path("result").path("cancelled").asBoolean());
+        } finally {
+            releaseWorker.countDown();
+        }
+
+        assertTrue(await(2).has("error"));
+        JsonNode tables = send(4, "executeQuery", params().put(
+            "sql",
+            "SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'CANCELLED_BEFORE_START'"
+        ));
+        assertEquals(0, tables.path("result").path("rows").path(0).path(0).asInt(), tables.toString());
+    }
+
+    @Test
+    void cancelStatementUsesTheStatementCapturedWithTheMatchingToken() throws Exception {
+        AtomicInteger abandonedCancels = new AtomicInteger();
+        AtomicInteger replacementCancels = new AtomicInteger();
+        Statement abandoned = statement(abandonedCancels);
+        Statement replacement = statement(replacementCancels);
+        DbxJdbcPlugin.CancellationRequest abandonedRequest = DbxJdbcPlugin.registerCancellationRequest(
+            request(1, "abandoned-token")
+        );
+        DbxJdbcPlugin.CancellationRequest replacementRequest = DbxJdbcPlugin.registerCancellationRequest(
+            request(2, "replacement-token")
+        );
+        DbxJdbcPlugin.trackActiveStatement(abandoned, abandonedRequest);
+
+        CountDownLatch statementCaptured = new CountDownLatch(1);
+        CountDownLatch replacementTracked = new CountDownLatch(1);
+        AtomicReference<Throwable> failure = new AtomicReference<>();
+        Thread cancel = new Thread(() -> {
+            try {
+                DbxJdbcPlugin.StatementCancellation cancellation =
+                    DbxJdbcPlugin.markCancellationRequested("abandoned-token");
+                assertTrue(cancellation.matched());
+                statementCaptured.countDown();
+                assertTrue(replacementTracked.await(5, TimeUnit.SECONDS));
+                DbxJdbcPlugin.cancelCapturedStatement(cancellation.statement());
+            } catch (Throwable error) {
+                failure.set(error);
+            }
+        });
+        cancel.start();
+
+        try {
+            assertTrue(statementCaptured.await(5, TimeUnit.SECONDS));
+            DbxJdbcPlugin.trackActiveStatement(replacement, replacementRequest);
+            replacementTracked.countDown();
+            cancel.join(TimeUnit.SECONDS.toMillis(5));
+            assertFalse(cancel.isAlive());
+            assertNull(failure.get());
+            assertEquals(1, abandonedCancels.get());
+            assertEquals(0, replacementCancels.get());
+        } finally {
+            replacementTracked.countDown();
+            DbxJdbcPlugin.finishCancellationRequest(abandonedRequest);
+            DbxJdbcPlugin.finishCancellationRequest(replacementRequest);
+        }
+    }
+
+    private void startServer() throws Exception {
+        startServer(Executors.newSingleThreadExecutor());
+    }
+
+    private void startServer(ExecutorService defaultRequests) throws Exception {
+        PipedOutputStream hostOut = new PipedOutputStream();
+        PipedInputStream pluginIn = new PipedInputStream(hostOut, 1 << 16);
+        PipedOutputStream pluginOut = new PipedOutputStream();
+        PipedInputStream hostIn = new PipedInputStream(pluginOut, 1 << 16);
+        toPlugin = new BufferedWriter(new OutputStreamWriter(hostOut, StandardCharsets.UTF_8));
+        server = new Thread(() -> {
+            try (
+                BufferedReader reader = new BufferedReader(new InputStreamReader(pluginIn, StandardCharsets.UTF_8));
+                BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(pluginOut, StandardCharsets.UTF_8))
+            ) {
+                DbxJdbcPlugin.serve(reader, writer, defaultRequests);
+            } catch (Exception ignored) {
+                // The pipe closes when the test ends.
+            }
+        });
+        server.setDaemon(true);
+        server.start();
+        Thread responseReader = new Thread(() -> {
+            try (BufferedReader reader = new BufferedReader(new InputStreamReader(hostIn, StandardCharsets.UTF_8))) {
+                String line;
+                while ((line = reader.readLine()) != null) {
+                    JsonNode response = MAPPER.readTree(line);
+                    queue(response.path("id").asInt()).add(response);
+                }
+            } catch (Exception ignored) {
+                // The pipe closes when the test ends.
+            }
+        });
+        responseReader.setDaemon(true);
+        responseReader.start();
+    }
+
+    private static ObjectNode params() {
+        ObjectNode connection = MAPPER.createObjectNode();
+        connection.put("connection_string", "jdbc:h2:mem:statement-cancel");
+        connection.put("jdbc_driver_class", "org.h2.Driver");
+        ObjectNode params = MAPPER.createObjectNode();
+        params.set("connection", connection);
+        return params;
+    }
+
+    private static ObjectNode request(int id, String token) {
+        ObjectNode request = MAPPER.createObjectNode().put("id", id).put("method", "executeQuery");
+        request.set("params", params().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, token));
+        return request;
+    }
+
+    private static Statement statement(AtomicInteger cancels) {
+        return (Statement) Proxy.newProxyInstance(
+            StatementCancelTest.class.getClassLoader(),
+            new Class<?>[] { Statement.class },
+            (proxy, method, args) -> {
+                if ("cancel".equals(method.getName())) cancels.incrementAndGet();
+                Class<?> returnType = method.getReturnType();
+                if (!returnType.isPrimitive()) return null;
+                if (returnType == boolean.class) return false;
+                if (returnType == byte.class) return (byte) 0;
+                if (returnType == short.class) return (short) 0;
+                if (returnType == int.class) return 0;
+                if (returnType == long.class) return 0L;
+                if (returnType == float.class) return 0F;
+                if (returnType == double.class) return 0D;
+                if (returnType == char.class) return '\0';
+                return null;
+            }
+        );
+    }
+
+    private BlockingQueue<JsonNode> queue(int id) {
+        return responses.computeIfAbsent(id, ignored -> new LinkedBlockingQueue<>());
+    }
+
+    private void write(int id, String method, ObjectNode params) throws Exception {
+        queue(id);
+        ObjectNode request = MAPPER.createObjectNode().put("jsonrpc", "2.0").put("id", id).put("method", method);
+        request.set("params", params);
+        toPlugin.write(request.toString());
+        toPlugin.newLine();
+        toPlugin.flush();
+    }
+
+    private JsonNode send(int id, String method, ObjectNode params) throws Exception {
+        write(id, method, params);
+        return await(id);
+    }
+
+    private JsonNode await(int id) throws Exception {
+        JsonNode response = queue(id).poll(10, TimeUnit.SECONDS);
+        assertNotNull(response, "no response for request " + id);
+        return response;
+    }
+}