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/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 700e219ea..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,7 +114,15 @@ 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 + // 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) { @@ -293,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); @@ -356,6 +367,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 +397,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 +409,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 +419,7 @@ public void submitQuery( pendingRequest.bindPayloadPtr = bindPayloadPtr; pendingRequest.bindPayloadLen = bindPayloadLen; pendingRequest.queryFlags = queryFlags; + pendingRequest.timeoutMs = timeoutMs; requests.put(pendingRequest); } @@ -404,6 +430,28 @@ 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 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 @@ -421,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)); } /** @@ -461,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)); } /** @@ -518,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)); } /** @@ -526,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); } } @@ -628,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 @@ -721,7 +780,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(); } @@ -771,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; @@ -783,5 +849,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..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 @@ -46,8 +46,10 @@ 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; +import java.util.concurrent.locks.LockSupport; /** * QWP egress (query results) client. @@ -90,11 +92,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 +116,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 +176,22 @@ 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; + // 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; + // 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 +225,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 +308,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 +390,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 +468,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 +529,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 +587,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 +886,7 @@ public synchronized void connect() { hostTracker.recordSuccess(i); currentEndpointIndex = i; connected = true; + reconnectPending = false; return; } if (lastObservedMismatch != null) { @@ -904,6 +971,55 @@ 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. 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 + * @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 +1031,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 +1078,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 +1142,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 +1170,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 +1487,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 +1621,22 @@ 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; + } + @SuppressWarnings("BooleanMethodIsAlwaysInverted") private static boolean matchesTarget(byte role, String target) { if (TARGET_ANY.equals(target)) { @@ -1452,6 +1653,107 @@ 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; + } + + /** + * 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 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 + * 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 abandonConnection(QwpEgressIoThread io, String failureMessage) { + GenerationListener listener = currentGenerationListener; + if (listener != null) { + listener.onTerminalFailure(WebSocketResponse.STATUS_INTERNAL_ERROR, failureMessage); + } + 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) { + discardEvent(io, ev); + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + /** + * 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 + * 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 +1851,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 +1875,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 +1914,11 @@ 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); + // 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; } @@ -1605,6 +1926,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 +1948,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 +1979,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 +1999,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 +2014,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 +2028,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 +2044,37 @@ 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. + *

    + * 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, 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 +2108,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 +2138,144 @@ 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; + 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(), - resolveQueryFlags(resetSymbolDict)); + queryFlags, wireTimeoutMs); + submitted = true; while (true) { - QueryEvent ev = io.takeEvent(); + 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) { + 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.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: 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 { @@ -1773,9 +2288,23 @@ private void executeOnce(CharSequence sql, QwpBindSetter binds, FailoverProbeHan } } 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. - 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 { + 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; } } @@ -1823,7 +2352,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 +2376,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 +2432,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 +2482,7 @@ private void reconnectViaTracker() { hostTracker.recordSuccess(i); currentEndpointIndex = i; connected = true; + reconnectPending = false; return; } if (lastMismatch != null) { @@ -1930,6 +2495,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 +2540,70 @@ 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; + } + + /** + * 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) { @@ -1977,7 +2615,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) { @@ -2042,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; @@ -2120,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/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..c15130015 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,26 @@ 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; + // 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(). + 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 +185,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 +248,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,11 +282,17 @@ 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(); try { done = false; + submissionSeq++; resultStatus = 0; resultMessage = null; unexpectedError = null; @@ -238,7 +302,19 @@ 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 + // 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); @@ -311,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) { @@ -333,6 +432,8 @@ void resetForBorrow() { userBinds = null; userHandler = null; sqlBuffer.clear(); + timeoutMillis = -1; + cancelRequested = false; resultStatus = 0; resultMessage = null; unexpectedError = null; @@ -340,12 +441,30 @@ 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 + // 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). + try { + client.execute(sqlBuffer, wireBinds, wrappingHandler, false, timeoutArg); + } catch (Throwable t) { + signalUnexpected(submission, t); + } } /** @@ -353,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/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..d034922e0 --- /dev/null +++ b/core/src/test/java/io/questdb/client/test/QueryTimeoutFacadeTest.java @@ -0,0 +1,582 @@ +/*+***************************************************************************** + * ___ _ ____ ____ + * / _ \ _ _ ___ ___| |_| _ \| __ ) + * | | | | | | |/ _ \/ __| __| | | | _ \ + * | |_| | |_| | __/\__ \ |_| |_| | |_) | + * \__\_\\__,_|\___||___/\__|____/|____/ + * + * 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.ByteArrayOutputStream; +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. A handler that throws -- on the + * timeout, or mid-result -- fails only its own submission. + */ +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 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(() -> { + 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 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(); + } 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 + } + + /** + * 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) tableCount); + bb.putInt(payloadLen); + bb.put(msgKind); + bb.putLong(requestId); + bb.put(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); + 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) { + } + } + + 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. + */ + 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..3fc7a3698 --- /dev/null +++ b/core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpQueryClientQueryTimeoutTest.java @@ -0,0 +1,1361 @@ +/*+***************************************************************************** + * ___ _ ____ ____ + * / _ \ _ _ ___ ___| |_| _ \| __ ) + * | | | | | | |/ _ \/ __| __| | | | _ \ + * | |_| | |_| | __/\__ \ |_| |_| | |_) | + * \__\_\\__,_|\___||___/\__|____/|____/ + * + * 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.lang.reflect.Field; +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.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; +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 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. + */ +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 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(() -> { + 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 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(() -> { + 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 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(() -> { + 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(); + } + }); + } + + /** + * 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]; + } + + 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 + } + + 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. + */ + 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 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); + 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(); + } + } + + /** + * 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(); + 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) { + } + } +}