From c35a3dd9970fb25219b4bd2cd24f2da6ed20fe0a Mon Sep 17 00:00:00 2001 From: Vlad Ilyushchenko Date: Tue, 6 Oct 2026 18:52:12 +0100 Subject: [PATCH 1/2] feat(qwp): add per-query timeout to the query client Adds a query timeout to QwpQueryClient and the QuestDB facade: query_timeout_ms (connection-string default), Query.timeout(...) per handle, and an execute() overload with an explicit timeout. A timed-out query ends with the new STATUS_QUERY_TIMEOUT (QueryException.isTimeout()) on a connection that stays open and authenticated for the next query. - A server advertising the new CAP_QUERY_TIMEOUT (0x08) capability gets the remaining budget as a timeout_ms field after query_flags (QUERY_FLAG_TIMEOUT, 0x02) and ends the query itself with STATUS_QUERY_TIMEOUT (0x0E). Older servers are cancelled by the client at the deadline; the cancellation is reported as the same status. - Once the deadline passes no further result batch reaches the handler. Discarded batches are still decoded, keeping the connection-scoped SYMBOL dict in step with the server. - If the server does not end the query within the grace period (query_close_timeout_ms on the facade), the caller is released anyway and the connection drains the aborted query in the background. Only a connection that stays silent through a second grace period is closed. - The deadline spans failover: replays carry the remaining budget, the backoff and every reconnect step are bounded by it, and a reconnect it cut short is retried by the next execute() instead of leaving the client disconnected. Pool: Query.close() waits for a worker still draining a timed-out query before returning it, discards a worker whose connection failed, and a cancel issued while a submission waits behind such a drain is no longer lost. The server side (decoding timeout_ms, the breaker timeout, the status mapping and the capability advert) lands in a tandem questdb PR. --- README.md | 37 +- .../java/io/questdb/client/Completion.java | 4 + .../main/java/io/questdb/client/Query.java | 44 +- .../io/questdb/client/QueryException.java | 15 +- .../io/questdb/client/QuestDBBuilder.java | 6 + .../cutlass/qwp/client/QwpEgressIoThread.java | 40 +- .../cutlass/qwp/client/QwpEgressMsgKind.java | 19 + .../cutlass/qwp/client/QwpQueryClient.java | 560 +++++++++++- .../cutlass/qwp/client/QwpSpscQueue.java | 43 + .../cutlass/qwp/protocol/QwpConstants.java | 10 + .../io/questdb/client/impl/ConfigSchema.java | 1 + .../io/questdb/client/impl/QueryImpl.java | 89 +- .../io/questdb/client/impl/QueryLease.java | 6 + .../io/questdb/client/impl/QueryWorker.java | 51 ++ .../client/test/QueryTimeoutFacadeTest.java | 391 ++++++++ .../QwpQueryClientQueryTimeoutTest.java | 845 ++++++++++++++++++ .../cutlass/qwp/client/QwpSpscQueueTest.java | 69 ++ .../client/test/impl/QueryCloseDrainTest.java | 80 ++ .../client/test/impl/QueryImplResetTest.java | 5 + .../impl/QwpQueryClientConfigHonoredTest.java | 1 + .../query/QueryCancellationExample.java | 3 + .../example/query/QueryTimeoutExample.java | 66 ++ 22 files changed, 2349 insertions(+), 36 deletions(-) create mode 100644 core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java create mode 100644 core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java create mode 100644 examples/src/main/java/com/example/query/QueryTimeoutExample.java diff --git a/README.md b/README.md index a013a15a8..df9962e63 100644 --- a/README.md +++ b/README.md @@ -219,8 +219,34 @@ try (Query q = db.borrowQuery()) { ### Cancel or Time Out a Query -`submit()` returns a `Completion`. `await(timeout, unit)` returns `false` if the query is still in flight; `cancel()` -stops it. +Give a query a timeout with `timeout(...)`, or every query a default with `query_timeout_ms` in the configuration +string. The timeout bounds the whole query, measured from `submit()`. When it expires the query is stopped, `await()` +throws a `QueryException` whose `isTimeout()` is `true`, and the pooled connection stays open for the next query. + +```java +import io.questdb.client.QueryException; + +try (Query q = db.borrowQuery()) { + q.sql("SELECT * FROM big_table ORDER BY ts") + .handler(handler) + .timeout(5, TimeUnit.SECONDS); + try { + q.submit().await(); + } catch (QueryException e) { + if (!e.isTimeout()) { + throw e; + } + // the query ran longer than 5 seconds and was stopped + } +} +``` + +Servers that support per-query timeouts stop the query themselves; against older servers the client cancels it. If +the server does not end the query within `query_close_timeout_ms` of the timeout, `await()` throws anyway while the +connection finishes the aborted query in the background. + +`submit()` returns a `Completion`. `await(timeout, unit)` only bounds the wait: it returns `false` while the query keeps +running. `cancel()` stops the query. ```java import io.questdb.client.Completion; @@ -540,6 +566,7 @@ schema::key1=value1;key2=value2; | `query_pool_min` | `1` | Minimum query connections kept warm (`0` under `lazy_connect`) | | `query_pool_max` | `4` | Maximum query connections | | `acquire_timeout_ms` | `5000` | How long `borrowSender()`/`borrowQuery()` waits for a free slot | +| `query_close_timeout_ms` | `5000` | How long `Query.close()` waits for a running query; grace of `query_timeout_ms` | | `idle_timeout_ms` | `60000` | How long a pooled connection may stay idle before it is reaped | | `max_lifetime_ms` | `1800000` | Maximum lifetime of a pooled connection before it is recycled | @@ -554,6 +581,12 @@ Applied by the query pool to select and fail over between the nodes in the `addr | `zone` | | Prefer same-zone endpoints for `target=any`/`replica` (opaque, case-insensitive) | | `client_id` | | Opaque client identifier surfaced server-side for observability | +### Query keys + +| Key | Default | Description | +| ------------------ | ------- | ------------------------------------------------------------------------------------------ | +| `query_timeout_ms` | `0` | Default per-query timeout in milliseconds (`0` = none); override per query with `timeout()` | + The ingest side also accepts store-and-forward and reconnection tuning keys (`auto_flush_*`, `initial_connect_retry`, `reconnect_*`, `request_durable_ack`, `sf_*`, `max_frame_rejections`, `poison_min_escalation_window_millis`, …). `sf_max_total_bytes` caps the unacknowledged data all pooled senders buffer together — 128 MiB of memory by default, diff --git a/core/src/main/java/io/questdb/client/Completion.java b/core/src/main/java/io/questdb/client/Completion.java index 615799e07..2d4ff1cfe 100644 --- a/core/src/main/java/io/questdb/client/Completion.java +++ b/core/src/main/java/io/questdb/client/Completion.java @@ -63,6 +63,10 @@ public interface Completion { /** * Blocks up to the given timeout. Returns {@code true} if the query * completed, {@code false} on timeout. + *

+ * This bounds only how long the caller waits: on {@code false} the query + * keeps running. To stop queries that run too long, set a query timeout + * ({@link Query#timeout(long, TimeUnit)} or {@code query_timeout_ms}). * * @throws QueryException if the server reported an error or * {@link #cancel()} won the race diff --git a/core/src/main/java/io/questdb/client/Query.java b/core/src/main/java/io/questdb/client/Query.java index 539d57117..55454dd26 100644 --- a/core/src/main/java/io/questdb/client/Query.java +++ b/core/src/main/java/io/questdb/client/Query.java @@ -27,8 +27,10 @@ import io.questdb.client.cutlass.qwp.client.QwpBindSetter; import io.questdb.client.cutlass.qwp.client.QwpColumnBatchHandler; import io.questdb.client.cutlass.qwp.client.QwpServerInfo; +import io.questdb.client.cutlass.qwp.protocol.QwpConstants; import java.io.Closeable; +import java.util.concurrent.TimeUnit; /** * A query handle leased from the {@link QuestDB} pool via @@ -43,9 +45,9 @@ * creates one small lease handle per borrow (often scalar-replaced by the JIT * when used with try-with-resources). *

- * Lifecycle: configure with {@link #sql}, optional {@link #binds}, and - * {@link #handler}, then call {@link #submit()} to obtain a {@link Completion} - * and {@code await()} it before the next {@link #submit()}. + * Lifecycle: configure with {@link #sql}, optional {@link #binds} and + * {@link #timeout}, and {@link #handler}, then call {@link #submit()} to obtain + * a {@link Completion} and {@code await()} it before the next {@link #submit()}. *

* Thread safety: not thread-safe and single-flight -- one in-flight query per * handle. To run queries concurrently, borrow one handle per concurrent query. @@ -132,4 +134,40 @@ public interface Query extends Closeable { * before the first successful bind */ QwpServerInfo serverInfo(); + + /** + * Sets the query timeout for subsequent {@link #submit()} calls on this + * handle, overriding the {@code query_timeout_ms} default from the + * connection string; {@code 0} runs queries without a timeout. It stays in + * effect until changed, and every borrow starts from the configured + * default. + *

+ * The timeout is measured from {@code submit()} and bounds the whole query: + * server execution, any failover, and the time the handler spends in its + * callbacks (it is checked between result batches, so a handler blocked + * inside {@code onBatch} is not interrupted). When it expires, the query is + * stopped -- by the server when it supports per-query timeouts, otherwise by + * a client-side cancel -- no further result batch reaches the handler, the + * handler's {@code onError} receives + * {@link QwpConstants#STATUS_QUERY_TIMEOUT}, and {@link Completion#await()} + * throws a {@link QueryException} whose {@link QueryException#isTimeout()} + * is {@code true}. The pooled connection stays open and authenticated, and + * serves the next query. + *

+ * Should the server not end the query within {@code query_close_timeout_ms} + * (see {@link QuestDBBuilder#queryCloseTimeoutMillis(long)}) of the timeout, + * the caller is released with the timeout anyway while the connection + * finishes draining the aborted query; only a connection that stays silent + * for that long once more is closed and replaced. + *

+ * Unlike {@link Completion#await(long, TimeUnit)}, which only bounds how + * long the caller waits, the query timeout stops the query. + * + * @param timeout the timeout, {@code 0} for none; a positive value below + * one millisecond counts as one millisecond + * @param unit the unit of {@code timeout} + * @return this handle + * @throws IllegalArgumentException when {@code timeout} is negative + */ + Query timeout(long timeout, TimeUnit unit); } diff --git a/core/src/main/java/io/questdb/client/QueryException.java b/core/src/main/java/io/questdb/client/QueryException.java index a3fdedbd1..010cba919 100644 --- a/core/src/main/java/io/questdb/client/QueryException.java +++ b/core/src/main/java/io/questdb/client/QueryException.java @@ -24,6 +24,8 @@ package io.questdb.client; +import io.questdb.client.cutlass.qwp.protocol.QwpConstants; + /** * Thrown from {@link Completion#await()} / {@link Completion#await(long, java.util.concurrent.TimeUnit)} * when the server reported an error for the corresponding {@link Query}, @@ -32,7 +34,9 @@ *

* The original wire-level status byte is exposed via {@link #getStatus()} so * callers can distinguish cancellation from schema errors etc. without - * string-matching the message. + * string-matching the message. A query that ran past its timeout (see + * {@link Query#timeout(long, java.util.concurrent.TimeUnit)}) reports + * {@link QwpConstants#STATUS_QUERY_TIMEOUT}; {@link #isTimeout()} tests for it. */ public class QueryException extends RuntimeException { @@ -56,4 +60,13 @@ public QueryException(byte status, String message, Throwable cause) { public byte getStatus() { return status; } + + /** + * Returns {@code true} when the query ran past its timeout, whether the + * server or the client detected it. The connection that ran it stays + * usable. + */ + public boolean isTimeout() { + return status == QwpConstants.STATUS_QUERY_TIMEOUT; + } } diff --git a/core/src/main/java/io/questdb/client/QuestDBBuilder.java b/core/src/main/java/io/questdb/client/QuestDBBuilder.java index 11364f9ba..d3c70cb87 100644 --- a/core/src/main/java/io/questdb/client/QuestDBBuilder.java +++ b/core/src/main/java/io/questdb/client/QuestDBBuilder.java @@ -105,6 +105,12 @@ public QuestDBBuilder acquireTimeoutMillis(long millis) { * letting the pool grow a fresh one. Bounds the close of a handle whose * {@code submit()} is still running -- e.g. when the caller's own * {@code await(timeout)} expired and they gave up. Defaults to 5000ms. + *

+ * Also the grace period of the query timeout ({@code query_timeout_ms}, + * {@link Query#timeout(long, java.util.concurrent.TimeUnit)}): how long a + * timed-out query may take to end before the caller is released with the + * timeout anyway, and how long its connection may then take to drain the + * aborted query before it is closed instead of reused. */ public QuestDBBuilder queryCloseTimeoutMillis(long millis) { if (millis < 0) { diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java index 700e219ea..02e0d787d 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java @@ -114,6 +114,12 @@ public class QwpEgressIoThread implements Runnable, WebSocketFrameHandler { private boolean creditEnabled; private boolean currentQueryDone; private long currentRequestId = -1L; + // requestId of the last QUERY_REQUEST fully encoded into sendScratch. Once a + // request's id is published here the I/O thread no longer reads the caller's + // SQL text or bind scratch for it, so the caller may reuse them even while the + // query itself is still running. -1 until the first request. Written by the + // I/O thread only; volatile so the executing thread observes it. + private volatile long encodedRequestId = -1L; private volatile boolean shutdown; public QwpEgressIoThread(WebSocketClient wsClient, int bufferPoolSize, TerminalFailureListener terminalFailureListener) { @@ -356,6 +362,15 @@ public void run() { } } + /** + * Returns {@code true} once the {@code QUERY_REQUEST} for {@code requestId} + * has been encoded, after which the I/O thread no longer reads the SQL text or + * bind payload that {@link #submitQuery} handed it. + */ + public boolean isRequestEncoded(long requestId) { + return encodedRequestId == requestId; + } + /** * Signals shutdown. Does not join the thread -- caller handles that. */ @@ -377,6 +392,10 @@ public void shutdown() { * during {@link #sendQueryRequest} and does not retain a reference after * the send completes. {@code bindCount} is the number of binds the * payload contains; zero when the user supplied no binds. + *

+ * {@code timeoutMs} is written as the {@code timeout_ms} field only when + * {@code queryFlags} carries {@link QwpEgressMsgKind#QUERY_FLAG_TIMEOUT}; + * otherwise it is ignored. */ public void submitQuery( CharSequence sql, @@ -385,7 +404,8 @@ public void submitQuery( int bindCount, long bindPayloadPtr, long bindPayloadLen, - long queryFlags + long queryFlags, + long timeoutMs ) throws InterruptedException { pendingRequest.sql = sql; pendingRequest.requestId = requestId; @@ -394,6 +414,7 @@ public void submitQuery( pendingRequest.bindPayloadPtr = bindPayloadPtr; pendingRequest.bindPayloadLen = bindPayloadLen; pendingRequest.queryFlags = queryFlags; + pendingRequest.timeoutMs = timeoutMs; requests.put(pendingRequest); } @@ -404,6 +425,16 @@ public QueryEvent takeEvent() throws InterruptedException { return events.take(); } + /** + * Pops the next event, waiting at most until the absolute + * {@link System#nanoTime()} {@code deadlineNanos}. Returns {@code null} when + * the deadline passes with no event available. Called by the user thread + * during {@code execute()} when the query has a timeout. + */ + public QueryEvent takeEvent(long deadlineNanos) throws InterruptedException { + return events.take(deadlineNanos); + } + /** * Takes a pre-allocated {@link QueryEvent} from the pool for the I/O thread * to populate before offering to {@link #events}. Falls back to a fresh @@ -721,7 +752,13 @@ private void sendQueryRequest(QueryRequest req) { // stays byte-identical and the server defaults the flags to 0. if (req.queryFlags != 0) { sendScratch.putVarint(req.queryFlags); + // Flag-gated fields follow the flags in flag-bit order. + if ((req.queryFlags & QwpEgressMsgKind.QUERY_FLAG_TIMEOUT) != 0) { + sendScratch.putVarint(req.timeoutMs); + } } + // The request no longer references the caller's SQL text or bind scratch. + encodedRequestId = req.requestId; wsClient.sendBinary(sendScratch.getBufferPtr(), sendScratch.getPosition()); sendScratch.reset(); } @@ -783,5 +820,6 @@ private static final class QueryRequest { long queryFlags; long requestId; CharSequence sql; + long timeoutMs; } } diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressMsgKind.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressMsgKind.java index e9b2c8934..5f4d6c15c 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressMsgKind.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressMsgKind.java @@ -45,6 +45,18 @@ public final class QwpEgressMsgKind { * constant {@code io.questdb.cutlass.qwp.codec.QwpEgressMsgKind#CAP_QUERY_FLAGS}. */ public static final int CAP_QUERY_FLAGS = 0x00000002; + /** + * {@code SERVER_INFO.capabilities} bit: the server enforces a per-query + * timeout carried on {@code QUERY_REQUEST}. When set (together with + * {@link #CAP_QUERY_FLAGS}), a client that wants a timeout sets + * {@link #QUERY_FLAG_TIMEOUT} and appends {@code timeout_ms:varint} after the + * {@code query_flags} trailer; the server then ends an over-budget query with + * a {@code QUERY_ERROR} carrying + * {@link io.questdb.client.cutlass.qwp.protocol.QwpConstants#STATUS_QUERY_TIMEOUT}, + * leaving the connection open. Mirrors the server-side constant + * {@code io.questdb.cutlass.qwp.codec.QwpEgressMsgKind#CAP_QUERY_TIMEOUT}. + */ + public static final int CAP_QUERY_TIMEOUT = 0x00000008; /** * {@code SERVER_INFO.capabilities} bit advertising that the frame ends with * an additional {@code zone_id:u16_len+utf8} field after {@code node_id}. @@ -67,6 +79,13 @@ public final class QwpEgressMsgKind { * advertised {@link #CAP_QUERY_FLAGS}. */ public static final int QUERY_FLAG_RESET_DICT = 0x01; + /** + * {@code QUERY_REQUEST.query_flags} bit: a {@code timeout_ms:varint} field + * follows the {@code query_flags} varint. The server runs the query under + * that timeout instead of its default {@code query.timeout}. Sent only when + * the server advertised {@link #CAP_QUERY_TIMEOUT}. + */ + public static final int QUERY_FLAG_TIMEOUT = 0x02; public static final byte QUERY_REQUEST = 0x10; /** * Reset mask bit: clear the connection-scoped SYMBOL dict. diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java index bfc924e52..c99ac5a79 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java @@ -46,6 +46,7 @@ import java.util.Base64; import java.util.List; import java.util.Random; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -90,11 +91,23 @@ * query {@code QUERY_ERROR} responses are NOT terminal -- the connection * remains usable for the next query. *

+ * Query timeout: {@code query_timeout_ms=} (or {@link #withQueryTimeout} / + * the {@code timeoutMs} argument of {@link #execute}) bounds each query. A + * server advertising {@link QwpEgressMsgKind#CAP_QUERY_TIMEOUT} enforces the + * timeout itself; otherwise the client cancels the query when its deadline + * expires. Either way the handler sees + * {@link QwpConstants#STATUS_QUERY_TIMEOUT} and the connection, still + * authenticated, stays open for the next query. Only a connection that does not + * answer at all within the grace period ({@link #withQueryTimeoutGrace}) is + * replaced. + *

* Status byte convention on {@link QwpColumnBatchHandler#onError}: server- * emitted {@code QUERY_ERROR} frames surface with the server's status code * ({@link WebSocketResponse#STATUS_PARSE_ERROR}, {@code STATUS_INTERNAL_ERROR}, * {@link QwpConstants#STATUS_CANCELLED}, {@link QwpConstants#STATUS_LIMIT_EXCEEDED}, - * etc.). Failures detected client-side (closed client, bind encoding error, + * {@link QwpConstants#STATUS_QUERY_TIMEOUT}, etc.). A query timeout reports + * {@code STATUS_QUERY_TIMEOUT} whether the server or the client detected it. + * Failures detected client-side (closed client, bind encoding error, * truncated / unknown frame, decoder out of sync, I/O thread interrupt) all * surface with {@link WebSocketResponse#STATUS_INTERNAL_ERROR} and the * specific cause in the message. @@ -102,6 +115,16 @@ public class QwpQueryClient implements QuietCloseable { public static final String DEFAULT_ENDPOINT_PATH = "/read/v1"; + /** + * Default grace period of the query timeout, in milliseconds: how long the + * client waits past an expired query timeout for the server to end the + * query before it releases the caller, and again before it gives up on an + * unresponsive connection. Matches the {@code QuestDB} facade's + * {@code query_close_timeout_ms} default, which the facade passes per query. + * + * @see #withQueryTimeoutGrace(long) + */ + public static final long DEFAULT_QUERY_TIMEOUT_GRACE_MS = 5_000L; public static final int DEFAULT_WS_PORT = 9000; /** * Hard ceiling on {@link #withMaxBatchRows}. Matches the client decoder's @@ -152,6 +175,19 @@ public class QwpQueryClient implements QuietCloseable { */ private static final int DEFAULT_SERVER_INFO_TIMEOUT_MS = 5_000; private static final Logger LOG = LoggerFactory.getLogger(QwpQueryClient.class); + // Upper bound on any timeout or grace period converted to nanoseconds. Keeps + // deadline arithmetic (deadline + 2 * grace) on System.nanoTime() values + // clear of overflow; ~73 years is "no limit" for every practical purpose. + private static final long MAX_TIMEOUT_NANOS = Long.MAX_VALUE / 8; + // States of a query running under a timeout (see executeOnce). + // RUNNING: before the deadline, events are delivered normally. + private static final int TIMEOUT_PHASE_RUNNING = 0; + // TIMED_OUT: the deadline passed. Result batches are discarded while waiting, + // up to one grace period, for the server to end the query. + private static final int TIMEOUT_PHASE_TIMED_OUT = 1; + // DRAINING: the caller has been told about the timeout; the connection keeps + // draining the aborted query, up to one more grace period, so it can be reused. + private static final int TIMEOUT_PHASE_DRAINING = 2; // Reusable typed bind-value sink. Populated on the user thread by the // {@link QwpBindSetter} passed to execute(); the pre-encoded bytes are // handed to the I/O thread via QueryRequest. Allocated once per client to @@ -185,6 +221,12 @@ public class QwpQueryClient implements QuietCloseable { // across zones). private String clientZone; private int compressionLevel = 1; + // Absolute System.nanoTime() deadline bounding the endpoint walk of a failover + // reconnect performed on behalf of a query that has a timeout. Honoured only + // while connectDeadlineActive is set, which reconnectViaTracker() does for the + // duration of the walk. Accessed only by the executing thread. + private boolean connectDeadlineActive; + private long connectDeadlineNanos; // User-facing compression preference from the connection string. "raw" is // the library default -- no compression, no handshake header, no server- // side CPU burn on payloads where the network isn't the bottleneck @@ -262,6 +304,18 @@ public class QwpQueryClient implements QuietCloseable { // re-clamp) so a misconfigured server is observable from user code. private int negotiatedZstdLevel; private long nextRequestId = 1; + // Default per-query timeout, in milliseconds, applied by the execute() + // overloads that take no explicit timeout. 0 (the default) means no timeout. + // Set via query_timeout_ms or withQueryTimeout(); read once per execute(). + private volatile long queryTimeoutMs; + // Grace period of the query timeout; see DEFAULT_QUERY_TIMEOUT_GRACE_MS. + private volatile long queryTimeoutGraceMs = DEFAULT_QUERY_TIMEOUT_GRACE_MS; + // True after a failover reconnect failed for a reason other than + // authentication -- every endpoint was unreachable, or a query timeout cut the + // reconnect short. The next execute() then reconnects before running its + // query instead of rejecting it as "not connected". Cleared on every + // successful connect. Accessed by the executing thread only. + private boolean reconnectPending; // Cancel intent latched between {@link #cancel} and the point where // {@link #executeOnce} assigns {@link #currentRequestId}. Without this // latch, a cancel arriving in the dispatch window (after the user thread's @@ -332,6 +386,9 @@ private QwpQueryClient(String host, int port) { * {@link #execute}, reconnect to another endpoint and re-submit the query. * The user handler sees {@link QwpColumnBatchHandler#onFailoverReset} before * replayed batches begin arriving (batch_seq restarts at 0 on the new node). + *

  • {@code query_timeout_ms=N} -- default per-query timeout in milliseconds, + * applied by every {@link #execute} overload that takes no explicit timeout. + * {@code 0} (the default) means no timeout. See {@link #withQueryTimeout(long)}.
    • *
  • {@code username=;password=} -- HTTP Basic authentication. The client builds the * {@code Authorization: Basic } header from these. Server verifies the credentials * against the same user store the Postgres wire protocol uses, so a user created via @@ -407,6 +464,7 @@ public static QwpQueryClient fromConfig(CharSequence configurationString) { // over-int value must reject, not wrap. Integer connectTimeout = view.has("connect_timeout") ? view.getInt("connect_timeout", 0) : null; Long initialCredit = view.has("initial_credit") ? view.getLong("initial_credit", 0) : null; + Long queryTimeoutMs = view.has("query_timeout_ms") ? view.getLong("query_timeout_ms", 0) : null; int poolSize = view.getInt("buffer_pool_size", DEFAULT_IO_BUFFER_POOL_SIZE); String compression = view.getEnum("compression"); if (compression == null) { @@ -467,6 +525,9 @@ public static QwpQueryClient fromConfig(CharSequence configurationString) { if (initialCredit != null) { client.withInitialCredit(initialCredit); } + if (queryTimeoutMs != null) { + client.withQueryTimeout(queryTimeoutMs); + } client.withBufferPoolSize(poolSize); client.withCompression(compression, compressionLevel); if (tls) { @@ -522,6 +583,7 @@ public static void validateConfig(ConfigView view, boolean tls) { long backoffMax = view.getLong("failover_backoff_max_ms", -1); view.getLong("failover_max_duration_ms", -1); view.getLong("initial_credit", -1); + view.getLong("query_timeout_ms", -1); view.getLong("auth_timeout_ms", -1); // getInt: connect_timeout feeds an int API, so validation must also // reject values that fit a long but not an int. @@ -820,6 +882,7 @@ public synchronized void connect() { hostTracker.recordSuccess(i); currentEndpointIndex = i; connected = true; + reconnectPending = false; return; } if (lastObservedMismatch != null) { @@ -904,6 +967,54 @@ public void execute(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHandler * {@link QwpEgressMsgKind#CAP_QUERY_FLAGS}; otherwise it is silently ignored. */ public void execute(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHandler handler, boolean resetSymbolDict) { + execute(sql, binds, handler, resetSymbolDict, queryTimeoutMs); + } + + /** + * As {@link #execute(CharSequence, QwpBindSetter, QwpColumnBatchHandler, boolean)}, + * with an explicit query timeout that overrides the client default + * ({@link #withQueryTimeout(long)}, {@code query_timeout_ms}). + *

    + * The timeout bounds the whole call, measured from entry: binding, every + * failover reconnect and replay, server execution, and the time spent in + * the handler's own callbacks. It is checked between result batches, so a + * handler that blocks inside {@code onBatch} is not interrupted. + *

    + * When the server advertises {@link QwpEgressMsgKind#CAP_QUERY_TIMEOUT}, the + * timeout travels with the query and the server ends an over-budget query + * itself with a {@code QUERY_ERROR} carrying + * {@link QwpConstants#STATUS_QUERY_TIMEOUT}; the connection stays open and is + * reused by the next query. Against an older server the client cancels the + * query when its deadline expires and reports the cancellation as the same + * status. Either way, once the deadline has passed no further result batch + * reaches the handler. + *

    + * If the server has not ended the query within the grace period + * ({@link #withQueryTimeoutGrace(long)}) after the deadline, the handler is + * told about the timeout anyway and this call keeps draining the aborted + * query for up to one more grace period, so the connection can still be + * reused. Only a connection that stays silent through both grace periods is + * treated as failed (see {@link #hasTerminalFailure()}). Unless a handler + * callback blocks, this call therefore returns about two grace periods after + * the timeout at the latest. + * + * @param timeoutMs query timeout in milliseconds; {@code 0} runs the query + * without a timeout + * @throws IllegalArgumentException when {@code timeoutMs} is negative + */ + public void execute( + CharSequence sql, + QwpBindSetter binds, + QwpColumnBatchHandler handler, + boolean resetSymbolDict, + long timeoutMs + ) { + if (timeoutMs < 0) { + throw new IllegalArgumentException("timeoutMs must be >= 0"); + } + // The clock starts on entry, so binding, dispatch and any failover all + // count against the timeout. + final long deadlineNanos = System.nanoTime() + toBoundedNanos(timeoutMs); if (!executing.compareAndSet(false, true)) { throw new IllegalStateException( "QwpQueryClient.execute called while another execute is in flight; one query at a time per client"); @@ -915,7 +1026,7 @@ public void execute(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHandler // is intentionally NOT cleared inside executeOnce(). pendingCancel = false; try { - executeImpl(sql, binds, handler, resetSymbolDict); + executeImpl(sql, binds, handler, resetSymbolDict, timeoutMs, deadlineNanos); } finally { executing.set(false); } @@ -962,6 +1073,7 @@ public java.util.Map configSnapshotForTest() { m.put("failover_max_duration_ms", failoverMaxDurationMs); m.put("max_batch_rows", maxBatchRows); m.put("initial_credit", initialCreditBytes); + m.put("query_timeout_ms", queryTimeoutMs); m.put("buffer_pool_size", bufferPoolSize); m.put("compression", compressionPreference); m.put("compression_level", compressionLevel); @@ -1025,6 +1137,25 @@ public int getNegotiatedZstdLevel() { return negotiatedZstdLevel; } + /** + * Returns the grace period of the query timeout, in milliseconds. + * + * @see #withQueryTimeoutGrace(long) + */ + public long getQueryTimeoutGraceMs() { + return queryTimeoutGraceMs; + } + + /** + * Returns the default per-query timeout in milliseconds, {@code 0} when + * queries run without a timeout. + * + * @see #withQueryTimeout(long) + */ + public long getQueryTimeoutMs() { + return queryTimeoutMs; + } + /** * Returns the {@link QwpServerInfo} decoded from the currently-bound * server's {@code SERVER_INFO} frame, or {@code null} if the client is not @@ -1034,6 +1165,19 @@ public QwpServerInfo getServerInfo() { return serverInfo; } + /** + * Returns {@code true} when the bound connection has latched a terminal + * transport failure: the server closed it, a protocol fault occurred, or it + * stopped responding after a query timeout. The next {@link #execute} then + * reconnects ({@code failover=on}) or reports the stored failure + * ({@code failover=off}). A query error reported by the server, including a + * query timeout, is not a terminal failure: the connection stays usable. + */ + public boolean hasTerminalFailure() { + GenerationListener listener = currentGenerationListener; + return listener != null && listener.get() != null; + } + public boolean isConnected() { return connected; } @@ -1338,6 +1482,42 @@ public QwpQueryClient withMaxBatchRows(int rows) { return this; } + /** + * Sets the default per-query timeout, in milliseconds, applied by every + * {@link #execute} overload that takes no explicit timeout; {@code 0} + * disables it (the default). Programmatic equivalent of the + * {@code query_timeout_ms=} connection-string key. Unlike the connection + * settings it may be changed at any time; it applies to subsequent + * {@code execute()} calls. See + * {@link #execute(CharSequence, QwpBindSetter, QwpColumnBatchHandler, boolean, long)} + * for the timeout semantics. + */ + public QwpQueryClient withQueryTimeout(long timeoutMs) { + if (timeoutMs < 0) { + throw new IllegalArgumentException("query timeout must be >= 0"); + } + this.queryTimeoutMs = timeoutMs; + return this; + } + + /** + * Sets the grace period of the query timeout, in milliseconds (default + * {@value #DEFAULT_QUERY_TIMEOUT_GRACE_MS}). After a query's timeout + * expires, the client waits up to this long for the server to end the + * query before it reports the timeout to the handler, then up to this long + * again for the connection to drain the aborted query before it gives up on + * the connection. The {@code QuestDB} facade sets it from + * {@code query_close_timeout_ms}. May be changed at any time; it applies to + * subsequent {@code execute()} calls. + */ + public QwpQueryClient withQueryTimeoutGrace(long graceMs) { + if (graceMs < 0) { + throw new IllegalArgumentException("query timeout grace must be >= 0"); + } + this.queryTimeoutGraceMs = graceMs; + return this; + } + /** * Overrides the {@link #DEFAULT_SERVER_INFO_TIMEOUT_MS} wait for the * {@code SERVER_INFO} frame. Must be called before {@link #connect}. @@ -1436,6 +1616,10 @@ private static String defaultClientId() { return "questdb-java-egress/1.0.0"; } + private static boolean isPast(long deadlineNanos) { + return System.nanoTime() - deadlineNanos >= 0; + } + @SuppressWarnings("BooleanMethodIsAlwaysInverted") private static boolean matchesTarget(byte role, String target) { if (TARGET_ANY.equals(target)) { @@ -1452,6 +1636,92 @@ private static boolean matchesTarget(byte role, String target) { return true; } + private static String queryTimeoutMessage(long timeoutMs) { + return "query timeout of " + timeoutMs + "ms exceeded"; + } + + // Milliseconds left until deadlineNanos, rounded up so a sub-millisecond + // remainder still counts as 1; 0 once the deadline has passed. + private static long remainingMillisCeil(long deadlineNanos) { + long remainingNanos = deadlineNanos - System.nanoTime(); + return remainingNanos <= 0L ? 0L : (remainingNanos + 999_999L) / 1_000_000L; + } + + private static long toBoundedNanos(long millis) { + return Math.min(TimeUnit.MILLISECONDS.toNanos(millis), MAX_TIMEOUT_NANOS); + } + + /** + * Gives up on a connection that left a timed-out query unanswered through + * both grace periods (hung server, black-holed network), or whose I/O thread + * never even sent the query. Latches a terminal failure, so the next + * {@link #execute} replaces the connection ({@code failover=on}) or reports + * the failure ({@code failover=off}) and a pool discards the client, then + * stops the I/O thread. + *

    + * Runs on the executing thread, the sole consumer of the event queue, and + * keeps releasing every batch the I/O thread still publishes while it winds + * down: the I/O thread waits uninterruptibly for each published batch to be + * released, and would otherwise outlive the joins in {@link #close()} and + * {@link #cleanupFailedConnect()}. + */ + private void abandonUnresponsiveConnection(QwpEgressIoThread io) { + GenerationListener listener = currentGenerationListener; + if (listener != null) { + listener.onTerminalFailure(WebSocketResponse.STATUS_INTERNAL_ERROR, + "connection stopped responding after a query timeout"); + } + LOG.warn("QwpQueryClient connection did not end a timed-out query within two grace periods of {}ms; " + + "closing it", queryTimeoutGraceMs); + io.shutdown(); + Thread handle = ioThreadHandle; + if (handle == null) { + return; + } + handle.interrupt(); + final long stopDeadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(shutdownJoinMs); + boolean interrupted = false; + while (handle.isAlive() && !isPast(stopDeadlineNanos)) { + long pollDeadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(10); + if (pollDeadlineNanos - stopDeadlineNanos > 0) { + pollDeadlineNanos = stopDeadlineNanos; + } + QueryEvent ev; + try { + ev = io.takeEvent(pollDeadlineNanos); + } catch (InterruptedException e) { + // Finish stopping the I/O thread first; the flag is restored below. + interrupted = true; + continue; + } + if (ev != null) { + if (ev.kind == QueryEvent.KIND_BATCH && ev.buffer != null) { + io.releaseBuffer(ev.buffer); + } + io.releaseEvent(ev); + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + /** + * Applies the deadline of a failover reconnect made on behalf of a query + * with a timeout (see {@link #connectDeadlineActive}) to one endpoint step's + * configured timeout. {@code configuredMs <= 0} means the step has no bound + * of its own ({@code connect_timeout}'s OS-default sentinel) and passes + * through unchanged when no reconnect deadline is active. + */ + private int boundToConnectDeadline(long configuredMs) { + long boundedMs = configuredMs; + if (connectDeadlineActive) { + long remainingMs = Math.max(1L, remainingMillisCeil(connectDeadlineNanos)); + boundedMs = configuredMs > 0L ? Math.min(configuredMs, remainingMs) : remainingMs; + } + return (int) Math.min(boundedMs, Integer.MAX_VALUE); + } + /** * Builds the {@code X-QWP-Accept-Encoding} header value from the user's * preference. {@code raw} (the library default) omits the header entirely @@ -1549,7 +1819,7 @@ private void connectToEndpoint(Endpoint ep, String authHeader) { webSocketClient.setQwpClientId(clientId != null ? clientId : defaultClientId()); webSocketClient.setQwpAcceptEncoding(buildAcceptEncodingHeader()); webSocketClient.setQwpMaxBatchRows(maxBatchRows); - webSocketClient.setConnectTimeout(connectTimeoutMs); + webSocketClient.setConnectTimeout(boundToConnectDeadline(connectTimeoutMs)); runUpgradeWithTimeout(ep, authHeader); negotiatedQwpVersion = webSocketClient.getServerQwpVersion(); negotiatedZstdLevel = webSocketClient.getServerNegotiatedZstdLevel(); @@ -1573,12 +1843,27 @@ private void connectToEndpoint(Endpoint ep, String authHeader) { } } - private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHandler handler, boolean resetSymbolDict) { + private void executeImpl( + CharSequence sql, + QwpBindSetter binds, + QwpColumnBatchHandler handler, + boolean resetSymbolDict, + long timeoutMs, + long deadlineNanos + ) { if (closedFlag.get()) { throw new IllegalStateException("QwpQueryClient is closed"); } + final boolean timed = timeoutMs > 0; if (!connected) { - throw new IllegalStateException("QwpQueryClient not connected; call connect() first"); + if (!reconnectPending) { + throw new IllegalStateException("QwpQueryClient not connected; call connect() first"); + } + // An earlier failover reconnect did not complete. Retry it for this + // query instead of leaving the client unusable. + if (!reconnectForQuery(handler, timed, timeoutMs, deadlineNanos)) { + return; + } } hostTracker.beginRound(false); long failoverDeadlineNanos; @@ -1597,7 +1882,7 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa while (true) { attempt++; FailoverProbeHandler probe = new FailoverProbeHandler(handler); - executeOnce(sql, binds, probe, resetSymbolDict); + executeOnce(sql, binds, probe, resetSymbolDict, timeoutMs, deadlineNanos); if (!probe.transportFailureIntercepted) { return; } @@ -1605,6 +1890,12 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa handler.onError(probe.interceptedRequestId, probe.interceptedStatus, probe.interceptedMessage); return; } + if (timed && isPast(deadlineNanos)) { + handler.onError(probe.interceptedRequestId, QwpConstants.STATUS_QUERY_TIMEOUT, + queryTimeoutMessage(timeoutMs) + " before failover could replay the query; last error: " + + probe.interceptedMessage); + return; + } if (attempt >= failoverMaxAttempts || System.nanoTime() - failoverDeadlineNanos >= 0) { int failovers = Math.max(0, attempt - 1); handler.onError(probe.interceptedRequestId, probe.interceptedStatus, @@ -1621,6 +1912,10 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa } cleanupFailedConnect(); connected = false; + // Cleared by a successful reconnect below. Any other way out of this + // iteration leaves the client disconnected, and the next execute() + // reconnects before running its query. + reconnectPending = true; if (failoverInitialBackoffMs > 0L) { long base = failoverInitialBackoffMs << Math.min(attempt - 1, 30); if (base < 0L) base = failoverMaxBackoffMs; @@ -1648,6 +1943,13 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa if (delay > remaining) { delay = remaining; } + if (timed) { + // Never back off past the query's own deadline. + long queryRemaining = (deadlineNanos - System.nanoTime()) / 1_000_000L; + if (delay > queryRemaining) { + delay = Math.max(0L, queryRemaining); + } + } if (delay > 0L) { try { Thread.sleep(delay); @@ -1661,12 +1963,14 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa } } try { - reconnectViaTracker(); + reconnectViaTracker(timed, deadlineNanos); } catch (QwpAuthFailedException authErr) { // failover.md S6: AuthError is terminal across all hosts. // Credentials are cluster-wide, so retrying floods server logs // without recovery. Surface a distinct message so monitoring // can pull auth incidents apart from generic transport failures. + // Not retried by a later execute() either. + reconnectPending = false; handler.onError(probe.interceptedRequestId, probe.interceptedStatus, "auth failure during failover reconnect [host=" + authErr.getHost() + ':' + authErr.getPort() @@ -1674,6 +1978,13 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa + ", last error: " + probe.interceptedMessage + ']'); return; } catch (RuntimeException reconnectErr) { + if (timed && isPast(deadlineNanos)) { + handler.onError(probe.interceptedRequestId, QwpConstants.STATUS_QUERY_TIMEOUT, + queryTimeoutMessage(timeoutMs) + " while reconnecting for failover [last error: " + + probe.interceptedMessage + ", reconnect error: " + + reconnectErr.getMessage() + ']'); + return; + } handler.onError(probe.interceptedRequestId, probe.interceptedStatus, "failover reconnect failed after " + attempt + " attempt" + (attempt == 1 ? "" : "s") + " [last error: " @@ -1681,6 +1992,14 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa + reconnectErr.getMessage() + ']'); return; } + if (timed && isPast(deadlineNanos)) { + // Reconnected, but no budget is left to replay the query. The new + // connection is idle and serves the next query. + handler.onError(probe.interceptedRequestId, QwpConstants.STATUS_QUERY_TIMEOUT, + queryTimeoutMessage(timeoutMs) + " before failover could replay the query; last error: " + + probe.interceptedMessage); + return; + } handler.onFailoverReset(probe.interceptedRequestId, serverInfo); } } @@ -1689,8 +2008,28 @@ private void executeImpl(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHa * Inner loop for a single query attempt. Driven by {@link #execute}; wraps * the user's handler in a {@link FailoverProbeHandler} so that the outer * loop can intercept transport failures before they reach the user. + *

    + * With a timeout ({@code timeoutMs > 0}) the attempt runs the + * {@code TIMEOUT_PHASE_*} state machine. Before {@code deadlineNanos} + * (RUNNING) events are delivered as usual. Once it passes (TIMED_OUT), + * result batches are discarded -- releasing them keeps the server streaming + * toward its next timeout or cancel check -- and the server gets up to one + * grace period to end the query, cancelled by the client unless the server + * enforces the timeout itself. The caller learns about the timeout from that + * terminal frame, or, failing it, at the end of the grace period; the attempt + * then keeps draining (DRAINING) for up to one more grace period, so that the + * connection, which carries one query at a time, can serve the next query. + * A connection still silent after that is abandoned. The handler sees + * exactly one terminal callback in every case. */ - private void executeOnce(CharSequence sql, QwpBindSetter binds, FailoverProbeHandler probe, boolean resetSymbolDict) { + private void executeOnce( + CharSequence sql, + QwpBindSetter binds, + FailoverProbeHandler probe, + boolean resetSymbolDict, + long timeoutMs, + long deadlineNanos + ) { // Cache the I/O thread reference at entry: close() may null the field while // we are inside this loop, so reading the field per-iteration would NPE // exactly when the user is mid-execute() and close() races. The queue and @@ -1724,6 +2063,27 @@ private void executeOnce(CharSequence sql, QwpBindSetter binds, FailoverProbeHan return; } } + final boolean timed = timeoutMs > 0; + long wireTimeoutMs = 0L; + if (timed) { + wireTimeoutMs = remainingMillisCeil(deadlineNanos); + if (wireTimeoutMs == 0L) { + // The budget ran out before the query could be sent (a slow bind + // setter, or a failover reconnect that used it up). Nothing is on + // the wire, so the connection is untouched. + bindValues.reset(); + probe.onError(-1L, QwpConstants.STATUS_QUERY_TIMEOUT, + queryTimeoutMessage(timeoutMs) + " before the query was sent"); + return; + } + } + final long queryFlags = resolveQueryFlags(resetSymbolDict, timed); + // A server that advertised CAP_QUERY_TIMEOUT receives the remaining budget + // with the query and ends it itself with STATUS_QUERY_TIMEOUT. The + // client-side deadline is then only a backstop and must not race the + // server's own report with a CANCEL. + final boolean serverEnforcesTimeout = (queryFlags & QwpEgressMsgKind.QUERY_FLAG_TIMEOUT) != 0; + final long graceNanos = timed ? toBoundedNanos(queryTimeoutGraceMs) : 0L; long requestId = nextRequestId++; currentRequestId = requestId; // Honor a cancel that arrived during the dispatch window. The latch @@ -1733,34 +2093,122 @@ private void executeOnce(CharSequence sql, QwpBindSetter binds, FailoverProbeHan if (pendingCancel) { io.requestCancel(requestId); } + int phase = TIMEOUT_PHASE_RUNNING; + // Set once a result batch is withheld from the handler because the + // deadline had passed; a later RESULT_END is then not a success. + boolean discardedBatch = false; + // Whether the user had already cancelled when the deadline passed. A + // STATUS_CANCELLED reply then reports that cancel, not the timeout. + boolean cancelledByUser = false; + long waitDeadlineNanos = deadlineNanos; try { io.submitQuery(sql, requestId, initialCreditBytes, bindValues.count(), bindValues.bufferPtr(), bindValues.bufferLen(), - resolveQueryFlags(resetSymbolDict)); + queryFlags, wireTimeoutMs); while (true) { - QueryEvent ev = io.takeEvent(); + QueryEvent ev = timed ? io.takeEvent(waitDeadlineNanos) : io.takeEvent(); + if (ev == null) { + // waitDeadlineNanos passed without an event. + if (phase == TIMEOUT_PHASE_RUNNING) { + cancelledByUser = pendingCancel; + phase = TIMEOUT_PHASE_TIMED_OUT; + if (!serverEnforcesTimeout) { + io.requestCancel(requestId); + } + waitDeadlineNanos = deadlineNanos + graceNanos; + continue; + } + // A request the I/O thread never encoded means the thread is + // wedged; it may still read the caller's SQL and bind buffers, + // so the caller must not be released before it is stopped. + if (phase == TIMEOUT_PHASE_TIMED_OUT && io.isRequestEncoded(requestId)) { + // The server has not ended the query within the grace + // period. Tell the caller now; keep draining the aborted + // query so the connection can serve the next one. + io.requestCancel(requestId); + phase = TIMEOUT_PHASE_DRAINING; + waitDeadlineNanos = deadlineNanos + 2 * graceNanos; + probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs) + + "; the server did not end the query within the " + queryTimeoutGraceMs + + "ms grace period"); + continue; + } + final boolean callerWaiting = phase != TIMEOUT_PHASE_DRAINING; + abandonUnresponsiveConnection(io); + if (callerWaiting) { + probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs) + + "; the connection did not respond and was closed"); + } + return; + } try { switch (ev.kind) { case QueryEvent.KIND_BATCH: try { - probe.onBatch(ev.buffer.batch); + if (phase == TIMEOUT_PHASE_RUNNING && timed && isPast(deadlineNanos)) { + // The deadline passed while this batch was queued, + // or while the handler worked on the previous one. + cancelledByUser = pendingCancel; + phase = TIMEOUT_PHASE_TIMED_OUT; + if (!serverEnforcesTimeout) { + io.requestCancel(requestId); + } + waitDeadlineNanos = deadlineNanos + graceNanos; + } + if (phase == TIMEOUT_PHASE_RUNNING) { + probe.onBatch(ev.buffer.batch); + } else { + // Still decoded by the I/O thread, which keeps the + // connection-scoped SYMBOL dict in step with the + // server; just never shown to the handler. + discardedBatch = true; + } } finally { io.releaseBuffer(ev.buffer); } break; case QueryEvent.KIND_END: - probe.onEnd(requestId, ev.totalRows); + if (phase == TIMEOUT_PHASE_RUNNING + || (phase == TIMEOUT_PHASE_TIMED_OUT && !discardedBatch)) { + // Past the deadline, a complete result that withheld + // nothing from the handler still counts as success. + probe.onEnd(requestId, ev.totalRows); + } else if (phase == TIMEOUT_PHASE_TIMED_OUT) { + probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs)); + } return; case QueryEvent.KIND_EXEC_DONE: - probe.onExecDone(requestId, ev.opType, ev.rowsAffected); + if (phase != TIMEOUT_PHASE_DRAINING) { + // A statement that completed took effect; report it as + // done even when the reply came past the deadline. + probe.onExecDone(requestId, ev.opType, ev.rowsAffected); + } return; case QueryEvent.KIND_ERROR: - probe.onError(requestId, ev.errorStatus, ev.errorMessage); + if (phase == TIMEOUT_PHASE_TIMED_OUT + && ev.errorStatus == QwpConstants.STATUS_CANCELLED + && !cancelledByUser) { + // The server honoured the cancel the deadline sent. + probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs)); + } else if (phase != TIMEOUT_PHASE_DRAINING) { + // Includes the server's own STATUS_QUERY_TIMEOUT report. + probe.onError(requestId, ev.errorStatus, ev.errorMessage); + } return; case QueryEvent.KIND_TRANSPORT_ERROR: - probe.markTransportFailure(requestId, ev.errorStatus, ev.errorMessage); + if (phase == TIMEOUT_PHASE_RUNNING) { + probe.markTransportFailure(requestId, ev.errorStatus, ev.errorMessage); + } else if (phase == TIMEOUT_PHASE_TIMED_OUT) { + // Past the deadline there is nothing left to fail over + // for. The I/O thread latched the failure, so the next + // execute() replaces the connection. + probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs) + + "; the connection then failed: " + ev.errorMessage); + } return; default: - probe.onError(requestId, WebSocketResponse.STATUS_INTERNAL_ERROR, "unknown event kind " + ev.kind); + if (phase != TIMEOUT_PHASE_DRAINING) { + probe.onError(requestId, WebSocketResponse.STATUS_INTERNAL_ERROR, "unknown event kind " + ev.kind); + } return; } } finally { @@ -1774,7 +2222,10 @@ private void executeOnce(CharSequence sql, QwpBindSetter binds, FailoverProbeHan } catch (InterruptedException ie) { Thread.currentThread().interrupt(); // Interrupt on the user thread is not a transport failure; surface directly. - probe.deliverFinal(requestId, "interrupted while waiting for server response"); + // A draining attempt has already given the handler its terminal callback. + if (phase != TIMEOUT_PHASE_DRAINING) { + probe.deliverFinal(requestId, "interrupted while waiting for server response"); + } } finally { currentRequestId = -1L; } @@ -1823,7 +2274,7 @@ private void probeZstdAvailable() { private QwpServerInfo receiveServerInfoSync() { ServerInfoReceiver receiver = new ServerInfoReceiver(); - long deadlineMs = System.currentTimeMillis() + serverInfoTimeoutMs; + long deadlineMs = System.currentTimeMillis() + boundToConnectDeadline(serverInfoTimeoutMs); while (receiver.info == null && receiver.decodeError == null && receiver.closeCode < 0 @@ -1847,6 +2298,35 @@ private QwpServerInfo receiveServerInfoSync() { return receiver.info; } + /** + * Retries, on behalf of the query about to run, a failover reconnect that an + * earlier {@code execute()} could not complete. On failure the handler gets + * the query's single terminal callback and the reconnect stays pending. + * + * @return {@code true} when the client is connected again + */ + private boolean reconnectForQuery(QwpColumnBatchHandler handler, boolean timed, long timeoutMs, long deadlineNanos) { + try { + reconnectViaTracker(timed, deadlineNanos); + return true; + } catch (QwpAuthFailedException authErr) { + reconnectPending = false; + handler.onError(-1L, WebSocketResponse.STATUS_INTERNAL_ERROR, + "auth failure during reconnect [host=" + authErr.getHost() + ':' + authErr.getPort() + + ", status=" + authErr.getStatusCode() + ']'); + } catch (RuntimeException e) { + if (timed && isPast(deadlineNanos)) { + handler.onError(-1L, QwpConstants.STATUS_QUERY_TIMEOUT, + queryTimeoutMessage(timeoutMs) + " while reconnecting [reconnect error: " + + e.getMessage() + ']'); + } else { + handler.onError(-1L, WebSocketResponse.STATUS_INTERNAL_ERROR, + "reconnect failed [reconnect error: " + e.getMessage() + ']'); + } + } + return false; + } + /** * Walks the endpoint list by tracker priority (HEALTHY → UNKNOWN → * TRANSIENT_REJECT → TRANSPORT_ERROR → TOPOLOGY_REJECT). The mid-stream @@ -1874,6 +2354,12 @@ private void reconnectViaTracker() { // as a per-endpoint transport error retried across every host. String authHeader = resolveAuthorizationHeader(); while (true) { + if (connectDeadlineActive && isPast(connectDeadlineNanos)) { + // The query this reconnect serves has run out of time. The walk + // resumes on the next execute() (see reconnectPending). + throw new HttpClientException("query timeout expired during the failover reconnect [lastError=" + + (lastError == null ? "" : lastError.getMessage()) + ']'); + } int i = hostTracker.pickNext(); if (i < 0) { if (!retriedAfterReset) { @@ -1918,6 +2404,7 @@ private void reconnectViaTracker() { hostTracker.recordSuccess(i); currentEndpointIndex = i; connected = true; + reconnectPending = false; return; } if (lastMismatch != null) { @@ -1930,6 +2417,23 @@ private void reconnectViaTracker() { + ", lastError=" + (lastError == null ? "" : lastError.getMessage()) + ']'); } + /** + * {@link #reconnectViaTracker()} on behalf of a query. When the query has a + * timeout ({@code timed}), every endpoint step -- TCP connect, upgrade, + * {@code SERVER_INFO} wait -- is bounded by the query's remaining budget and + * the walk stops once {@code deadlineNanos} passes, so a black-holed endpoint + * cannot hold the caller past the timeout. + */ + private void reconnectViaTracker(boolean timed, long deadlineNanos) { + connectDeadlineActive = timed; + connectDeadlineNanos = deadlineNanos; + try { + reconnectViaTracker(); + } finally { + connectDeadlineActive = false; + } + } + private String resolveAuthorizationHeader() { // With a token provider, query it once per connect()/reconnect (the caller resolves before the // endpoint walk) so a reconnect presents a freshly refreshed token; validateToken rejects a @@ -1958,14 +2462,22 @@ private String resolveAuthorizationHeader() { return authorizationHeader; } - private long resolveQueryFlags(boolean resetSymbolDict) { - if (!resetSymbolDict) { + private long resolveQueryFlags(boolean resetSymbolDict, boolean timed) { + if (!resetSymbolDict && !timed) { return 0L; } QwpServerInfo info = serverInfo; - return info != null && (info.getCapabilities() & QwpEgressMsgKind.CAP_QUERY_FLAGS) != 0 - ? QwpEgressMsgKind.QUERY_FLAG_RESET_DICT - : 0L; + if (info == null || (info.getCapabilities() & QwpEgressMsgKind.CAP_QUERY_FLAGS) == 0) { + return 0L; + } + long flags = 0L; + if (resetSymbolDict) { + flags |= QwpEgressMsgKind.QUERY_FLAG_RESET_DICT; + } + if (timed && (info.getCapabilities() & QwpEgressMsgKind.CAP_QUERY_TIMEOUT) != 0) { + flags |= QwpEgressMsgKind.QUERY_FLAG_TIMEOUT; + } + return flags; } private void runUpgradeWithTimeout(Endpoint ep, String authHeader) { @@ -1977,7 +2489,7 @@ private void runUpgradeWithTimeout(Endpoint ep, String authHeader) { // as a transport error and moves on to the next endpoint. webSocketClient.connect(ep.host, ep.port); - int timeoutMs = (int) Math.min(authTimeoutMs, Integer.MAX_VALUE); + int timeoutMs = boundToConnectDeadline(authTimeoutMs); try { webSocketClient.upgrade(DEFAULT_ENDPOINT_PATH, timeoutMs, authHeader); } catch (HttpClientException ex) { diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpSpscQueue.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpSpscQueue.java index c5c292d32..f8fb6c3a7 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpSpscQueue.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpSpscQueue.java @@ -146,4 +146,47 @@ public T take() throws InterruptedException { consumerThread = null; } } + + /** + * Spin-then-park take bounded by an absolute {@link System#nanoTime()} + * deadline. Returns the next value, or {@code null} once the deadline has + * passed with the ring still empty. A deadline already in the past + * degenerates to a single {@link #poll()}: a value that is already + * available is still returned. Allocation-free, like {@link #take()}. + * + * @throws InterruptedException when the consumer thread was interrupted + * while waiting + */ + public T take(long deadlineNanos) throws InterruptedException { + T value = poll(); + if (value != null) { + return value; + } + if (deadlineNanos - System.nanoTime() <= 0) { + return null; + } + for (int i = 0; i < SPIN_ITERATIONS; i++) { + Compat.onSpinWait(); + if ((value = poll()) != null) { + return value; + } + } + // Same publish-then-re-poll handshake as take(); see the comment there. + consumerThread = Thread.currentThread(); + try { + while ((value = poll()) == null) { + if (Thread.interrupted()) { + throw new InterruptedException(); + } + final long remainingNanos = deadlineNanos - System.nanoTime(); + if (remainingNanos <= 0) { + return null; + } + LockSupport.parkNanos(remainingNanos); + } + return value; + } finally { + consumerThread = null; + } + } } diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/protocol/QwpConstants.java b/core/src/main/java/io/questdb/client/cutlass/qwp/protocol/QwpConstants.java index 074cf0e3f..31f7c0aba 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/protocol/QwpConstants.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/protocol/QwpConstants.java @@ -112,6 +112,16 @@ public final class QwpConstants { * the ingress {@code STATUS_*} namespace (0x00-0x09). */ public static final byte STATUS_LIMIT_EXCEEDED = 0x0B; + /** + * Status byte on a {@code QUERY_ERROR} frame: the query ran past its + * per-query timeout. Sent by a server that enforces the timeout carried on + * {@code QUERY_REQUEST} (see {@code QwpEgressMsgKind#CAP_QUERY_TIMEOUT}), and + * synthesized by the client when its own deadline for the query expires. + * The connection stays usable for the next query. Egress extension of the + * ingress {@code STATUS_*} namespace; the next free value after + * {@code STATUS_DICTIONARY_GAP} (0x0D). + */ + public static final byte STATUS_QUERY_TIMEOUT = 0x0E; /** * Column type: BINARY (length-prefixed opaque bytes). * Wire format: identical to VARCHAR — (N+1) x uint32 offsets + concatenated bytes. diff --git a/core/src/main/java/io/questdb/client/impl/ConfigSchema.java b/core/src/main/java/io/questdb/client/impl/ConfigSchema.java index c9529a13e..75c0821bc 100644 --- a/core/src/main/java/io/questdb/client/impl/ConfigSchema.java +++ b/core/src/main/java/io/questdb/client/impl/ConfigSchema.java @@ -102,6 +102,7 @@ public final class ConfigSchema { longRange("failover_max_duration_ms", Side.EGRESS, 0, OPEN_MAX, false, false); // >= 0 intRange("max_batch_rows", Side.EGRESS, 1, 1_048_576, false, false); // [1, 1048576] longRange("initial_credit", Side.EGRESS, 0, OPEN_MAX, false, false); // >= 0 + longRange("query_timeout_ms", Side.EGRESS, 0, OPEN_MAX, false, false); // >= 0; 0 = no timeout intRange("buffer_pool_size", Side.EGRESS, 1, OPEN_MAX, false, false); // >= 1 enumKey("compression", Side.EGRESS, "zstd", "raw", "auto"); intRange("compression_level", Side.EGRESS, 1, 22, false, false); // [1, 22] diff --git a/core/src/main/java/io/questdb/client/impl/QueryImpl.java b/core/src/main/java/io/questdb/client/impl/QueryImpl.java index 82088fa8b..3f5aa2e4d 100644 --- a/core/src/main/java/io/questdb/client/impl/QueryImpl.java +++ b/core/src/main/java/io/questdb/client/impl/QueryImpl.java @@ -72,9 +72,22 @@ final class QueryImpl { private final QueryWorker worker; private final QwpBindSetter wireBinds = this::applyBinds; private final WrappingHandler wrappingHandler = new WrappingHandler(); + // Set when this lease cancels its current submission; cleared by submit(). + // applyBinds() re-applies it from inside execute(), so a cancel issued while + // the submission still waits for the worker is not lost. See requestCancel(). + private volatile boolean cancelRequested; private volatile boolean done = true; private volatile String resultMessage; private volatile byte resultStatus; + // Stamped by submit() for the worker thread: when the query was submitted and + // the timeout it runs under (0 = none), so runOn() can hand the client only + // the part of the budget that is left. Published by worker.dispatch(). + private long submitNanos; + private long submitTimeoutMillis; + // Query timeout set on this handle, in milliseconds: -1 = the pooled client's + // default (query_timeout_ms), 0 = none. Kept across submits like the other + // builder state; resetForBorrow() restores -1. + private long timeoutMillis = -1; private volatile Throwable unexpectedError; private QwpBindSetter userBinds; private QwpColumnBatchHandler userHandler; @@ -168,13 +181,31 @@ void close(long gen) { // is now uncertain -- a late RESULT_* for the abandoned query could // corrupt the next borrower's stream -- so it is discarded rather than // returned. The pool grows a fresh worker on the next borrow. + final long budgetMillis = worker.closeQueryTimeoutMillis(); + final long startNanos = System.nanoTime(); if (!done) { worker.cancelInFlight(gen); - if (!awaitDone(worker.closeQueryTimeoutMillis())) { + if (!awaitDone(budgetMillis)) { worker.discardFromPool(gen); return; } } + // done means the caller has its outcome, not that the worker is idle: + // after a query timeout the worker may still be draining the aborted + // query so that its connection stays reusable. It must not reach another + // borrower before that finishes, so wait for it within what is left of + // the same budget; one that does not finish is discarded as above. + final long remainingMillis = budgetMillis - (System.nanoTime() - startNanos) / 1_000_000L; + if (!worker.awaitIdle(remainingMillis)) { + worker.discardFromPool(gen); + return; + } + // A connection the client gave up on -- it stopped responding after a + // query timeout, or the server closed it -- must not go back to the pool. + if (worker.client().hasTerminalFailure()) { + worker.discardFromPool(gen); + return; + } worker.releaseToPool(gen); } @@ -213,6 +244,29 @@ void setSql(long gen, CharSequence sql) { sqlBuffer.put(sql); } + /** + * Records that the current submission was cancelled. Called by + * {@link QueryWorker#cancelInFlight()} under the pool lock, after the lease + * generation was validated, so a stale handle cannot mark a later borrower's + * query. Needed because {@code execute()} starts by clearing the client's + * cancel latch: a cancel that lands while the submission still waits for the + * worker -- for instance behind a timed-out query the worker is draining -- + * would otherwise be dropped. + */ + void requestCancel() { + cancelRequested = true; + } + + void setTimeout(long gen, long timeout, TimeUnit unit) { + checkLive(gen); + if (timeout < 0) { + throw new IllegalArgumentException("timeout must be >= 0"); + } + long millis = unit.toMillis(timeout); + // A positive timeout below one millisecond must not round down to "none". + this.timeoutMillis = millis == 0 && timeout > 0 ? 1 : millis; + } + void submit(long gen) { checkLive(gen); if (sqlBuffer.length() == 0) { @@ -224,6 +278,11 @@ void submit(long gen) { if (!done) { throw new IllegalStateException("a previous submit() is still in flight; await the Completion first"); } + // The timeout runs from here, so a dispatch that has to wait for the + // worker (still draining a previously timed-out query) counts against it. + submitTimeoutMillis = timeoutMillis >= 0 ? timeoutMillis : worker.client().getQueryTimeoutMs(); + submitNanos = System.nanoTime(); + cancelRequested = false; // Reset terminal state under the lock so a stale signal from a prior // run can't be observed by the upcoming await(). doneLock.lock(); @@ -239,6 +298,13 @@ void submit(long gen) { } private void applyBinds(QwpBindValues binds) { + // Runs inside execute() after it cleared the client's cancel latch for + // this query, and before the request is sent: set the latch again if the + // lease cancelled the submission meanwhile, so the query goes out with + // its CANCEL right behind it. + if (cancelRequested) { + worker.client().cancel(); + } QwpBindSetter setter = userBinds; if (setter != null) { setter.apply(binds); @@ -333,6 +399,8 @@ void resetForBorrow() { userBinds = null; userHandler = null; sqlBuffer.clear(); + timeoutMillis = -1; + cancelRequested = false; resultStatus = 0; resultMessage = null; unexpectedError = null; @@ -340,12 +408,23 @@ void resetForBorrow() { } void runOn(QwpQueryClient client) { + long timeoutArg = 0; + if (submitTimeoutMillis > 0) { + // Only what is left of the budget since submit(); at least 1ms so the + // query still runs under a timeout rather than without one. + long elapsedMillis = (System.nanoTime() - submitNanos) / 1_000_000L; + timeoutArg = Math.max(1L, submitTimeoutMillis - elapsedMillis); + } + // The grace period after an expired timeout is query_close_timeout_ms, + // read per query: the pool learns it only after prewarming its clients. + client.withQueryTimeoutGrace(worker.closeQueryTimeoutMillis()); // Pass the StringSink directly as a CharSequence -- the wire encoder // reads chars and writes UTF-8 bytes straight into the send buffer. - // sqlBuffer is stable for the duration of execute(): the calling - // worker thread is blocked here until a terminal event arrives, and - // sql(...) cannot be invoked again until done==true. - client.execute(sqlBuffer, wireBinds, wrappingHandler); + // sqlBuffer must stay stable until the request is encoded: the caller + // can call sql(...) again only once done==true, which execute() signals + // either at the terminal event or, after a query timeout, once the I/O + // thread has encoded the request (see QwpQueryClient.executeOnce). + client.execute(sqlBuffer, wireBinds, wrappingHandler, false, timeoutArg); } /** diff --git a/core/src/main/java/io/questdb/client/impl/QueryLease.java b/core/src/main/java/io/questdb/client/impl/QueryLease.java index 4b07e9654..267d848a9 100644 --- a/core/src/main/java/io/questdb/client/impl/QueryLease.java +++ b/core/src/main/java/io/questdb/client/impl/QueryLease.java @@ -113,4 +113,10 @@ public Completion submit() { public QwpServerInfo serverInfo() { return impl.serverInfo(generation); } + + @Override + public Query timeout(long timeout, TimeUnit unit) { + impl.setTimeout(generation, timeout, unit); + return this; + } } diff --git a/core/src/main/java/io/questdb/client/impl/QueryWorker.java b/core/src/main/java/io/questdb/client/impl/QueryWorker.java index 568ac0dfe..2e9966a34 100644 --- a/core/src/main/java/io/questdb/client/impl/QueryWorker.java +++ b/core/src/main/java/io/questdb/client/impl/QueryWorker.java @@ -28,6 +28,7 @@ import io.questdb.client.QueryException; import io.questdb.client.cutlass.qwp.client.QwpQueryClient; +import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; @@ -51,6 +52,9 @@ public final class QueryWorker { static final long SHUTDOWN_JOIN_MILLIS = 5_000; private final QwpQueryClient client; private final long createdAtMillis; + // Signalled, under signalLock, whenever the worker becomes idle: a job's + // runOn() returned, or shutdown stranded the pending job. See awaitIdle(). + private final Condition idleCondition; private final QueryClientPool pool; private final QueryImpl query; private final Condition signalCondition; @@ -78,6 +82,10 @@ public final class QueryWorker { // thread observes the latest value without taking the pool lock. private volatile long generation; private volatile long idleSinceMillis; + // True while the dispatch thread runs a job's runOn(). Guarded by signalLock. + // Outlives the job's terminal callback: after a query timeout the caller is + // released while runOn() still drains the aborted query. + private boolean running; private volatile boolean shuttingDown; public QueryWorker(QwpQueryClient client, QueryClientPool pool, int slotIndex) { @@ -85,12 +93,43 @@ public QueryWorker(QwpQueryClient client, QueryClientPool pool, int slotIndex) { this.pool = pool; this.query = new QueryImpl(this); this.signalCondition = signalLock.newCondition(); + this.idleCondition = signalLock.newCondition(); this.thread = new Thread(this::runLoop, "questdb-query-worker-" + slotIndex); this.thread.setDaemon(true); this.createdAtMillis = System.currentTimeMillis(); this.idleSinceMillis = this.createdAtMillis; } + /** + * Waits up to {@code timeoutMillis} for this worker to become idle: no job + * running and none pending. Used by {@link QueryImpl#close(long)} before + * returning the worker to the pool, because a job's terminal callback can + * precede the end of its {@code runOn()} -- after a query timeout the + * client keeps draining the aborted query so its connection stays + * reusable. Returns {@code false} on timeout or interrupt; an interrupt + * re-raises the caller's flag, like {@link QueryImpl}'s close drain. + */ + boolean awaitIdle(long timeoutMillis) { + long remainingNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0L, timeoutMillis)); + signalLock.lock(); + try { + while (running || current != null) { + if (remainingNanos <= 0L) { + return false; + } + try { + remainingNanos = idleCondition.awaitNanos(remainingNanos); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + } + return true; + } finally { + signalLock.unlock(); + } + } + long createdAtMillis() { return createdAtMillis; } @@ -139,6 +178,9 @@ void markIdleAt(long nowMillis) { * generation first. Lease code must use {@link #cancelInFlight(long)}. */ void cancelInFlight() { + // Remembered on the lease too: a submission still waiting for this worker + // (behind a timed-out query being drained) has not reached the client yet. + query.requestCancel(); try { client.cancel(); } catch (RuntimeException ignored) { @@ -317,6 +359,7 @@ private void runLoop() { // caller so its Completion.await() does not hang. QueryImpl stranded = current; current = null; + idleCondition.signalAll(); if (stranded != null) { stranded.signalUnexpected( new QueryException((byte) 0, "QuestDB handle is closed")); @@ -335,6 +378,7 @@ private void runLoop() { // already-consumed signal, and park the worker forever while // the user thread waits on a Completion that never fires. current = null; + running = true; } finally { signalLock.unlock(); } @@ -343,6 +387,13 @@ private void runLoop() { } catch (Throwable t) { q.signalUnexpected(t); } + signalLock.lock(); + try { + running = false; + idleCondition.signalAll(); + } finally { + signalLock.unlock(); + } // Test-only barrier: deterministically reproduce the busy-worker // shutdown-drop race (df6f7ca) at its exact site. Null in production. Runnable hook = busyWorkerTestHook; diff --git a/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java b/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java new file mode 100644 index 000000000..2e8913737 --- /dev/null +++ b/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java @@ -0,0 +1,391 @@ +/*+***************************************************************************** + * ___ _ ____ ____ + * / _ \ _ _ ___ ___| |_| _ \| __ ) + * | | | | | | |/ _ \/ __| __| | | | _ \ + * | |_| | |_| | __/\__ \ |_| |_| | |_) | + * \__\_\\__,_|\___||___/\__|____/|____/ + * + * Copyright (c) 2014-2019 Appsicle + * Copyright (c) 2019-2026 QuestDB + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + ******************************************************************************/ + + +package io.questdb.client.test; + +import io.questdb.client.Completion; +import io.questdb.client.Query; +import io.questdb.client.QueryException; +import io.questdb.client.QuestDB; +import io.questdb.client.cutlass.qwp.client.QwpColumnBatch; +import io.questdb.client.cutlass.qwp.client.QwpColumnBatchHandler; +import io.questdb.client.cutlass.qwp.client.QwpEgressMsgKind; +import io.questdb.client.cutlass.qwp.protocol.QwpConstants; +import io.questdb.client.test.cutlass.qwp.websocket.TestWebSocketServer; +import io.questdb.client.test.tools.TestUtils; +import org.junit.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * {@link Query#timeout} through the {@link QuestDB} facade: a timed-out query is + * reported as a {@link QueryException} whose {@link QueryException#isTimeout()} + * holds, while the pooled, authenticated connection stays open and serves the + * following queries. A worker still draining a timed-out query is not handed to + * the next borrower before it is idle, and one whose connection stopped + * responding is replaced rather than reused. + */ +public class QueryTimeoutFacadeTest { + + @Test(timeout = 30_000) + public void testCancelOfASubmissionQueuedBehindADrainIsNotLost() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 1_000; + Script script = new Script(); + script.onQuery = (s, c, id, n) -> { + if (n == 1) { + // Past the grace period: the caller is released while the worker + // is still draining this query. + s.replyOnceLater(c, id, timeoutMs + graceMs + 600, + queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, "timeout, query aborted")); + } + // n == 2: held until cancelled + }; + script.onCancel = (s, c, id) -> s.replyOnce(c, id, + queryError(id, QwpConstants.STATUS_CANCELLED, "cancelled by client")); + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT); + QuestDB db = QuestDB.connect(config(server) + "query_close_timeout_ms=" + graceMs + ";"); + Query q = db.borrowQuery()) { + q.sql("SELECT slow()").handler(new NoopHandler()).timeout(timeoutMs, TimeUnit.MILLISECONDS); + assertTimesOut(q); + + // The worker is still draining the first query, so this submission + // waits for it -- and is cancelled while it waits. + Completion c = q.timeout(5, TimeUnit.SECONDS).submit(); + c.cancel(); + try { + c.await(); + Assert.fail("the cancelled query must not complete"); + } catch (QueryException e) { + Assert.assertEquals("the cancel must reach the queued query: " + e.getMessage(), + QwpConstants.STATUS_CANCELLED, e.getStatus()); + } + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testConfiguredDefaultAppliesAndHandleCanOverrideIt() throws Exception { + TestUtils.assertMemoryLeak(() -> { + Script script = new Script(); + // An older server: the client cancels at the deadline. Query 2 is the + // one run without a timeout, so it is allowed to take its time. + script.onQuery = (s, c, id, n) -> { + if (n == 2) { + s.sendLater(c, 300, resultEnd(id)); + } + }; + script.onCancel = (s, c, id) -> s.replyOnce(c, id, + queryError(id, QwpConstants.STATUS_CANCELLED, "cancelled by client")); + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS); + QuestDB db = QuestDB.connect(config(server) + "query_timeout_ms=100;")) { + try (Query q = db.borrowQuery()) { + q.sql("SELECT slow()").handler(new NoopHandler()); + assertTimesOut(q); + + q.timeout(0, TimeUnit.MILLISECONDS); + q.submit().await(); // no timeout: completes after 300ms + } + // A fresh borrow starts from the configured default again. + try (Query q = db.borrowQuery()) { + q.sql("SELECT slow()").handler(new NoopHandler()); + assertTimesOut(q); + } + Assert.assertEquals("all queries must share one connection", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testTimedOutQueryKeepsThePooledConnection() throws Exception { + TestUtils.assertMemoryLeak(() -> { + Script script = new Script(); + script.onQuery = (s, c, id, n) -> { + if (n == 1) { + s.sendLater(c, 150, queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, + "timeout, query aborted [runtime=150ms, timeout=100ms]")); + } else { + s.send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT); + QuestDB db = QuestDB.connect(config(server))) { + try (Query q = db.borrowQuery()) { + q.sql("SELECT slow()").handler(new NoopHandler()).timeout(100, TimeUnit.MILLISECONDS); + QueryException e = assertTimesOut(q); + Assert.assertTrue("the server's report must reach the caller: " + e.getMessage(), + e.getMessage().startsWith("timeout, query aborted")); + + // The same handle, and so the same connection, runs the next query. + q.sql("INSERT INTO t VALUES (1)").submit().await(); + } + try (Query q = db.borrowQuery()) { + q.sql("INSERT INTO t VALUES (2)").handler(new NoopHandler()).submit().await(); + } + Assert.assertEquals("the authenticated connection must be kept across the timeout", + 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testWorkerDrainingATimedOutQueryIsReturnedOnlyWhenIdle() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 1_000; + Script script = new Script(); + script.onQuery = (s, c, id, n) -> { + if (n == 1) { + // Past the grace period but within the second one: the caller is + // released first, the connection drains the reply afterwards. + s.sendLater(c, timeoutMs + graceMs + 400, + queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, "timeout, query aborted")); + } else { + s.send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT); + QuestDB db = QuestDB.connect(config(server) + "query_close_timeout_ms=" + graceMs + ";")) { + Query q = db.borrowQuery(); + q.sql("SELECT slow()").handler(new NoopHandler()).timeout(timeoutMs, TimeUnit.MILLISECONDS); + long start = System.nanoTime(); + assertTimesOut(q); + long awaitMs = elapsedMs(start); + Assert.assertTrue("the caller is released at the end of the grace period, after " + awaitMs + "ms", + awaitMs >= timeoutMs + graceMs); + Assert.assertEquals("the caller must not wait for the server's late reply", 0, script.sends.get()); + + q.close(); + Assert.assertEquals("close() must wait for the drain before returning the worker", + 1, script.sends.get()); + + try (Query next = db.borrowQuery()) { + next.sql("INSERT INTO t VALUES (1)").handler(new NoopHandler()).submit().await(); + } + Assert.assertEquals("the drained connection must be reused", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testWorkerWhoseConnectionStoppedRespondingIsReplaced() throws Exception { + TestUtils.assertMemoryLeak(() -> { + Script script = new Script(); + script.onQuery = (s, c, id, n) -> { + if (n > 1) { + s.send(c, execDone(id)); + } + // n == 1: never answered, not even after a CANCEL. + }; + // failover=off: the client cannot replace the connection by itself, so + // the next query succeeds only if the pool discarded the worker. + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT); + QuestDB db = QuestDB.connect(config(server) + "query_close_timeout_ms=200;failover=off;")) { + try (Query q = db.borrowQuery()) { + q.sql("SELECT slow()").handler(new NoopHandler()).timeout(100, TimeUnit.MILLISECONDS); + assertTimesOut(q); + } + try (Query q = db.borrowQuery()) { + q.sql("INSERT INTO t VALUES (1)").handler(new NoopHandler()).submit().await(); + } + Assert.assertEquals("the silent connection must be replaced by a new one", 2, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + private static QueryException assertTimesOut(Query q) throws InterruptedException { + try { + q.submit().await(); + } catch (QueryException e) { + Assert.assertTrue("expected a query timeout, got status=" + e.getStatus() + ": " + e.getMessage(), + e.isTimeout()); + return e; + } + Assert.fail("the query must time out"); + return null; + } + + private static String config(TestWebSocketServer server) { + return "ws::addr=localhost:" + server.getPort() + + ";auth_timeout_ms=2000;sender_pool_min=0;query_pool_min=1;query_pool_max=1;"; + } + + private static long elapsedMs(long startNanos) { + return TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos); + } + + private static byte[] execDone(long requestId) { + return frame(QwpEgressMsgKind.EXEC_DONE, requestId, new byte[]{0, 0}); // op_type, rows_affected + } + + private static byte[] frame(byte msgKind, long requestId, byte[] body) { + int payloadLen = 1 + 8 + body.length; + ByteBuffer bb = ByteBuffer.allocate(QwpConstants.HEADER_SIZE + payloadLen).order(ByteOrder.LITTLE_ENDIAN); + bb.putInt(QwpConstants.MAGIC_MESSAGE); + bb.put(QwpConstants.VERSION); + bb.put((byte) 0); // flags + bb.putShort((short) 0); // table_count + bb.putInt(payloadLen); + bb.put(msgKind); + bb.putLong(requestId); + bb.put(body); + return bb.array(); + } + + private static byte[] queryError(long requestId, byte status, String message) { + byte[] msg = message.getBytes(StandardCharsets.UTF_8); + ByteBuffer body = ByteBuffer.allocate(1 + 2 + msg.length).order(ByteOrder.LITTLE_ENDIAN); + body.put(status).putShort((short) msg.length).put(msg); + return frame(QwpEgressMsgKind.QUERY_ERROR, requestId, body.array()); + } + + private static byte[] resultEnd(long requestId) { + return frame(QwpEgressMsgKind.RESULT_END, requestId, new byte[]{0, 0}); // final_seq, total_rows + } + + private static TestWebSocketServer startServer(Script script, int capabilities) throws Exception { + TestWebSocketServer server = new TestWebSocketServer(script); + server.setSendServerInfo(true); + server.setCapabilities(capabilities); + server.start(); + Assert.assertTrue(server.awaitStart(5, TimeUnit.SECONDS)); + return server; + } + + @FunctionalInterface + private interface OnCancel { + void run(Script script, TestWebSocketServer.ClientHandler client, long requestId); + } + + @FunctionalInterface + private interface OnQuery { + void run(Script script, TestWebSocketServer.ClientHandler client, long requestId, int queryNumber); + } + + private static final class NoopHandler implements QwpColumnBatchHandler { + @Override + public void onBatch(QwpColumnBatch batch) { + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + } + } + + /** + * Scripted egress server. Like a real one it sends one terminal frame per + * request, so a cancel of a query that already ended gets no reply. + */ + private static final class Script implements TestWebSocketServer.WebSocketServerHandler, AutoCloseable { + // Delayed sends that have started, for ordering assertions. Counted before + // the frame goes out, so a client that has received it sees the count. + final AtomicInteger sends = new AtomicInteger(); + private final Set answered = ConcurrentHashMap.newKeySet(); + private final AtomicInteger queryCount = new AtomicInteger(); + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(r -> { + Thread t = new Thread(r, "scripted-egress-server"); + t.setDaemon(true); + return t; + }); + volatile OnCancel onCancel; + volatile OnQuery onQuery; + + @Override + public void close() { + scheduler.shutdownNow(); + } + + @Override + public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] data) { + if (data.length < 9) { + return; + } + long requestId = ByteBuffer.wrap(data, 1, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(); + if (data[0] == QwpEgressMsgKind.QUERY_REQUEST) { + OnQuery script = onQuery; + if (script != null) { + script.run(this, client, requestId, queryCount.incrementAndGet()); + } + } else if (data[0] == QwpEgressMsgKind.CANCEL) { + OnCancel script = onCancel; + if (script != null) { + script.run(this, client, requestId); + } + } + } + + void replyOnce(TestWebSocketServer.ClientHandler client, long requestId, byte[] frame) { + if (answered.add(requestId)) { + send(client, frame); + } + } + + void replyOnceLater(TestWebSocketServer.ClientHandler client, long requestId, long delayMs, byte[] frame) { + if (answered.add(requestId)) { + sendLater(client, delayMs, frame); + } + } + + void send(TestWebSocketServer.ClientHandler client, byte[] frame) { + try { + client.sendBinary(frame); + } catch (IOException ignore) { + // the client went away; the test's assertions report the outcome + } + } + + void sendLater(TestWebSocketServer.ClientHandler client, long delayMs, byte[] frame) { + scheduler.schedule(() -> { + sends.incrementAndGet(); + send(client, frame); + }, delayMs, TimeUnit.MILLISECONDS); + } + } +} diff --git a/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java new file mode 100644 index 000000000..9240a1fc9 --- /dev/null +++ b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java @@ -0,0 +1,845 @@ +/*+***************************************************************************** + * ___ _ ____ ____ + * / _ \ _ _ ___ ___| |_| _ \| __ ) + * | | | | | | |/ _ \/ __| __| | | | _ \ + * | |_| | |_| | __/\__ \ |_| |_| | |_) | + * \__\_\\__,_|\___||___/\__|____/|____/ + * + * Copyright (c) 2014-2019 Appsicle + * Copyright (c) 2019-2026 QuestDB + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + ******************************************************************************/ + + +package io.questdb.client.test.cutlass.qwp.client; + +import io.questdb.client.cutlass.qwp.client.QwpColumnBatch; +import io.questdb.client.cutlass.qwp.client.QwpColumnBatchHandler; +import io.questdb.client.cutlass.qwp.client.QwpEgressMsgKind; +import io.questdb.client.cutlass.qwp.client.QwpQueryClient; +import io.questdb.client.cutlass.qwp.client.QwpServerInfo; +import io.questdb.client.cutlass.qwp.client.WebSocketResponse; +import io.questdb.client.cutlass.qwp.protocol.QwpConstants; +import io.questdb.client.test.cutlass.qwp.websocket.TestWebSocketServer; +import io.questdb.client.test.tools.TestUtils; +import org.junit.Assert; +import org.junit.Test; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.net.SocketTimeoutException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Per-query timeout of {@link QwpQueryClient}, against a scripted mock server. + * Pins the wire contract (the capability-gated {@code timeout_ms} field after + * {@code query_flags}), the client-side deadline and its grace periods, and the + * central promise of the feature: a timed-out query ends gracefully -- reported + * as {@link QwpConstants#STATUS_QUERY_TIMEOUT} -- on a connection that stays + * open and authenticated for the next query. Only a connection that does not + * answer at all is replaced. + *

    + * The server-side enforcement (breaker timeout, status mapping) is covered + * against a live server in the questdb repository. + */ +public class QwpQueryClientQueryTimeoutTest { + + private static final int CAPS_FLAGS_ONLY = QwpEgressMsgKind.CAP_QUERY_FLAGS; + private static final int CAPS_WITH_TIMEOUT = QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT; + + @Test(timeout = 30_000) + public void testBatchesArrivingAfterTheDeadlineAreDiscarded() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + send(c, firstBatch(id, "a")); + s.sendLater(c, 400, nextBatch(id, 1, "b"), nextBatch(id, 2, "c"), + queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, "timeout, query aborted")); + } else { + send(c, firstBatch(id, "z")); + send(c, resultEnd(id, 1)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + RecordingHandler h = new RecordingHandler(); + client.execute("SELECT s FROM t", null, h, false, 150); + Assert.assertEquals("only the batch that arrived before the deadline reaches the handler", + 1, h.batches.get()); + h.assertTimedOut(); + + // The discarded batches were still decoded, so the connection-scoped + // state stays in step with the server and the next query decodes. + RecordingHandler next = new RecordingHandler(); + client.execute("SELECT s FROM t", null, next, false, 0); + Assert.assertTrue("the next query must complete on the same connection", next.ended); + Assert.assertEquals(1, next.batches.get()); + Assert.assertEquals("the connection must be reused, not replaced", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testCallerIsReleasedAfterTheGraceWhileTheConnectionDrains() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 1_000; + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + // Ends the query well after the grace period, but well within the + // second one: the caller must not wait for it, the connection must. + s.sendLater(c, timeoutMs + graceMs + 400, + queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, "timeout, query aborted")); + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.withQueryTimeoutGrace(graceMs); + RecordingHandler h = new RecordingHandler(); + long start = System.nanoTime(); + client.execute("SELECT slow()", null, h, false, timeoutMs); + long executeMs = elapsedMs(start); + + h.assertTimedOut(); + Assert.assertTrue("the message must say the server did not end the query in time: " + h.errorMessage, + h.errorMessage.contains("grace period")); + long releasedMs = TimeUnit.NANOSECONDS.toMillis(h.terminalNanos - start); + Assert.assertTrue("the caller must be released only once the grace period is over, after " + + releasedMs + "ms", releasedMs >= timeoutMs + graceMs); + Assert.assertTrue("the caller must be released before the server's late reply was even sent", + h.terminalNanos < script.lastSendNanos); + Assert.assertTrue("execute() must return only after draining the server's reply, after " + + executeMs + "ms", executeMs >= timeoutMs + graceMs + 400); + Assert.assertEquals("one terminal callback per query", 1, h.terminals.get()); + Assert.assertNotNull("the end of the grace period must nudge the server with a CANCEL", + script.cancels.poll(5, TimeUnit.SECONDS)); + Assert.assertFalse("a drained connection is healthy", client.hasTerminalFailure()); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue(next.execDone); + Assert.assertEquals("the drained connection must be reused", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testCompleteResultJustPastTheDeadlineIsASuccess() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + // No result batch is withheld from the handler, so a reply that lands + // after the client's deadline still reports the real outcome. + s.sendLater(c, 300, n == 1 ? resultEnd(id, 0) : execDone(id)); + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + RecordingHandler select = new RecordingHandler(); + client.execute("SELECT 1 WHERE false", null, select, false, 100); + Assert.assertTrue("an empty result past the deadline is still complete", select.ended); + Assert.assertEquals(0, select.errorStatus); + + RecordingHandler insert = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, insert, false, 100); + Assert.assertTrue("a statement that completed past the deadline took effect", insert.execDone); + Assert.assertEquals(0, insert.errorStatus); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testFailoverReplayCarriesTheRemainingBudget() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + if (n == 1) { + // A transport failure mid-query: the client fails over and replays. + s.closeLater(c, 200); + } else { + send(c, execDone(requestIdOf(frame))); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, + "failover_backoff_initial_ms=0;failover_backoff_max_ms=0;")) { + RecordingHandler h = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, h, false, 10_000); + Assert.assertTrue("the replay must complete", h.execDone); + Assert.assertEquals(1, h.failoverResets.get()); + + long firstTimeout = trailerOf(script.queries.take())[1]; + long replayTimeout = trailerOf(script.queries.take())[1]; + Assert.assertTrue("the first attempt carries about the whole budget: " + firstTimeout, + firstTimeout > 9_000 && firstTimeout <= 10_000); + Assert.assertTrue("the replay must carry only the remaining budget, not a fresh one: " + + replayTimeout, replayTimeout > 0 && replayTimeout <= firstTimeout - 150); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testNoTimeoutFieldWithoutTheCapabilityOrTimeout() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> send(c, execDone(requestIdOf(frame))); + try (TestWebSocketServer server = startServer(script, CAPS_FLAGS_ONLY); + QwpQueryClient client = connect(server, "")) { + client.execute("SELECT 1", null, new RecordingHandler(), false, 30_000); + Assert.assertNull("a server without CAP_QUERY_TIMEOUT must get no timeout field", + trailerOf(script.queries.take())); + } finally { + script.close(); + } + + ScriptedServer script2 = new ScriptedServer(); + script2.onQuery = (s, c, frame, n) -> send(c, execDone(requestIdOf(frame))); + try (TestWebSocketServer server = startServer(script2, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.execute("SELECT 1", null, new RecordingHandler(), false, 0); + Assert.assertNull("a query without a timeout must send no timeout field", + trailerOf(script2.queries.take())); + } finally { + script2.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testOlderServerIsCancelledAtTheDeadlineAndTheConnectionIsKept() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + if (n > 1) { + send(c, execDone(requestIdOf(frame))); + } + // n == 1: hold the query; only a CANCEL ends it. + }; + script.onCancel = (s, c, frame, n) -> s.replyOnce(c, requestIdOf(frame), 0, + queryError(requestIdOf(frame), QwpConstants.STATUS_CANCELLED, "cancelled by client")); + try (TestWebSocketServer server = startServer(script, CAPS_FLAGS_ONLY); + QwpQueryClient client = connect(server, "")) { + RecordingHandler h = new RecordingHandler(); + long start = System.nanoTime(); + client.execute("SELECT slow()", null, h, false, 150); + long elapsed = elapsedMs(start); + + h.assertTimedOut(); + Assert.assertTrue("the timeout cannot fire early, fired after " + elapsed + "ms", elapsed >= 150); + Assert.assertTrue("the cancelled query must end promptly, took " + elapsed + "ms", + elapsed < QwpQueryClient.DEFAULT_QUERY_TIMEOUT_GRACE_MS); + byte[] cancel = script.cancels.poll(5, TimeUnit.SECONDS); + Assert.assertNotNull("an older server must be sent a CANCEL at the deadline", cancel); + Assert.assertEquals("the CANCEL must target the timed-out query", + requestIdOf(script.queries.take()), requestIdOf(cancel)); + Assert.assertFalse(client.hasTerminalFailure()); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue(next.execDone); + Assert.assertEquals("the connection must be reused, not replaced", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testReconnectWalkIsBoundedByTheQueryDeadline() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + if (n == 1) { + send(c, closeFrame()); + } else { + send(c, execDone(requestIdOf(frame))); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + BlackHole blackHole = new BlackHole()) { + // The failover prefers the never-tried black hole over the endpoint + // that just failed. It accepts TCP but never answers the upgrade, so + // without the query deadline the walk would wait out auth_timeout_ms. + try (QwpQueryClient client = QwpQueryClient.fromConfig("ws::addr=localhost:" + server.getPort() + + ",localhost:" + blackHole.port() + ";auth_timeout_ms=5000;" + + "failover_backoff_initial_ms=0;failover_backoff_max_ms=0;")) { + client.connect(); + RecordingHandler h = new RecordingHandler(); + long start = System.nanoTime(); + client.execute("INSERT INTO t VALUES (1)", null, h, false, 500); + long elapsed = elapsedMs(start); + + h.assertTimedOut(); + Assert.assertTrue("the reconnect must stop at the query deadline, took " + elapsed + "ms", + elapsed < 3_000); + Assert.assertEquals(1, blackHole.accepted()); + + // The cut-short reconnect left the client disconnected; the next + // query reconnects instead of failing with "not connected". + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (2)", null, next, false, 0); + Assert.assertTrue("the next query must reconnect and complete: " + next.errorMessage, next.execDone); + Assert.assertTrue(client.isConnected()); + } + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testServerReportedTimeoutKeepsTheConnectionOpen() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + s.sendLater(c, 150, queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, + "timeout, query aborted [runtime=150ms, timeout=100ms]")); + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + RecordingHandler h = new RecordingHandler(); + client.execute("SELECT slow()", null, h, false, 100); + + h.assertTimedOut(); + Assert.assertTrue("the server's own report must reach the handler: " + h.errorMessage, + h.errorMessage.startsWith("timeout, query aborted")); + Assert.assertTrue("the client must leave a server-enforced timeout to the server", + script.cancels.isEmpty()); + Assert.assertFalse(client.hasTerminalFailure()); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue(next.execDone); + Assert.assertEquals("the authenticated connection must be reused", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testTimeoutFieldCarriesTheBudget() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> send(c, execDone(requestIdOf(frame))); + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "query_timeout_ms=20000;")) { + Assert.assertEquals(20_000, client.getQueryTimeoutMs()); + + client.execute("SELECT 1", null, new RecordingHandler(), false, 30_000); + long[] trailer = trailerOf(script.queries.take()); + Assert.assertNotNull(trailer); + Assert.assertEquals(QwpEgressMsgKind.QUERY_FLAG_TIMEOUT, trailer[0]); + Assert.assertTrue("timeout_ms must carry the budget left: " + trailer[1], + trailer[1] > 29_000 && trailer[1] <= 30_000); + + // The overloads without a timeout apply the configured default, and + // the field composes with the other query flags. + client.execute("SELECT 1", new RecordingHandler(), true); + trailer = trailerOf(script.queries.take()); + Assert.assertNotNull(trailer); + Assert.assertEquals(QwpEgressMsgKind.QUERY_FLAG_RESET_DICT | QwpEgressMsgKind.QUERY_FLAG_TIMEOUT, + trailer[0]); + Assert.assertTrue("the configured default must apply: " + trailer[1], + trailer[1] > 19_000 && trailer[1] <= 20_000); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testTimeoutWhileFailingOverIsReportedAsTimeout() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + // Every attempt fails at the transport level. + script.onQuery = (s, c, frame, n) -> send(c, closeFrame()); + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, + "failover_max_attempts=1000000;failover_backoff_initial_ms=10;failover_backoff_max_ms=20;")) { + RecordingHandler h = new RecordingHandler(); + long start = System.nanoTime(); + client.execute("INSERT INTO t VALUES (1)", null, h, false, 300); + long elapsed = elapsedMs(start); + + h.assertTimedOut(); + Assert.assertTrue("failover must stop at the deadline, took " + elapsed + "ms", + elapsed >= 300 && elapsed < 3_000); + Assert.assertTrue("there must have been replays", script.queryCount.get() > 1); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testUnresponsiveConnectionIsReplaced() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + if (n > 1) { + send(c, execDone(requestIdOf(frame))); + } + // n == 1: never answered, not even after a CANCEL. + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.withQueryTimeoutGrace(200); + RecordingHandler h = new RecordingHandler(); + long start = System.nanoTime(); + client.execute("SELECT slow()", null, h, false, 100); + long elapsed = elapsedMs(start); + + h.assertTimedOut(); + Assert.assertEquals(1, h.terminals.get()); + Assert.assertTrue("execute() must give up after two grace periods, took " + elapsed + "ms", + elapsed >= 100 + 2 * 200 && elapsed < 5_000); + Assert.assertTrue("a silent connection must be marked failed", client.hasTerminalFailure()); + + // failover=on (the default) replaces the connection for the next query. + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue(next.execDone); + Assert.assertEquals("a replacement connection must have been opened", 2, server.handshakeCount()); + Assert.assertFalse(client.hasTerminalFailure()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testUnresponsiveConnectionWithFailoverOffReportsTheFailure() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + // never answered + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "failover=off;")) { + client.withQueryTimeoutGrace(100); + RecordingHandler h = new RecordingHandler(); + client.execute("SELECT slow()", null, h, false, 100); + h.assertTimedOut(); + Assert.assertTrue(client.hasTerminalFailure()); + + RecordingHandler next = new RecordingHandler(); + client.execute("SELECT 1", null, next, false, 0); + Assert.assertEquals(WebSocketResponse.STATUS_INTERNAL_ERROR, next.errorStatus); + Assert.assertTrue(next.errorMessage, next.errorMessage.contains("stopped responding")); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testUserCancelBeforeTheDeadlineIsReportedAsCancel() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + // held until cancelled + }; + // The cancel is acknowledged only after the deadline has passed. + script.onCancel = (s, c, frame, n) -> s.replyOnce(c, requestIdOf(frame), 400, + queryError(requestIdOf(frame), QwpConstants.STATUS_CANCELLED, "cancelled by client")); + try (TestWebSocketServer server = startServer(script, CAPS_FLAGS_ONLY); + QwpQueryClient client = connect(server, "")) { + RecordingHandler h = new RecordingHandler(); + Thread canceller = new Thread(() -> { + try { + script.queries.take(); + client.cancel(); + } catch (InterruptedException ignore) { + Thread.currentThread().interrupt(); + } + }); + canceller.start(); + client.execute("SELECT slow()", null, h, false, 200); + canceller.join(); + + Assert.assertEquals("the user's cancel came first and must be reported as such: " + h.errorMessage, + QwpConstants.STATUS_CANCELLED, h.errorStatus); + } finally { + script.close(); + } + }); + } + + private static byte[] closeFrame() { + // Marker understood by ScriptedServer's send(): close the connection. + return new byte[0]; + } + + private static QwpQueryClient connect(TestWebSocketServer server, String extraConfig) { + QwpQueryClient client = QwpQueryClient.fromConfig( + "ws::addr=localhost:" + server.getPort() + ";auth_timeout_ms=2000;" + extraConfig); + try { + client.connect(); + } catch (RuntimeException e) { + client.close(); + throw e; + } + return client; + } + + private static long elapsedMs(long startNanos) { + return TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos); + } + + private static byte[] execDone(long requestId) { + return serverFrame(QwpEgressMsgKind.EXEC_DONE, requestId, 0, new byte[]{0, 0}); // op_type, rows_affected + } + + /** + * RESULT_BATCH {@code batch_seq == 0}: the schema (one VARCHAR column) and a + * single row. + */ + private static byte[] firstBatch(long requestId, String value) { + ByteArrayOutputStream body = new ByteArrayOutputStream(); + putVarint(body, 0); // batch_seq + putVarint(body, 0); // table name length + putVarint(body, 1); // row_count + putVarint(body, 1); // column_count + putVarint(body, 1); // column name length + body.write('s'); + body.write(QwpConstants.TYPE_VARCHAR); + putVarcharCell(body, value); + return serverFrame(QwpEgressMsgKind.RESULT_BATCH, requestId, 1, body.toByteArray()); + } + + /** + * Continuation RESULT_BATCH ({@code batch_seq > 0}): rows only, against the + * schema of {@link #firstBatch}. + */ + private static byte[] nextBatch(long requestId, long batchSeq, String value) { + ByteArrayOutputStream body = new ByteArrayOutputStream(); + putVarint(body, batchSeq); + putVarint(body, 0); // table name length + putVarint(body, 1); // row_count + putVarcharCell(body, value); + return serverFrame(QwpEgressMsgKind.RESULT_BATCH, requestId, 1, body.toByteArray()); + } + + private static void putVarcharCell(ByteArrayOutputStream body, String value) { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + body.write(0); // null_flag: no nulls + ByteBuffer offsets = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN); + offsets.putInt(0).putInt(bytes.length); + body.write(offsets.array(), 0, 8); + body.write(bytes, 0, bytes.length); + } + + private static void putVarint(ByteArrayOutputStream out, long value) { + while ((value & ~0x7FL) != 0) { + out.write((int) ((value & 0x7F) | 0x80)); + value >>>= 7; + } + out.write((int) value); + } + + private static byte[] queryError(long requestId, byte status, String message) { + byte[] msg = message.getBytes(StandardCharsets.UTF_8); + ByteBuffer body = ByteBuffer.allocate(1 + 2 + msg.length).order(ByteOrder.LITTLE_ENDIAN); + body.put(status).putShort((short) msg.length).put(msg); + return serverFrame(QwpEgressMsgKind.QUERY_ERROR, requestId, 0, body.array()); + } + + private static long readVarint(byte[] buf, int[] pos) { + long value = 0; + int shift = 0; + while (true) { + byte b = buf[pos[0]++]; + value |= (long) (b & 0x7F) << shift; + if ((b & 0x80) == 0) { + return value; + } + shift += 7; + } + } + + private static long requestIdOf(byte[] clientFrame) { + return ByteBuffer.wrap(clientFrame, 1, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(); + } + + private static byte[] resultEnd(long requestId, long totalRows) { + ByteArrayOutputStream body = new ByteArrayOutputStream(); + putVarint(body, 0); // final_seq + putVarint(body, totalRows); + return serverFrame(QwpEgressMsgKind.RESULT_END, requestId, 0, body.toByteArray()); + } + + private static void send(TestWebSocketServer.ClientHandler client, byte[] frame) { + try { + if (frame.length == 0) { + client.sendClose(1011, "scripted failure"); + } else { + client.sendBinary(frame); + } + } catch (IOException ignore) { + // the client went away; the test's assertions report the outcome + } + } + + private static byte[] serverFrame(byte msgKind, long requestId, int tableCount, byte[] body) { + int payloadLen = 1 + 8 + body.length; + ByteBuffer bb = ByteBuffer.allocate(QwpConstants.HEADER_SIZE + payloadLen).order(ByteOrder.LITTLE_ENDIAN); + bb.putInt(QwpConstants.MAGIC_MESSAGE); + bb.put(QwpConstants.VERSION); + bb.put((byte) 0); // flags + bb.putShort((short) tableCount); + bb.putInt(payloadLen); + bb.put(msgKind); + bb.putLong(requestId); + bb.put(body); + return bb.array(); + } + + private static TestWebSocketServer startServer(ScriptedServer script, int capabilities) throws Exception { + TestWebSocketServer server = new TestWebSocketServer(script); + server.setSendServerInfo(true); + server.setCapabilities(capabilities); + server.start(); + Assert.assertTrue(server.awaitStart(5, TimeUnit.SECONDS)); + return server; + } + + /** + * Decodes the trailer of a captured {@code QUERY_REQUEST}: {@code {query_flags, + * timeout_ms}}, with {@code timeout_ms == -1} when the timeout flag is not + * set, or {@code null} when the frame carries no trailer at all. Assumes no + * binds. + */ + private static long[] trailerOf(byte[] f) { + Assert.assertEquals(QwpEgressMsgKind.QUERY_REQUEST, f[0]); + int[] p = {1 + 8}; + long sqlLen = readVarint(f, p); + p[0] += (int) sqlLen; + readVarint(f, p); // initial_credit + Assert.assertEquals("test frames carry no binds", 0L, readVarint(f, p)); + if (p[0] >= f.length) { + return null; + } + long flags = readVarint(f, p); + long timeoutMs = (flags & QwpEgressMsgKind.QUERY_FLAG_TIMEOUT) != 0 ? readVarint(f, p) : -1L; + Assert.assertEquals("the trailer must end the frame", f.length, p[0]); + return new long[]{flags, timeoutMs}; + } + + @FunctionalInterface + private interface Script { + void run(ScriptedServer server, TestWebSocketServer.ClientHandler client, byte[] frame, int queryNumber); + } + + /** + * A TCP listener that accepts connections and never answers them -- an + * endpoint whose WebSocket upgrade hangs. + */ + private static final class BlackHole implements AutoCloseable { + private final List accepted = Collections.synchronizedList(new ArrayList<>()); + private final ServerSocket socket; + private final Thread thread; + private volatile boolean running = true; + + BlackHole() throws IOException { + socket = new ServerSocket(0, 50, InetAddress.getLoopbackAddress()); + socket.setSoTimeout(50); + thread = new Thread(() -> { + while (running) { + try { + accepted.add(socket.accept()); + } catch (SocketTimeoutException ignore) { + // poll the running flag + } catch (IOException e) { + return; + } + } + }, "black-hole-accept"); + thread.setDaemon(true); + thread.start(); + } + + int accepted() { + return accepted.size(); + } + + @Override + public void close() throws Exception { + running = false; + thread.join(5_000); + socket.close(); + synchronized (accepted) { + for (Socket s : accepted) { + s.close(); + } + } + } + + int port() { + return socket.getLocalPort(); + } + } + + private static final class RecordingHandler implements QwpColumnBatchHandler { + final AtomicInteger batches = new AtomicInteger(); + final AtomicInteger failoverResets = new AtomicInteger(); + final AtomicInteger terminals = new AtomicInteger(); + volatile boolean ended; + volatile String errorMessage; + volatile byte errorStatus; + volatile boolean execDone; + volatile long terminalNanos; + + @Override + public void onBatch(QwpColumnBatch batch) { + batches.incrementAndGet(); + } + + @Override + public void onEnd(long totalRows) { + ended = true; + terminal(); + } + + @Override + public void onError(byte status, String message) { + errorStatus = status; + errorMessage = message; + terminal(); + } + + @Override + public void onExecDone(short opType, long rowsAffected) { + execDone = true; + terminal(); + } + + @Override + public void onFailoverReset(QwpServerInfo newNode) { + failoverResets.incrementAndGet(); + } + + void assertTimedOut() { + Assert.assertEquals("expected a query timeout, got status=" + errorStatus + ", message=" + errorMessage, + QwpConstants.STATUS_QUERY_TIMEOUT, errorStatus); + Assert.assertFalse("a timed-out query must not also complete", ended || execDone); + } + + private void terminal() { + terminalNanos = System.nanoTime(); + terminals.incrementAndGet(); + } + } + + /** + * Mock egress server scripted per test. Records every {@code QUERY_REQUEST} + * and {@code CANCEL}, and replies like a real server: one terminal frame per + * request, so a cancel of a query that already ended gets no reply. + */ + private static final class ScriptedServer implements TestWebSocketServer.WebSocketServerHandler, AutoCloseable { + final LinkedBlockingQueue cancels = new LinkedBlockingQueue<>(); + final LinkedBlockingQueue queries = new LinkedBlockingQueue<>(); + final AtomicInteger queryCount = new AtomicInteger(); + private final Set answered = ConcurrentHashMap.newKeySet(); + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(r -> { + Thread t = new Thread(r, "scripted-egress-server"); + t.setDaemon(true); + return t; + }); + volatile long lastSendNanos; + volatile Script onCancel; + volatile Script onQuery; + + @Override + public void close() { + scheduler.shutdownNow(); + } + + @Override + public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] data) { + if (data.length == 0) { + return; + } + if (data[0] == QwpEgressMsgKind.QUERY_REQUEST) { + int n = queryCount.incrementAndGet(); + queries.add(data); + Script script = onQuery; + if (script != null) { + script.run(this, client, data, n); + } + } else if (data[0] == QwpEgressMsgKind.CANCEL) { + cancels.add(data); + Script script = onCancel; + if (script != null) { + script.run(this, client, data, 0); + } + } + } + + void closeLater(TestWebSocketServer.ClientHandler client, long delayMs) { + sendLater(client, delayMs, closeFrame()); + } + + // Sends the terminal frame for requestId unless one was already sent. + void replyOnce(TestWebSocketServer.ClientHandler client, long requestId, long delayMs, byte[] frame) { + if (answered.add(requestId)) { + sendLater(client, delayMs, frame); + } + } + + void sendLater(TestWebSocketServer.ClientHandler client, long delayMs, byte[]... frames) { + scheduler.schedule(() -> { + lastSendNanos = System.nanoTime(); + for (byte[] frame : frames) { + send(client, frame); + } + }, delayMs, TimeUnit.MILLISECONDS); + } + } +} diff --git a/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpSpscQueueTest.java b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpSpscQueueTest.java index 50a930b19..9bbab56d5 100644 --- a/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpSpscQueueTest.java +++ b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpSpscQueueTest.java @@ -209,6 +209,75 @@ public void testNoUnparkWhenConsumerNotParked() { Assert.assertNull(q.poll()); } + @Test + public void testTimedTakeReturnsNullAtTheDeadline() throws InterruptedException { + QwpSpscQueue q = new QwpSpscQueue<>(4); + long start = System.nanoTime(); + Assert.assertNull("an empty queue must yield null once the deadline passes", + q.take(start + TimeUnit.MILLISECONDS.toNanos(100))); + long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); + Assert.assertTrue("timed take must wait until the deadline, waited " + elapsedMs + " ms", elapsedMs >= 100); + // The queue stays fully usable afterwards. + Assert.assertTrue(q.offer("after")); + Assert.assertEquals("after", q.take(System.nanoTime() + TimeUnit.SECONDS.toNanos(5))); + } + + @Test + public void testTimedTakeWithPastDeadlineStillReturnsAvailableValue() throws InterruptedException { + QwpSpscQueue q = new QwpSpscQueue<>(4); + long past = System.nanoTime() - TimeUnit.SECONDS.toNanos(1); + Assert.assertNull("a past deadline on an empty queue must not block", q.take(past)); + Assert.assertTrue(q.offer("ready")); + Assert.assertEquals("a value already available is returned even past the deadline", "ready", q.take(past)); + } + + @Test + public void testTimedTakeWakesOnOfferBeforeTheDeadline() throws Exception { + QwpSpscQueue q = new QwpSpscQueue<>(4); + AtomicReference taken = new AtomicReference<>(); + CountDownLatch done = new CountDownLatch(1); + Thread consumer = new Thread(() -> { + try { + taken.set(q.take(System.nanoTime() + TimeUnit.SECONDS.toNanos(30))); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + done.countDown(); + } + }, "spsc-timed-consumer"); + consumer.start(); + // Past the spin window, so the consumer is parked in parkNanos. + Thread.sleep(100); + Assert.assertTrue(q.offer("delivered")); + Assert.assertTrue("the producer's offer must unpark a timed take", done.await(5, TimeUnit.SECONDS)); + Assert.assertEquals("delivered", taken.get()); + consumer.join(1_000); + } + + @Test + public void testTimedTakeInterruptedThrows() throws Exception { + QwpSpscQueue q = new QwpSpscQueue<>(4); + AtomicReference caught = new AtomicReference<>(); + CountDownLatch done = new CountDownLatch(1); + Thread consumer = new Thread(() -> { + try { + q.take(System.nanoTime() + TimeUnit.SECONDS.toNanos(30)); + caught.set(new AssertionError("expected InterruptedException")); + } catch (Throwable t) { + caught.set(t); + } finally { + done.countDown(); + } + }, "spsc-timed-interrupt"); + consumer.start(); + Thread.sleep(100); + consumer.interrupt(); + Assert.assertTrue(done.await(5, TimeUnit.SECONDS)); + Assert.assertTrue("expected InterruptedException, got " + caught.get(), + caught.get() instanceof InterruptedException); + consumer.join(1_000); + } + @Test public void testCapacityOneIsRoundedToOne() { // Edge case: requesting capacity 1 still yields a power-of-two ring (1). diff --git a/core/src/test/java/io/questdb/client/test/impl/QueryCloseDrainTest.java b/core/src/test/java/io/questdb/client/test/impl/QueryCloseDrainTest.java index 76a4adb71..228c2a9fd 100644 --- a/core/src/test/java/io/questdb/client/test/impl/QueryCloseDrainTest.java +++ b/core/src/test/java/io/questdb/client/test/impl/QueryCloseDrainTest.java @@ -34,6 +34,8 @@ import java.lang.reflect.Field; import java.lang.reflect.Method; import java.util.ArrayList; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.Consumer; /** @@ -127,6 +129,65 @@ public void testCloseReturnsWorkerWhenAlreadyDrained() throws Exception { }); } + @Test(timeout = 30_000) + public void testCloseDiscardsWorkerThatStaysBusy() throws Exception { + TestUtils.assertMemoryLeak(() -> { + try (QueryClientPool pool = new QueryClientPool( + CFG, 0, 2, 1_000L, Long.MAX_VALUE, Long.MAX_VALUE, NO_CONNECT)) { + setCloseQueryTimeout(pool, 150L); + QueryWorker w = pool.acquire(); + long gen = generation(w); + // done stays true (the caller has its outcome), but the worker never + // finishes the job -- a drain that outlives the close budget. + setRunning(w, true); + + long startNanos = System.nanoTime(); + closeQuery(w, gen); + long elapsedMs = (System.nanoTime() - startNanos) / 1_000_000; + + Assert.assertTrue("close() must wait about the close budget, elapsed=" + elapsedMs, elapsedMs >= 120); + Assert.assertFalse("a worker that is still busy must be discarded, not returned to the pool", + allWorkers(pool).contains(w)); + } + }); + } + + @Test(timeout = 30_000) + public void testCloseWaitsForBusyWorkerToFinishBeforeReturningIt() throws Exception { + TestUtils.assertMemoryLeak(() -> { + try (QueryClientPool pool = new QueryClientPool( + CFG, 0, 2, 1_000L, Long.MAX_VALUE, Long.MAX_VALUE, NO_CONNECT)) { + setCloseQueryTimeout(pool, 5_000L); + QueryWorker w = pool.acquire(); + long gen = generation(w); + // done stays true, but the worker is still inside the job's runOn() -- + // as after a query timeout, when it drains the aborted query after + // the caller has been released. + setRunning(w, true); + Thread finisher = new Thread(() -> { + try { + Thread.sleep(200); + setRunning(w, false); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + finisher.start(); + + long startNanos = System.nanoTime(); + closeQuery(w, gen); + long elapsedMs = (System.nanoTime() - startNanos) / 1_000_000; + finisher.join(); + + Assert.assertTrue("close() must wait for the job to finish, elapsed=" + elapsedMs, elapsedMs >= 150); + Assert.assertTrue("a worker that finished in time must be returned to the pool", + allWorkers(pool).contains(w)); + QueryWorker again = pool.acquire(); + Assert.assertSame("the returned worker must be handed out again", w, again); + } + }); + } + @SuppressWarnings("unchecked") private static ArrayList allWorkers(QueryClientPool pool) throws Exception { Field f = QueryClientPool.class.getDeclaredField("all"); @@ -165,6 +226,25 @@ private static void setCloseQueryTimeout(QueryClientPool pool, long millis) thro f.setLong(pool, millis); } + // Sets the worker's running flag under its signalLock and wakes idle waiters, + // exactly as its run loop does around a job's runOn(). + private static void setRunning(QueryWorker w, boolean running) throws Exception { + Field lockF = QueryWorker.class.getDeclaredField("signalLock"); + Field runningF = QueryWorker.class.getDeclaredField("running"); + Field idleF = QueryWorker.class.getDeclaredField("idleCondition"); + lockF.setAccessible(true); + runningF.setAccessible(true); + idleF.setAccessible(true); + ReentrantLock lock = (ReentrantLock) lockF.get(w); + lock.lock(); + try { + runningF.setBoolean(w, running); + ((Condition) idleF.get(w)).signalAll(); + } finally { + lock.unlock(); + } + } + private static void setDone(QueryWorker w, boolean done) throws Exception { Object impl = queryImpl(w); Field doneF = impl.getClass().getDeclaredField("done"); diff --git a/core/src/test/java/io/questdb/client/test/impl/QueryImplResetTest.java b/core/src/test/java/io/questdb/client/test/impl/QueryImplResetTest.java index bedbba92f..5ce9841d9 100644 --- a/core/src/test/java/io/questdb/client/test/impl/QueryImplResetTest.java +++ b/core/src/test/java/io/questdb/client/test/impl/QueryImplResetTest.java @@ -71,10 +71,12 @@ public void testResetForBorrowClearsBuilderState() throws Exception { Field bindsF = queryImplClass.getDeclaredField("userBinds"); Field sqlBufF = queryImplClass.getDeclaredField("sqlBuffer"); Field doneF = queryImplClass.getDeclaredField("done"); + Field timeoutF = queryImplClass.getDeclaredField("timeoutMillis"); handlerF.setAccessible(true); bindsF.setAccessible(true); sqlBufF.setAccessible(true); doneF.setAccessible(true); + timeoutF.setAccessible(true); // Seed builder state as a prior borrow would have left it. handlerF.set(q, new NoopHandler()); @@ -82,6 +84,7 @@ public void testResetForBorrowClearsBuilderState() throws Exception { // no-op }); ((StringSink) sqlBufF.get(q)).put("SELECT 1"); + timeoutF.setLong(q, 0L); // a prior borrow disabled the timeout doneF.setBoolean(q, false); Method reset = queryImplClass.getDeclaredMethod("resetForBorrow"); @@ -97,6 +100,8 @@ public void testResetForBorrowClearsBuilderState() throws Exception { 0, sqlBuffer.length()); Assert.assertTrue("done must be true so the handle starts idle, not in flight", doneF.getBoolean(q)); + Assert.assertEquals("the query timeout must fall back to the configured default (-1)", + -1L, timeoutF.getLong(q)); }); } diff --git a/core/src/test/java/io/questdb/client/test/impl/QwpQueryClientConfigHonoredTest.java b/core/src/test/java/io/questdb/client/test/impl/QwpQueryClientConfigHonoredTest.java index 10b1e3c03..b8650afec 100644 --- a/core/src/test/java/io/questdb/client/test/impl/QwpQueryClientConfigHonoredTest.java +++ b/core/src/test/java/io/questdb/client/test/impl/QwpQueryClientConfigHonoredTest.java @@ -62,6 +62,7 @@ public void testEveryEgressKeyIsHonored() throws Exception { assertHonored("failover_max_duration_ms=56000", "failover_max_duration_ms", 56000L); assertHonored("max_batch_rows=512", "max_batch_rows", 512); assertHonored("initial_credit=65536", "initial_credit", 65536L); + assertHonored("query_timeout_ms=1500", "query_timeout_ms", 1500L); assertHonored("buffer_pool_size=3", "buffer_pool_size", 3); assertHonored("compression=zstd", "compression", "zstd"); assertHonored("compression_level=9", "compression_level", 9); diff --git a/examples/src/main/java/com/example/query/QueryCancellationExample.java b/examples/src/main/java/com/example/query/QueryCancellationExample.java index b9def4bb9..57bda2880 100644 --- a/examples/src/main/java/com/example/query/QueryCancellationExample.java +++ b/examples/src/main/java/com/example/query/QueryCancellationExample.java @@ -19,6 +19,9 @@ * blocking {@link Completion#await()} throws {@link QueryException} carrying * that status. If the query finished before the cancel landed, {@code await()} * returns normally -- {@code cancel()} is a no-op on an already-terminal query. + *

    + * To stop queries that run too long without driving the cancel yourself, give + * them a timeout instead; see {@link QueryTimeoutExample}. */ public class QueryCancellationExample { diff --git a/examples/src/main/java/com/example/query/QueryTimeoutExample.java b/examples/src/main/java/com/example/query/QueryTimeoutExample.java new file mode 100644 index 000000000..1408b8730 --- /dev/null +++ b/examples/src/main/java/com/example/query/QueryTimeoutExample.java @@ -0,0 +1,66 @@ +package com.example.query; + +import io.questdb.client.Query; +import io.questdb.client.QueryException; +import io.questdb.client.QuestDB; +import io.questdb.client.cutlass.qwp.client.QwpColumnBatch; +import io.questdb.client.cutlass.qwp.client.QwpColumnBatchHandler; + +import java.util.concurrent.TimeUnit; + +/** + * Giving a query a timeout. + *

    + * {@link Query#timeout(long, TimeUnit)} bounds the whole query, measured from + * {@code submit()}; {@code query_timeout_ms} in the connection string sets a + * default for every query. When the timeout expires the query is stopped -- + * by the server when it supports per-query timeouts, otherwise by a cancel the + * client sends -- and {@code await()} throws a {@link QueryException} whose + * {@link QueryException#isTimeout()} is {@code true}. The pooled connection + * stays open, so the same handle goes on to run the next query on it. + *

    + * Compare {@link QueryCancellationExample}: {@code await(timeout)} only bounds + * how long the caller waits, the query keeps running until cancelled. + */ +public class QueryTimeoutExample { + + public static void main(String[] args) throws InterruptedException { + // query_timeout_ms: default timeout for every query of this handle's pool. + try (QuestDB db = QuestDB.connect("ws::addr=localhost:9000;query_timeout_ms=30000;"); + Query q = db.borrowQuery()) { + + q.sql("SELECT * FROM trades ORDER BY price") + .handler(new PrintingHandler()) + // Tighter than the default, for this handle only. + .timeout(5, TimeUnit.SECONDS); + try { + q.submit().await(); + System.out.println("finished within the timeout"); + } catch (QueryException e) { + if (!e.isTimeout()) { + throw e; + } + System.out.println("timed out: " + e.getMessage()); + } + + // The connection survived the timeout; run the next query on it. + q.sql("SELECT count() FROM trades").timeout(0, TimeUnit.SECONDS).submit().await(); + } + } + + private static final class PrintingHandler implements QwpColumnBatchHandler { + @Override + public void onBatch(QwpColumnBatch batch) { + // Process rows... kept minimal here. + } + + @Override + public void onEnd(long totalRows) { + System.out.println("done: " + totalRows + " rows"); + } + + @Override + public void onError(byte status, String message) { + } + } +} From 7baaf87de02c53595234772649ba1d508468f2b9 Mon Sep 17 00:00:00 2001 From: Vlad Ilyushchenko Date: Wed, 7 Oct 2026 20:49:43 +0100 Subject: [PATCH 2/2] Stop late frames from answering the next query A query that ended early on the client side left its remaining frames on the connection, and the next query read them as its own result. This fixes the three ways in. Handler throws on a query timeout. At the end of the grace period execute() reports the timeout through onError while the aborted query still runs on the connection. If onError threw, execute() skipped the drain that follows. FailoverProbeHandler now holds the exception, executeOnce keeps draining as it does for a handler that returns, and executeImpl rethrows the original throwable once the drain ends or the connection is given up. The facade worker reports an exception that escapes execute() as the outcome of the current submission. Rethrown after the drain, it could land after the caller, already released by the timeout, had submitted again on the same Query, and complete that newer submission with the old exception. QueryImpl now numbers submissions, and runOn() reports an escaping exception only while its own submission is current. The check and the report share doneLock with submit(). This also closes the same race, until now microseconds wide, for a handler that throws from onEnd or onExecDone. Handler throws from onBatch, or the thread is interrupted. The I/O thread now stamps every event of a query with that query's request id; connection-level events carry QueryEvent.ANY_REQUEST and still reach whichever query waits. executeOnce skips the events of other queries, so the next query no longer receives the leftovers of an abandoned one, and it handles its own deadline even while leftovers keep coming. An attempt that ends before its query's last event now cancels that query, so the leftovers stop coming. While the I/O thread still works through an abandoned query, the next query can be queued and cancelled. The I/O thread sent that CANCEL at once, ahead of the query itself, and the server drops a cancel of a query it does not know, so the cancel was lost. drainPendingCancel now sends only a cancel of the query it serves, holds one of a later query, and drops one of a query that has already ended. Interrupt before the I/O thread encodes the request. execute() returned while the I/O thread could still read the request holder, the bind buffer and the caller's SQL text, and the next execute() overwrote them: the interrupted query was never sent and the next one was sent twice, possibly mixing the two queries' SQL and binds. With leftovers of an abandoned query on the connection, the next execute() could also block forever on the single request slot, while the I/O thread waited for it to release a leftover batch. An interrupted execute() now withdraws a request the I/O thread has not picked up yet, or waits for the encoding the I/O thread has started, before it reports the interrupt. An I/O thread still not done after shutdownJoinMs gives up the connection, as on a query timeout. Behavior visible to callers: an exception from onError at the end of the grace period now surfaces when the drain ends, up to one more grace period later, which is when execute() already returns for a handler that does not throw. A handler that throws from onBatch, or an interrupt, now cancels its query, and the next query skips that query's leftovers instead of returning them. A handler that does not throw sees no change, and the per-batch path gains only one field store on the I/O thread and one comparison on the reading thread. New tests drive each path against the scripted mock server, through both QwpQueryClient and the QuestDB facade. --- .../client/cutlass/qwp/client/QueryEvent.java | 38 ++ .../qwp/client/QwpColumnBatchHandler.java | 7 +- .../cutlass/qwp/client/QwpEgressIoThread.java | 59 +- .../cutlass/qwp/client/QwpQueryClient.java | 199 ++++++- .../io/questdb/client/impl/QueryImpl.java | 44 +- .../client/test/QueryTimeoutFacadeTest.java | 195 ++++++- .../QwpQueryClientQueryTimeoutTest.java | 518 +++++++++++++++++- 7 files changed, 1019 insertions(+), 41 deletions(-) diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QueryEvent.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QueryEvent.java index d3083e4c7..71441d6a7 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QueryEvent.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QueryEvent.java @@ -29,9 +29,21 @@ * One event per {@code RESULT_BATCH} / {@code RESULT_END} / {@code QUERY_ERROR} * received from the server, plus a synthetic error event if the connection drops * mid-query. + *

    + * The I/O thread stamps the events of a query with that query's id + * ({@link #requestId}), so {@code execute()} can tell its own events from the + * leftovers of an earlier query that its {@code execute()} stopped waiting for. + * Connection-level events -- the connection failed, the I/O thread stopped, the + * client closed -- carry {@link #ANY_REQUEST} and reach whichever query waits. */ public class QueryEvent { + /** + * {@link #requestId} of an event that belongs to no query in particular, such + * as a connection failure. Whichever query waits for events must see it. + */ + public static final long ANY_REQUEST = -1L; + public static final int KIND_BATCH = 0; public static final int KIND_END = 1; public static final int KIND_ERROR = 2; @@ -50,12 +62,14 @@ public class QueryEvent { public byte errorStatus; // valid for KIND_ERROR public int kind; public short opType; // valid for KIND_EXEC_DONE (matches CompiledQuery.SELECT/INSERT/etc.) + public long requestId = ANY_REQUEST; // query the event belongs to; see forRequest() public long rowsAffected; // valid for KIND_EXEC_DONE public long totalRows; // valid for KIND_END public QueryEvent asBatch(QwpBatchBuffer buffer) { this.kind = KIND_BATCH; this.buffer = buffer; + this.requestId = ANY_REQUEST; return this; } @@ -63,6 +77,7 @@ public QueryEvent asEnd(long totalRows) { this.kind = KIND_END; this.buffer = null; this.totalRows = totalRows; + this.requestId = ANY_REQUEST; return this; } @@ -71,6 +86,7 @@ public QueryEvent asError(byte status, String message) { this.buffer = null; this.errorStatus = status; this.errorMessage = message; + this.requestId = ANY_REQUEST; return this; } @@ -79,6 +95,7 @@ public QueryEvent asExecDone(short opType, long rowsAffected) { this.buffer = null; this.opType = opType; this.rowsAffected = rowsAffected; + this.requestId = ANY_REQUEST; return this; } @@ -87,9 +104,29 @@ public QueryEvent asTransportError(byte status, String message) { this.buffer = null; this.errorStatus = status; this.errorMessage = message; + this.requestId = ANY_REQUEST; + return this; + } + + /** + * Ties the event to the query with id {@code requestId}. Call it after the + * {@code asX()} builder, which leaves the event at {@link #ANY_REQUEST}. + */ + public QueryEvent forRequest(long requestId) { + this.requestId = requestId; return this; } + /** + * Returns {@code true} when the event belongs to a query other than + * {@code requestId}: a leftover of an earlier query that its + * {@code execute()} stopped waiting for. An {@link #ANY_REQUEST} event + * belongs to every query. + */ + public boolean isForOtherRequest(long requestId) { + return this.requestId != ANY_REQUEST && this.requestId != requestId; + } + /** * Clears object references and resets primitive fields so a pooled event is * safe to reuse across queries. The I/O thread calls the {@code asX(...)} @@ -106,5 +143,6 @@ public void reset() { this.opType = 0; this.rowsAffected = 0; this.totalRows = 0; + this.requestId = ANY_REQUEST; } } diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpColumnBatchHandler.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpColumnBatchHandler.java index 6237e1de0..86feaf128 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpColumnBatchHandler.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpColumnBatchHandler.java @@ -38,7 +38,12 @@ * Exception contract: if any callback method throws, the * exception propagates out of the {@link QwpQueryClient#execute} call on the * caller's thread and no further callbacks fire for that query. The connection - * remains usable for subsequent queries. + * remains usable for subsequent queries: when {@link #onBatch} throws while the + * query is still running, {@code execute} cancels the query, and the next query + * skips whatever the cancelled one still sends. When {@link #onError} reports a + * query timeout at the end of the grace period, the aborted query is still + * running as well; if that callback throws, {@code execute} first drains the + * query, as it does when the callback returns, and only then rethrows. */ public interface QwpColumnBatchHandler { diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java index 02e0d787d..3c51f44dd 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpEgressIoThread.java @@ -89,8 +89,9 @@ public class QwpEgressIoThread implements Runnable, WebSocketFrameHandler { // reads the fields synchronously in {@link #sendQueryRequest} and does not // retain a reference past that call. Reuse avoids a per-submit allocation // -- one in-flight query per client makes this safe: the worker that - // mutates pendingRequest is blocked on the events queue until the I/O - // thread has finished consuming the previous instance. + // mutates pendingRequest does not return from execute() before the I/O + // thread has encoded the previous instance, or it was withdrawn + // (withdrawRequest). private final QueryRequest pendingRequest = new QueryRequest(); // Single-slot request queue (Phase-1 allows one in-flight query). private final BlockingQueue requests = new ArrayBlockingQueue<>(1); @@ -113,6 +114,8 @@ public class QwpEgressIoThread implements Runnable, WebSocketFrameHandler { private volatile boolean closed; private boolean creditEnabled; private boolean currentQueryDone; + // Query being served. Stamped on the events of that query (QueryEvent.requestId); + // connection-level events stay at QueryEvent.ANY_REQUEST. private long currentRequestId = -1L; // requestId of the last QUERY_REQUEST fully encoded into sendScratch. Once a // request's id is published here the I/O thread no longer reads the caller's @@ -299,9 +302,11 @@ public void releaseEvent(QueryEvent event) { /** * Queues a CANCEL frame for {@code requestId} to be sent by the I/O thread * between the next two {@code receiveFrame} iterations (typically within - * {@link #POLL_TIMEOUT_MS}). Safe to call from any thread. If a CANCEL for - * the same (or another) requestId is already pending, the newer id wins -- - * multiple concurrent cancels coalesce into one send. + * {@link #POLL_TIMEOUT_MS}) once it serves that query. Safe to call from any + * thread. If a CANCEL for the same (or another) requestId is already pending, + * the newer id wins -- multiple concurrent cancels coalesce into one send. + * See {@link #drainPendingCancel()} for a cancel of a query that has not been + * sent yet, or has already ended. */ public void requestCancel(long requestId) { pendingCancelRequestId.set(requestId); @@ -435,6 +440,18 @@ public QueryEvent takeEvent(long deadlineNanos) throws InterruptedException { return events.take(deadlineNanos); } + /** + * Takes back the request {@code requestId} that {@link #submitQuery} queued, + * provided this thread has not picked it up yet. Returns {@code true} when it + * was still queued: it is never sent, and the caller may reuse its SQL text + * and bind payload at once. Returns {@code false} when this thread has it; + * it reads them until {@link #isRequestEncoded} holds. Removal and pick-up + * are atomic with respect to each other. + */ + public boolean withdrawRequest(long requestId) { + return pendingRequest.requestId == requestId && requests.remove(pendingRequest); + } + /** * Takes a pre-allocated {@link QueryEvent} from the pool for the I/O thread * to populate before offering to {@link #events}. Falls back to a fresh @@ -452,8 +469,7 @@ private QueryEvent borrowEvent() { } private void decodeAndEmitError(long payload, int payloadLen) { - QueryEvent ev = decodeError(payload, payloadLen); - events.offer(ev); + events.offer(decodeError(payload, payloadLen).forRequest(currentRequestId)); } /** @@ -492,7 +508,7 @@ private void decodeAndEmitExecDone(long payload, int payloadLen) { emitTerminalTransportError("EXEC_DONE frame truncated mid rows_affected varint"); return; } - events.offer(new QueryEvent().asExecDone(opType, rowsAffected)); + events.offer(new QueryEvent().asExecDone(opType, rowsAffected).forRequest(currentRequestId)); } /** @@ -549,7 +565,7 @@ private void decodeAndEmitResultEnd(long payload, int payloadLen) { emitTerminalTransportError("RESULT_END frame truncated mid total_rows varint"); return; } - events.offer(new QueryEvent().asEnd(total)); + events.offer(new QueryEvent().asEnd(total).forRequest(currentRequestId)); } /** @@ -557,10 +573,22 @@ private void decodeAndEmitResultEnd(long payload, int payloadLen) { * thread at every loop boundary so a cancel set by a user thread reaches * the server regardless of whether the I/O thread was waiting on a frame * or on a free buffer. + *

    + * Only a cancel of the query being served goes out. One for a later query + * stays pending: the user thread can queue that query, and cancel it, while + * this thread still works through an earlier query whose {@code execute()} + * stopped waiting for it. Sent now, the CANCEL would reach the server before + * the query does, and the server drops a cancel of a query it does not know. + * One for an earlier query, which has already ended, is dropped. Request ids + * only grow, so comparing them tells the three apart. */ private void drainPendingCancel() { - long id = pendingCancelRequestId.getAndSet(-1L); - if (id >= 0L) { + long id = pendingCancelRequestId.get(); + if (id < 0L || id > currentRequestId) { + return; + } + // The CAS keeps a cancel the user thread set meanwhile, which is newer. + if (pendingCancelRequestId.compareAndSet(id, -1L) && id == currentRequestId) { sendCancel(id); } } @@ -659,7 +687,7 @@ private void handleResultBatch(long payloadPtr, int payloadLen) { currentQueryDone = true; return; } - events.offer(borrowEvent().asBatch(buf)); + events.offer(borrowEvent().asBatch(buf).forRequest(currentRequestId)); // Park on the release latch. Returning sooner would let receiveFrame // compact the WebSocket recv buffer, overwriting the bytes that the // user-visible column pointers still reference. User thread's @@ -808,9 +836,10 @@ public interface TerminalFailureListener { /** * Mutable request holder reused across submits. Safe to reuse because at * most one query is in flight per client: the worker thread mutates fields - * and offers the instance into {@link #requests}, then blocks on the events - * queue until the I/O thread has fully consumed the previous instance and - * delivered a terminal event. + * and offers the instance into {@link #requests}, and does not return from + * {@code execute()} before this thread has encoded the instance (see + * {@link #isRequestEncoded}) or the worker withdrew it + * ({@link #withdrawRequest}). */ private static final class QueryRequest { int bindCount; diff --git a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java index c99ac5a79..94d11e1d9 100644 --- a/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java +++ b/core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java @@ -49,6 +49,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; /** * QWP egress (query results) client. @@ -179,6 +180,9 @@ public class QwpQueryClient implements QuietCloseable { // deadline arithmetic (deadline + 2 * grace) on System.nanoTime() values // clear of overflow; ~73 years is "no limit" for every practical purpose. private static final long MAX_TIMEOUT_NANOS = Long.MAX_VALUE / 8; + // How often an interrupted execute() checks whether the I/O thread has + // encoded its request (see retractRequest). + private static final long REQUEST_ENCODED_POLL_NANOS = TimeUnit.MICROSECONDS.toNanos(100); // States of a query running under a timeout (see executeOnce). // RUNNING: before the deadline, events are delivered normally. private static final int TIMEOUT_PHASE_RUNNING = 0; @@ -993,10 +997,11 @@ public void execute(CharSequence sql, QwpBindSetter binds, QwpColumnBatchHandler * ({@link #withQueryTimeoutGrace(long)}) after the deadline, the handler is * told about the timeout anyway and this call keeps draining the aborted * query for up to one more grace period, so the connection can still be - * reused. Only a connection that stays silent through both grace periods is - * treated as failed (see {@link #hasTerminalFailure()}). Unless a handler - * callback blocks, this call therefore returns about two grace periods after - * the timeout at the latest. + * reused. That drain happens even if the handler's {@code onError} throws; + * the exception then propagates once the drain ends. Only a connection that + * stays silent through both grace periods is treated as failed (see + * {@link #hasTerminalFailure()}). Unless a handler callback blocks, this call + * therefore returns about two grace periods after the timeout at the latest. * * @param timeoutMs query timeout in milliseconds; {@code 0} runs the query * without a timeout @@ -1616,6 +1621,18 @@ private static String defaultClientId() { return "questdb-java-egress/1.0.0"; } + /** + * Hands an event the executing thread does not deliver back to the I/O + * thread, with its batch buffer: the I/O thread waits for each published + * batch to be released. + */ + private static void discardEvent(QwpEgressIoThread io, QueryEvent ev) { + if (ev.kind == QueryEvent.KIND_BATCH && ev.buffer != null) { + io.releaseBuffer(ev.buffer); + } + io.releaseEvent(ev); + } + private static boolean isPast(long deadlineNanos) { return System.nanoTime() - deadlineNanos >= 0; } @@ -1647,17 +1664,27 @@ private static long remainingMillisCeil(long deadlineNanos) { return remainingNanos <= 0L ? 0L : (remainingNanos + 999_999L) / 1_000_000L; } + /** + * Throws {@code t} as is, without wrapping it. Callbacks of + * {@link QwpColumnBatchHandler} declare no checked exceptions, so this only + * matters for a handler that throws one anyway: it reaches the caller of + * {@link #execute} unchanged, as it does when it is not held back. + */ + @SuppressWarnings("unchecked") + private static void throwUnchecked(Throwable t) throws T { + throw (T) t; + } + private static long toBoundedNanos(long millis) { return Math.min(TimeUnit.MILLISECONDS.toNanos(millis), MAX_TIMEOUT_NANOS); } /** - * Gives up on a connection that left a timed-out query unanswered through - * both grace periods (hung server, black-holed network), or whose I/O thread - * never even sent the query. Latches a terminal failure, so the next - * {@link #execute} replaces the connection ({@code failover=on}) or reports - * the failure ({@code failover=off}) and a pool discards the client, then - * stops the I/O thread. + * Gives up on the connection. Latches a terminal failure carrying + * {@code failureMessage}, so the next {@link #execute} replaces the + * connection ({@code failover=on}) or reports the failure + * ({@code failover=off}) and a pool discards the client, then stops the I/O + * thread. *

    * Runs on the executing thread, the sole consumer of the event queue, and * keeps releasing every batch the I/O thread still publishes while it winds @@ -1665,14 +1692,11 @@ private static long toBoundedNanos(long millis) { * released, and would otherwise outlive the joins in {@link #close()} and * {@link #cleanupFailedConnect()}. */ - private void abandonUnresponsiveConnection(QwpEgressIoThread io) { + private void abandonConnection(QwpEgressIoThread io, String failureMessage) { GenerationListener listener = currentGenerationListener; if (listener != null) { - listener.onTerminalFailure(WebSocketResponse.STATUS_INTERNAL_ERROR, - "connection stopped responding after a query timeout"); + listener.onTerminalFailure(WebSocketResponse.STATUS_INTERNAL_ERROR, failureMessage); } - LOG.warn("QwpQueryClient connection did not end a timed-out query within two grace periods of {}ms; " - + "closing it", queryTimeoutGraceMs); io.shutdown(); Thread handle = ioThreadHandle; if (handle == null) { @@ -1695,10 +1719,7 @@ private void abandonUnresponsiveConnection(QwpEgressIoThread io) { continue; } if (ev != null) { - if (ev.kind == QueryEvent.KIND_BATCH && ev.buffer != null) { - io.releaseBuffer(ev.buffer); - } - io.releaseEvent(ev); + discardEvent(io, ev); } } if (interrupted) { @@ -1706,6 +1727,17 @@ private void abandonUnresponsiveConnection(QwpEgressIoThread io) { } } + /** + * Gives up on a connection that left a timed-out query unanswered through + * both grace periods (hung server, black-holed network), or whose I/O thread + * never even sent the query. See {@link #abandonConnection}. + */ + private void abandonUnresponsiveConnection(QwpEgressIoThread io) { + LOG.warn("QwpQueryClient connection did not end a timed-out query within two grace periods of {}ms; " + + "closing it", queryTimeoutGraceMs); + abandonConnection(io, "connection stopped responding after a query timeout"); + } + /** * Applies the deadline of a failover reconnect made on behalf of a query * with a timeout (see {@link #connectDeadlineActive}) to one endpoint step's @@ -1883,6 +1915,10 @@ private void executeImpl( attempt++; FailoverProbeHandler probe = new FailoverProbeHandler(handler); executeOnce(sql, binds, probe, resetSymbolDict, timeoutMs, deadlineNanos); + // A handler failure held back while the attempt drained a timed-out + // query propagates now: the query has ended, or the connection was + // given up on. + probe.rethrowDeferredFailure(); if (!probe.transportFailureIntercepted) { return; } @@ -2021,6 +2057,15 @@ private void executeImpl( * connection, which carries one query at a time, can serve the next query. * A connection still silent after that is abandoned. The handler sees * exactly one terminal callback in every case. + *

    + * An attempt that ends before its query's last frame -- the handler threw + * from {@code onBatch}, or the thread was interrupted -- leaves the query + * running on the connection, and cancels it. The I/O thread still works + * through it before it sends the next query, so the next attempt skips the + * leftover events: each event carries the id of its query (see + * {@link QueryEvent#requestId}). An interrupt that comes before the I/O + * thread has encoded the request takes the request back, or waits for the + * encoding, before the caller hears about it (see {@link #retractRequest}). */ private void executeOnce( CharSequence sql, @@ -2101,11 +2146,27 @@ private void executeOnce( // STATUS_CANCELLED reply then reports that cancel, not the timeout. boolean cancelledByUser = false; long waitDeadlineNanos = deadlineNanos; + boolean submitted = false; + // Set once the attempt has taken its query's last event, or given up on + // the connection: the query no longer runs on it. + boolean settled = false; try { io.submitQuery(sql, requestId, initialCreditBytes, bindValues.count(), bindValues.bufferPtr(), bindValues.bufferLen(), queryFlags, wireTimeoutMs); + submitted = true; while (true) { QueryEvent ev = timed ? io.takeEvent(waitDeadlineNanos) : io.takeEvent(); + if (ev != null && ev.isForOtherRequest(requestId)) { + // A leftover of an earlier query whose attempt ended before that + // query did. Not this query's; skip it. + discardEvent(io, ev); + if (!timed || !isPast(waitDeadlineNanos)) { + continue; + } + // Leftovers can keep coming past the wait deadline, and the + // queue still returns them then: handle the deadline now. + ev = null; + } if (ev == null) { // waitDeadlineNanos passed without an event. if (phase == TIMEOUT_PHASE_RUNNING) { @@ -2127,19 +2188,25 @@ private void executeOnce( io.requestCancel(requestId); phase = TIMEOUT_PHASE_DRAINING; waitDeadlineNanos = deadlineNanos + 2 * graceNanos; - probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs) + probe.onErrorWhileDraining(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs) + "; the server did not end the query within the " + queryTimeoutGraceMs + "ms grace period"); continue; } final boolean callerWaiting = phase != TIMEOUT_PHASE_DRAINING; abandonUnresponsiveConnection(io); + settled = true; if (callerWaiting) { probe.onError(requestId, QwpConstants.STATUS_QUERY_TIMEOUT, queryTimeoutMessage(timeoutMs) + "; the connection did not respond and was closed"); } return; } + if (ev.kind != QueryEvent.KIND_BATCH) { + // Every other event ends the query on the connection, and the + // attempt returns once the switch below handles it. + settled = true; + } try { switch (ev.kind) { case QueryEvent.KIND_BATCH: @@ -2221,12 +2288,23 @@ private void executeOnce( } } catch (InterruptedException ie) { Thread.currentThread().interrupt(); + // Before the caller hears about the interrupt: the I/O thread must be + // done with the request (see retractRequest). + if (submitted && !io.isRequestEncoded(requestId) && retractRequest(io, requestId)) { + settled = true; + } // Interrupt on the user thread is not a transport failure; surface directly. // A draining attempt has already given the handler its terminal callback. if (phase != TIMEOUT_PHASE_DRAINING) { probe.deliverFinal(requestId, "interrupted while waiting for server response"); } } finally { + if (submitted && !settled) { + // The attempt ends while its query still runs: the handler threw, + // or the thread was interrupted. Cancel the query, so the leftovers + // the next attempt skips stop coming. + io.requestCancel(requestId); + } currentRequestId = -1L; } } @@ -2480,6 +2558,54 @@ private long resolveQueryFlags(boolean resetSymbolDict, boolean timed) { return flags; } + /** + * Makes sure the I/O thread no longer reads the request {@code requestId}, + * which an interrupted {@code execute()} stops waiting for before the I/O + * thread has encoded it. Until then the I/O thread reads the caller's SQL + * text, which the caller may change once it learns about the interrupt, and + * this client's request holder and bind buffer, which the next query reuses. + *

    + * A request the I/O thread has not picked up yet is withdrawn: it is never + * sent. One it has picked up is awaited -- encoding is CPU work, done + * microseconds after the pick-up -- without regard to interrupts. An I/O + * thread that is still not done after {@link #shutdownJoinMs} is stuck, and + * the connection is given up, as on a query timeout. + * + * @return {@code true} when the query does not run on the connection: it was + * withdrawn, or the connection was given up + */ + private boolean retractRequest(QwpEgressIoThread io, long requestId) { + if (io.withdrawRequest(requestId)) { + return true; + } + final Thread handle = ioThreadHandle; + // parkNanos() returns at once while the interrupt flag is set, so clear it + // for the wait, and restore it after. + boolean interrupted = Thread.interrupted(); + try { + final long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(shutdownJoinMs); + while (!io.isRequestEncoded(requestId)) { + if (handle == null || !handle.isAlive()) { + // A stopped I/O thread reads nothing more. + return true; + } + if (isPast(deadlineNanos)) { + LOG.warn("QwpQueryClient I/O thread did not finish sending an interrupted query within {}ms; " + + "closing the connection", shutdownJoinMs); + abandonConnection(io, "I/O thread stopped responding"); + return true; + } + LockSupport.parkNanos(REQUEST_ENCODED_POLL_NANOS); + interrupted |= Thread.interrupted(); + } + return false; + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + private void runUpgradeWithTimeout(Endpoint ep, String authHeader) { // Connect first, OUTSIDE the upgrade try. A connect-phase failure -- // including a connect_timeout overage flagged via flagAsTimeout() -- must @@ -2554,6 +2680,9 @@ private static final class Endpoint { */ private static final class FailoverProbeHandler implements QwpColumnBatchHandler { final QwpColumnBatchHandler delegate; + // Thrown by the handler while its query was still running; see + // onErrorWhileDraining(). + Throwable deferredFailure; String interceptedMessage; long interceptedRequestId = -1L; byte interceptedStatus; @@ -2632,6 +2761,36 @@ void markTransportFailure(long requestId, byte status, String message) { interceptedStatus = status; interceptedMessage = message; } + + /** + * Delivers the error that ends a query for the caller while the query is + * still running on the connection: the timeout reported at the end of the + * grace period, before the aborted query has drained. The connection + * carries one query at a time, so a throw from the handler must not cut + * that drain short -- the next query would read this query's remaining + * frames as its own. The throw is held instead, and + * {@link #rethrowDeferredFailure()} raises it once the attempt is over. + */ + void onErrorWhileDraining(long requestId, byte status, String message) { + try { + delegate.onError(requestId, status, message); + } catch (Throwable t) { + deferredFailure = t; + } + } + + /** + * Raises the handler failure held back by {@link #onErrorWhileDraining}, + * if any. Called once the attempt has drained its query, or given up on + * the connection. + */ + void rethrowDeferredFailure() { + Throwable t = deferredFailure; + if (t != null) { + deferredFailure = null; + QwpQueryClient.throwUnchecked(t); + } + } } /** diff --git a/core/src/main/java/io/questdb/client/impl/QueryImpl.java b/core/src/main/java/io/questdb/client/impl/QueryImpl.java index 3f5aa2e4d..c15130015 100644 --- a/core/src/main/java/io/questdb/client/impl/QueryImpl.java +++ b/core/src/main/java/io/questdb/client/impl/QueryImpl.java @@ -79,6 +79,10 @@ final class QueryImpl { private volatile boolean done = true; private volatile String resultMessage; private volatile byte resultStatus; + // Counts submit() calls. Written under doneLock; volatile so runOn() can read + // it without the lock. Lets runOn() tell its own submission from a later one, + // see signalUnexpected(long, Throwable). + private volatile long submissionSeq; // Stamped by submit() for the worker thread: when the query was submitted and // the timeout it runs under (0 = none), so runOn() can hand the client only // the part of the budget that is left. Published by worker.dispatch(). @@ -288,6 +292,7 @@ void submit(long gen) { doneLock.lock(); try { done = false; + submissionSeq++; resultStatus = 0; resultMessage = null; unexpectedError = null; @@ -297,6 +302,11 @@ void submit(long gen) { worker.dispatch(this); } + private static String describe(Throwable t) { + String message = t.getMessage(); + return message != null ? message : t.getClass().getSimpleName(); + } + private void applyBinds(QwpBindValues binds) { // Runs inside execute() after it cleared the client's cancel latch for // this query, and before the request is sent: set the latch again if the @@ -377,6 +387,29 @@ private void signalDone(byte status, String message, Throwable unexpected) { } } + /** + * Signals an error that escaped the run of {@code submission}, unless a later + * {@link #submit} has taken that submission's place. A run can end with an + * exception after it signalled its outcome: when a handler throws on a query + * timeout, the client rethrows the exception only once the aborted query has + * drained, and the caller, already released, may have submitted again + * meanwhile. That exception must not become the newer submission's outcome. + * The check and the signal share {@code doneLock}, which {@code submit()} + * takes to start a submission. + */ + private void signalUnexpected(long submission, Throwable t) { + // getMessage() may be user code: run it before taking the lock submit() needs. + final String message = describe(t); + doneLock.lock(); + try { + if (submission == submissionSeq) { + signalDone((byte) 0, message, t); + } + } finally { + doneLock.unlock(); + } + } + private void throwIfFailed() { Throwable unexpected = unexpectedError; if (unexpected != null) { @@ -408,6 +441,9 @@ void resetForBorrow() { } void runOn(QwpQueryClient client) { + // The submission this run executes. No later submit() can happen before + // the run signals this one's outcome, so the value is stable until then. + final long submission = submissionSeq; long timeoutArg = 0; if (submitTimeoutMillis > 0) { // Only what is left of the budget since submit(); at least 1ms so the @@ -424,7 +460,11 @@ void runOn(QwpQueryClient client) { // can call sql(...) again only once done==true, which execute() signals // either at the terminal event or, after a query timeout, once the I/O // thread has encoded the request (see QwpQueryClient.executeOnce). - client.execute(sqlBuffer, wireBinds, wrappingHandler, false, timeoutArg); + try { + client.execute(sqlBuffer, wireBinds, wrappingHandler, false, timeoutArg); + } catch (Throwable t) { + signalUnexpected(submission, t); + } } /** @@ -432,7 +472,7 @@ void runOn(QwpQueryClient client) { * exception escaping {@code execute()} before any handler callback). */ void signalUnexpected(Throwable t) { - signalDone((byte) 0, t.getMessage() != null ? t.getMessage() : t.getClass().getSimpleName(), t); + signalDone((byte) 0, describe(t), t); } private final class WrappingHandler implements QwpColumnBatchHandler { diff --git a/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java b/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java index 2e8913737..d034922e0 100644 --- a/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java +++ b/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java @@ -38,6 +38,7 @@ import org.junit.Assert; import org.junit.Test; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; @@ -55,7 +56,8 @@ * holds, while the pooled, authenticated connection stays open and serves the * following queries. A worker still draining a timed-out query is not handed to * the next borrower before it is idle, and one whose connection stopped - * responding is replaced rather than reused. + * responding is replaced rather than reused. A handler that throws -- on the + * timeout, or mid-result -- fails only its own submission. */ public class QueryTimeoutFacadeTest { @@ -133,6 +135,112 @@ public void testConfiguredDefaultAppliesAndHandleCanOverrideIt() throws Exceptio }); } + @Test(timeout = 30_000) + public void testThrowingBatchHandlerFailsOnlyItsOwnSubmission() throws Exception { + TestUtils.assertMemoryLeak(() -> { + Script script = new Script(); + script.onQuery = (s, c, id, n) -> { + if (n == 1) { + s.send(c, firstBatch(id)); + // The rest of the result arrives after the handler has thrown. + s.sendLater(c, 300, nextBatch(id)); + s.sendLater(c, 300, resultEnd(id)); + } else { + s.send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT); + QuestDB db = QuestDB.connect(config(server)); + Query q = db.borrowQuery()) { + final IllegalStateException failure = new IllegalStateException("handler failed"); + q.sql("SELECT s FROM t").handler(new QwpColumnBatchHandler() { + @Override + public void onBatch(QwpColumnBatch batch) { + throw failure; + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + } + }); + try { + q.submit().await(); + Assert.fail("the submission whose handler threw must fail"); + } catch (QueryException e) { + Assert.assertSame("the handler's exception must be its own submission's outcome", + failure, e.getCause()); + } + + RecordingHandler next = new RecordingHandler(); + q.sql("INSERT INTO t VALUES (1)").handler(next); + assertOwnOutcome(q); + Assert.assertTrue("the next submission must get its own reply", next.execDone); + Assert.assertEquals("the abandoned query's rows must not reach the next submission", + 0, next.batches.get()); + Assert.assertEquals("the connection must be reused", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testThrowingHandlerLeavesTheNextSubmissionIntact() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 1_000; + Script script = new Script(); + script.onQuery = (s, c, id, n) -> { + if (n == 1) { + // Ends the query halfway through the second grace period: the + // caller is released first, while the worker still drains it. + s.replyOnceLater(c, id, timeoutMs + graceMs + graceMs / 2, + queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, "timeout, query aborted")); + } else { + s.send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, QwpEgressMsgKind.CAP_QUERY_FLAGS | QwpEgressMsgKind.CAP_QUERY_TIMEOUT); + QuestDB db = QuestDB.connect(config(server) + "query_close_timeout_ms=" + graceMs + ";")) { + try (Query q = db.borrowQuery()) { + q.sql("SELECT slow()").handler(new QwpColumnBatchHandler() { + @Override + public void onBatch(QwpColumnBatch batch) { + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + throw new IllegalStateException("handler failed: " + message); + } + }).timeout(timeoutMs, TimeUnit.MILLISECONDS); + assertTimesOut(q); + + // Submitted while the worker still drains the timed-out query: it runs + // once the drain ends and must get its own outcome -- neither the + // aborted query's late reply nor the handler's failure. + q.sql("INSERT INTO t VALUES (1)").handler(new NoopHandler()).timeout(0, TimeUnit.MILLISECONDS); + assertOwnOutcome(q); + } + // The pooled connection is handed to the next borrower in step too. + try (Query q = db.borrowQuery()) { + q.sql("INSERT INTO t VALUES (2)").handler(new NoopHandler()); + assertOwnOutcome(q); + } + Assert.assertEquals("the drained connection must be reused", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + @Test(timeout = 30_000) public void testTimedOutQueryKeepsThePooledConnection() throws Exception { TestUtils.assertMemoryLeak(() -> { @@ -236,6 +344,15 @@ public void testWorkerWhoseConnectionStoppedRespondingIsReplaced() throws Except }); } + private static void assertOwnOutcome(Query q) throws InterruptedException { + try { + q.submit().await(); + } catch (QueryException e) { + Assert.fail("the submission must get its own outcome, got status=" + e.getStatus() + + ", message=" + e.getMessage() + ", cause=" + e.getCause()); + } + } + private static QueryException assertTimesOut(Query q) throws InterruptedException { try { q.submit().await(); @@ -261,13 +378,34 @@ private static byte[] execDone(long requestId) { return frame(QwpEgressMsgKind.EXEC_DONE, requestId, new byte[]{0, 0}); // op_type, rows_affected } + /** + * RESULT_BATCH {@code batch_seq == 0}: the schema (one VARCHAR column) and a + * single row. + */ + private static byte[] firstBatch(long requestId) { + ByteArrayOutputStream body = new ByteArrayOutputStream(); + putVarint(body, 0); // batch_seq + putVarint(body, 0); // table name length + putVarint(body, 1); // row_count + putVarint(body, 1); // column_count + putVarint(body, 1); // column name length + body.write('s'); + body.write(QwpConstants.TYPE_VARCHAR); + putVarcharCell(body, "a"); + return frame(QwpEgressMsgKind.RESULT_BATCH, requestId, 1, body.toByteArray()); + } + private static byte[] frame(byte msgKind, long requestId, byte[] body) { + return frame(msgKind, requestId, 0, body); + } + + private static byte[] frame(byte msgKind, long requestId, int tableCount, byte[] body) { int payloadLen = 1 + 8 + body.length; ByteBuffer bb = ByteBuffer.allocate(QwpConstants.HEADER_SIZE + payloadLen).order(ByteOrder.LITTLE_ENDIAN); bb.putInt(QwpConstants.MAGIC_MESSAGE); bb.put(QwpConstants.VERSION); bb.put((byte) 0); // flags - bb.putShort((short) 0); // table_count + bb.putShort((short) tableCount); bb.putInt(payloadLen); bb.put(msgKind); bb.putLong(requestId); @@ -275,6 +413,36 @@ private static byte[] frame(byte msgKind, long requestId, byte[] body) { return bb.array(); } + /** + * Continuation RESULT_BATCH ({@code batch_seq == 1}): one row against the + * schema of {@link #firstBatch}. + */ + private static byte[] nextBatch(long requestId) { + ByteArrayOutputStream body = new ByteArrayOutputStream(); + putVarint(body, 1); // batch_seq + putVarint(body, 0); // table name length + putVarint(body, 1); // row_count + putVarcharCell(body, "b"); + return frame(QwpEgressMsgKind.RESULT_BATCH, requestId, 1, body.toByteArray()); + } + + private static void putVarcharCell(ByteArrayOutputStream body, String value) { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + body.write(0); // null_flag: no nulls + ByteBuffer offsets = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN); + offsets.putInt(0).putInt(bytes.length); + body.write(offsets.array(), 0, 8); + body.write(bytes, 0, bytes.length); + } + + private static void putVarint(ByteArrayOutputStream out, long value) { + while ((value & ~0x7FL) != 0) { + out.write((int) ((value & 0x7F) | 0x80)); + value >>>= 7; + } + out.write((int) value); + } + private static byte[] queryError(long requestId, byte status, String message) { byte[] msg = message.getBytes(StandardCharsets.UTF_8); ByteBuffer body = ByteBuffer.allocate(1 + 2 + msg.length).order(ByteOrder.LITTLE_ENDIAN); @@ -319,6 +487,29 @@ public void onError(byte status, String message) { } } + private static final class RecordingHandler implements QwpColumnBatchHandler { + final AtomicInteger batches = new AtomicInteger(); + volatile boolean execDone; + + @Override + public void onBatch(QwpColumnBatch batch) { + batches.incrementAndGet(); + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + } + + @Override + public void onExecDone(short opType, long rowsAffected) { + execDone = true; + } + } + /** * Scripted egress server. Like a real one it sends one terminal frame per * request, so a cancel of a query that already ended gets no reply. diff --git a/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java index 9240a1fc9..3fc7a3698 100644 --- a/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java +++ b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java @@ -39,6 +39,7 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.lang.reflect.Field; import java.net.InetAddress; import java.net.ServerSocket; import java.net.Socket; @@ -47,10 +48,12 @@ import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; @@ -64,7 +67,9 @@ * central promise of the feature: a timed-out query ends gracefully -- reported * as {@link QwpConstants#STATUS_QUERY_TIMEOUT} -- on a connection that stays * open and authenticated for the next query. Only a connection that does not - * answer at all is replaced. + * answer at all is replaced. The same holds for a query its caller abandons -- + * the handler throws, or the thread is interrupted -- before the query ends: the + * next query never sees the abandoned one's leftovers. *

    * The server-side enforcement (breaker timeout, status mapping) is covered * against a live server in the questdb repository. @@ -160,6 +165,52 @@ public void testCallerIsReleasedAfterTheGraceWhileTheConnectionDrains() throws E }); } + @Test(timeout = 30_000) + public void testCancelOfAQueryQueuedBehindLeftoversIsNotLost() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final Set received = ConcurrentHashMap.newKeySet(); + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + send(c, firstBatch(id, "a")); + // The rest arrives after the handler has thrown. + s.sendLater(c, 300, nextBatch(id, 1, "b"), resultEnd(id, 2)); + } else { + // Held: only a CANCEL ends it. + received.add(id); + } + }; + // Like a real server, ignores a cancel of a query it has not received. + script.onCancel = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (received.contains(id)) { + s.replyOnce(c, id, 0, queryError(id, QwpConstants.STATUS_CANCELLED, "cancelled by client")); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.withQueryTimeoutGrace(500); + try { + client.execute("SELECT s FROM t", null, failingOnBatch(new IllegalStateException("handler failed")), false, 0); + Assert.fail("the handler's exception must propagate out of execute()"); + } catch (IllegalStateException expected) { + // the first query is abandoned, its rest still on the way + } + + // Cancelled before it is sent: the I/O thread still works through the + // abandoned query, and must send the CANCEL only after this query. + RecordingHandler next = new RecordingHandler(); + client.execute("SELECT s FROM t", binds -> client.cancel(), next, false, 2_000); + Assert.assertEquals("the cancel must reach the server after its query, got message=" + next.errorMessage, + QwpConstants.STATUS_CANCELLED, next.errorStatus); + Assert.assertEquals("the abandoned query's rows must not reach the next query", 0, next.batches.get()); + } finally { + script.close(); + } + }); + } + @Test(timeout = 30_000) public void testCompleteResultJustPastTheDeadlineIsASuccess() throws Exception { TestUtils.assertMemoryLeak(() -> { @@ -219,6 +270,229 @@ public void testFailoverReplayCarriesTheRemainingBudget() throws Exception { }); } + @Test(timeout = 30_000) + public void testInterruptReleasesTheCallerOnlyOnceItsRequestIsEncoded() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> send(c, execDone(requestIdOf(frame))); + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + // The I/O thread picks the request up, and stops inside its SQL text. + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CharSequence sql = new BlockingSql("INSERT INTO t VALUES (1)", reading, release); + RecordingHandler interrupted = new RecordingHandler(); + Thread caller = new Thread(() -> client.execute(sql, null, interrupted, false, 0)); + caller.start(); + try { + Assert.assertTrue("the I/O thread must start encoding the request", reading.await(5, TimeUnit.SECONDS)); + caller.interrupt(); + // The caller may change the SQL text once released, and the next + // query reuses the bind buffer: not before the I/O thread is done. + caller.join(300); + Assert.assertTrue("execute() must wait for the I/O thread to finish reading the request", + caller.isAlive()); + } finally { + release.countDown(); + } + caller.join(5_000); + Assert.assertFalse("execute() must return once the request is encoded", caller.isAlive()); + Assert.assertEquals(WebSocketResponse.STATUS_INTERNAL_ERROR, interrupted.errorStatus); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (2)", null, next, false, 0); + Assert.assertTrue("the next statement must get its own reply", next.execDone); + Assert.assertEquals("the connection must be reused, not replaced", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testInterruptWithAStuckIoThreadGivesUpTheConnection() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> send(c, execDone(requestIdOf(frame))); + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + final long shutdownJoinMs = 300; + setShutdownJoinMs(client, shutdownJoinMs); + // The I/O thread picks the request up, and stays inside its SQL text + // until it is interrupted: stuck. + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CharSequence sql = new BlockingSql("INSERT INTO t VALUES (1)", reading, release); + RecordingHandler interrupted = new RecordingHandler(); + Thread caller = new Thread(() -> client.execute(sql, null, interrupted, false, 0)); + caller.start(); + long elapsed; + try { + Assert.assertTrue("the I/O thread must start encoding the request", reading.await(5, TimeUnit.SECONDS)); + long start = System.nanoTime(); + caller.interrupt(); + caller.join(10_000); + elapsed = elapsedMs(start); + } finally { + release.countDown(); + } + Assert.assertFalse("execute() must return", caller.isAlive()); + Assert.assertEquals(WebSocketResponse.STATUS_INTERNAL_ERROR, interrupted.errorStatus); + Assert.assertTrue("execute() must first wait for the I/O thread, returned after " + elapsed + "ms", + elapsed >= shutdownJoinMs); + Assert.assertTrue("a connection whose I/O thread is stuck must be given up", client.hasTerminalFailure()); + + // failover=on (the default) replaces the connection for the next query. + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (2)", null, next, false, 0); + Assert.assertTrue("the next statement must get its own reply", next.execDone); + Assert.assertEquals("a replacement connection must have been opened", 2, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testInterruptedQueryDoesNotAnswerTheNextQuery() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + // Answers only after its caller has stopped waiting. + s.sendLater(c, 500, firstBatch(id, "late"), resultEnd(id, 1)); + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + RecordingHandler first = new RecordingHandler(); + Thread caller = new Thread(() -> client.execute("SELECT s FROM t", null, first, false, 0)); + caller.start(); + long firstId = requestIdOf(script.queries.take()); + caller.interrupt(); + caller.join(5_000); + Assert.assertFalse("execute() must return when its thread is interrupted", caller.isAlive()); + Assert.assertEquals(WebSocketResponse.STATUS_INTERNAL_ERROR, first.errorStatus); + byte[] cancel = script.cancels.poll(5, TimeUnit.SECONDS); + Assert.assertNotNull("the abandoned query must be cancelled", cancel); + Assert.assertEquals(firstId, requestIdOf(cancel)); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue("the next statement must get its own reply, got status=" + next.errorStatus + + ", message=" + next.errorMessage, next.execDone); + Assert.assertEquals("the abandoned query's rows must not reach the next statement", 0, next.batches.get()); + Assert.assertEquals("the connection must be reused, not replaced", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testInterruptedQueryNotYetSentIsWithdrawn() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + send(c, firstBatch(id, "a")); + // The rest arrives late; until then the I/O thread does not pick up + // the next request. + s.sendLater(c, 500, nextBatch(id, 1, "b"), resultEnd(id, 2)); + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + try { + client.execute("SELECT s FROM t", null, failingOnBatch(new IllegalStateException("handler failed")), false, 0); + Assert.fail("the handler's exception must propagate out of execute()"); + } catch (IllegalStateException expected) { + // the first query is abandoned, its rest still on the way + } + + // Queued behind the abandoned query's leftovers, and interrupted there. + RecordingHandler interrupted = new RecordingHandler(); + Thread caller = new Thread(() -> client.execute("INSERT INTO t VALUES (2)", null, interrupted, false, 0)); + caller.start(); + awaitParked(caller); + caller.interrupt(); + caller.join(5_000); + Assert.assertFalse("execute() must return when its thread is interrupted", caller.isAlive()); + Assert.assertEquals(WebSocketResponse.STATUS_INTERNAL_ERROR, interrupted.errorStatus); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (3)", null, next, false, 0); + Assert.assertTrue("the next statement must get its own reply", next.execDone); + RecordingHandler last = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (4)", null, last, false, 0); + Assert.assertTrue("the last statement must get its own reply", last.execDone); + + // The interrupted statement, withdrawn before it was sent, never reaches + // the server; every other statement reaches it exactly once. + List received = new ArrayList<>(); + for (byte[] frame : script.queries) { + received.add(sqlOf(frame)); + } + Assert.assertEquals(Arrays.asList("SELECT s FROM t", "INSERT INTO t VALUES (3)", "INSERT INTO t VALUES (4)"), + received); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testLeftoversOfAnAbandonedQueryDoNotPostponeTheNextTimeout() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 200; + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + send(c, firstBatch(id, "a")); + // Streams on, ignoring the CANCEL, for 3 seconds: far past the + // next query's deadline and grace period. + for (int i = 1; i <= 150; i++) { + s.sendLater(c, 20L * i, nextBatch(id, i, "b")); + } + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.withQueryTimeoutGrace(graceMs); + try { + client.execute("SELECT s FROM t", null, failingOnBatch(new IllegalStateException("handler failed")), false, 0); + Assert.fail("the handler's exception must propagate out of execute()"); + } catch (IllegalStateException expected) { + // the first query is abandoned and keeps streaming + } + + RecordingHandler next = new RecordingHandler(); + long start = System.nanoTime(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, timeoutMs); + long elapsed = elapsedMs(start); + + next.assertTimedOut(); + Assert.assertEquals("the abandoned query's rows must not reach the next statement", 0, next.batches.get()); + Assert.assertTrue("the leftovers must not postpone the timeout, took " + elapsed + "ms", + elapsed < 1_500); + Assert.assertTrue("a connection still busy with the abandoned query is given up", + client.hasTerminalFailure()); + } finally { + script.close(); + } + }); + } + @Test(timeout = 30_000) public void testNoTimeoutFieldWithoutTheCapabilityOrTimeout() throws Exception { TestUtils.assertMemoryLeak(() -> { @@ -363,6 +637,159 @@ public void testServerReportedTimeoutKeepsTheConnectionOpen() throws Exception { }); } + @Test(timeout = 30_000) + public void testThrowingBatchHandlerDoesNotLeakRowsIntoTheNextQuery() throws Exception { + TestUtils.assertMemoryLeak(() -> { + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + send(c, firstBatch(id, "a")); + // The rest of the result arrives after the handler has thrown. + s.sendLater(c, 300, nextBatch(id, 1, "b"), resultEnd(id, 2)); + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + final IllegalStateException failure = new IllegalStateException("handler failed"); + try { + client.execute("SELECT s FROM t", null, failingOnBatch(failure), false, 0); + Assert.fail("the handler's exception must propagate out of execute()"); + } catch (IllegalStateException e) { + Assert.assertSame(failure, e); + } + long firstId = requestIdOf(script.queries.take()); + byte[] cancel = script.cancels.poll(5, TimeUnit.SECONDS); + Assert.assertNotNull("the abandoned query must be cancelled", cancel); + Assert.assertEquals(firstId, requestIdOf(cancel)); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertEquals("the abandoned query's rows must not reach the next statement", 0, next.batches.get()); + Assert.assertTrue("the next statement must get its own reply, got status=" + next.errorStatus + + ", message=" + next.errorMessage, next.execDone); + Assert.assertFalse(next.ended); + Assert.assertEquals("the connection must be reused, not replaced", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testThrowingHandlerIsRethrownWhenTheConnectionIsGivenUp() throws Exception { + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 200; + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + if (n > 1) { + send(c, execDone(requestIdOf(frame))); + } + // n == 1: never answered, not even after a CANCEL. + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.withQueryTimeoutGrace(graceMs); + // An Error, not an exception: the handler's throwable must come back as is. + final Error failure = new Error("handler failed"); + QwpColumnBatchHandler throwing = new QwpColumnBatchHandler() { + @Override + public void onBatch(QwpColumnBatch batch) { + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + throw failure; + } + }; + Throwable thrown = null; + long start = System.nanoTime(); + try { + client.execute("SELECT slow()", null, throwing, false, timeoutMs); + } catch (Throwable t) { + thrown = t; + } + long elapsed = elapsedMs(start); + + Assert.assertSame("the handler's throwable must propagate unwrapped", failure, thrown); + Assert.assertTrue("execute() must give up on the connection before rethrowing, took " + elapsed + "ms", + elapsed >= timeoutMs + 2 * graceMs); + Assert.assertTrue("a silent connection must be marked failed", client.hasTerminalFailure()); + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue(next.execDone); + Assert.assertEquals("a replacement connection must have been opened", 2, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + + @Test(timeout = 30_000) + public void testThrowingHandlerStillDrainsTheTimedOutQuery() throws Exception { + // The handler contract lets any callback throw and keeps the connection + // usable. At the end of the grace period onError fires while the aborted + // query is still running, so a throw there must not skip the drain: + // otherwise the aborted query's late reply answers the next query. + TestUtils.assertMemoryLeak(() -> { + final long timeoutMs = 100; + final long graceMs = 1_000; + ScriptedServer script = new ScriptedServer(); + script.onQuery = (s, c, frame, n) -> { + long id = requestIdOf(frame); + if (n == 1) { + // Ends the query halfway through the second grace period: after the + // handler is told about the timeout, before the connection is given up. + s.sendLater(c, timeoutMs + graceMs + graceMs / 2, + queryError(id, QwpConstants.STATUS_QUERY_TIMEOUT, "timeout, query aborted")); + } else { + send(c, execDone(id)); + } + }; + try (TestWebSocketServer server = startServer(script, CAPS_WITH_TIMEOUT); + QwpQueryClient client = connect(server, "")) { + client.withQueryTimeoutGrace(graceMs); + QwpColumnBatchHandler throwing = new QwpColumnBatchHandler() { + @Override + public void onBatch(QwpColumnBatch batch) { + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + throw new IllegalStateException(message); + } + }; + try { + client.execute("SELECT slow()", null, throwing, false, timeoutMs); + Assert.fail("the handler's exception must propagate out of execute()"); + } catch (IllegalStateException e) { + Assert.assertTrue("the handler must throw at the end of the grace period, while the query " + + "is still running: " + e.getMessage(), e.getMessage().contains("grace period")); + } + + RecordingHandler next = new RecordingHandler(); + client.execute("INSERT INTO t VALUES (1)", null, next, false, 0); + Assert.assertTrue("the next statement must get its own reply, not the aborted query's: status=" + + next.errorStatus + ", message=" + next.errorMessage, next.execDone); + Assert.assertEquals("the drained connection must be reused", 1, server.handshakeCount()); + } finally { + script.close(); + } + }); + } + @Test(timeout = 30_000) public void testTimeoutFieldCarriesTheBudget() throws Exception { TestUtils.assertMemoryLeak(() -> { @@ -512,6 +939,18 @@ public void testUserCancelBeforeTheDeadlineIsReportedAsCancel() throws Exception }); } + /** + * Waits until {@code thread} blocks, here: parked waiting for its query's + * events, with the request queued. + */ + private static void awaitParked(Thread thread) throws InterruptedException { + final long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (thread.getState() != Thread.State.WAITING && thread.getState() != Thread.State.TIMED_WAITING) { + Assert.assertTrue("the thread must block", System.nanoTime() - deadlineNanos < 0); + Thread.sleep(1); + } + } + private static byte[] closeFrame() { // Marker understood by ScriptedServer's send(): close the connection. return new byte[0]; @@ -537,6 +976,23 @@ private static byte[] execDone(long requestId) { return serverFrame(QwpEgressMsgKind.EXEC_DONE, requestId, 0, new byte[]{0, 0}); // op_type, rows_affected } + private static QwpColumnBatchHandler failingOnBatch(RuntimeException failure) { + return new QwpColumnBatchHandler() { + @Override + public void onBatch(QwpColumnBatch batch) { + throw failure; + } + + @Override + public void onEnd(long totalRows) { + } + + @Override + public void onError(byte status, String message) { + } + }; + } + /** * RESULT_BATCH {@code batch_seq == 0}: the schema (one VARCHAR column) and a * single row. @@ -641,6 +1097,21 @@ private static byte[] serverFrame(byte msgKind, long requestId, int tableCount, return bb.array(); } + private static void setShutdownJoinMs(QwpQueryClient client, long millis) throws Exception { + Field field = QwpQueryClient.class.getDeclaredField("shutdownJoinMs"); + field.setAccessible(true); + field.setLong(client, millis); + } + + /** + * The SQL text of a captured {@code QUERY_REQUEST}. + */ + private static String sqlOf(byte[] queryRequest) { + int[] p = {1 + 8}; + int len = (int) readVarint(queryRequest, p); + return new String(queryRequest, p[0], len, StandardCharsets.UTF_8); + } + private static TestWebSocketServer startServer(ScriptedServer script, int capabilities) throws Exception { TestWebSocketServer server = new TestWebSocketServer(script); server.setSendServerInfo(true); @@ -726,6 +1197,51 @@ int port() { } } + /** + * SQL text that holds up whoever reads its first character -- the I/O + * thread, which reads it while encoding the request -- until released or + * interrupted. + */ + private static final class BlockingSql implements CharSequence { + private final CountDownLatch reading; + private final CountDownLatch release; + private final String text; + + BlockingSql(String text, CountDownLatch reading, CountDownLatch release) { + this.text = text; + this.reading = reading; + this.release = release; + } + + @Override + public char charAt(int index) { + if (index == 0) { + reading.countDown(); + try { + release.await(10, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + return text.charAt(index); + } + + @Override + public int length() { + return text.length(); + } + + @Override + public CharSequence subSequence(int start, int end) { + return text.subSequence(start, end); + } + + @Override + public String toString() { + return text; + } + } + private static final class RecordingHandler implements QwpColumnBatchHandler { final AtomicInteger batches = new AtomicInteger(); final AtomicInteger failoverResets = new AtomicInteger();