diff --git a/crates/dbx-plugin-runtime/src/plugins.rs b/crates/dbx-plugin-runtime/src/plugins.rs index ed06c84da9..ae557bef9d 100644 --- a/crates/dbx-plugin-runtime/src/plugins.rs +++ b/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 { + 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, + driver_id: String, + token: Option, +} + +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 = 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::("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") diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 4deb4e3d26..69e0679848 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/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 CONNECTION_STATE = ThreadLocal.withInitial(() -> DEFAULT_CONNECTION_STATE); private static final Map LOGICAL_SESSIONS = new java.util.concurrent.ConcurrentHashMap<>(); + private static final ThreadLocal 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 querySessions = new HashMap<>(); + final Map 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(); diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java new file mode 100644 index 0000000000..377abff9b3 --- /dev/null +++ b/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> 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 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 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; + } +}