diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts index 7268a8bc64572..f70eca0c544c9 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts @@ -9,6 +9,7 @@ import { QueryStageStateResponse, AddToQueueOptions, QueuePriority, + QueueDriverConnectionInterface, RetrieveForProcessingSuccess } from '@cubejs-backend/base-driver'; import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver'; @@ -854,6 +855,9 @@ export class QueryQueue { try { let executionResult; let queryExecutionFinished = false; + // Set once the queue item is known to be gone, so that the rejection the cancellation + // causes isn't reported as a query error. + let queryCancelled = false; // Set by the query handler's setCancelHandler callback once execution begins. // Not available on the original query def from retrieveForProcessing. let localCancelHandler: unknown = null; @@ -902,6 +906,7 @@ export class QueryQueue { try { const currentDef = await queueConnection.getQueryDef(queryKeyHashed, queueId); if (!currentDef && !queryExecutionFinished) { + queryCancelled = true; this.logger('Cancelling query due to external cancellation', { queueId, queryKey: query.queryKey, @@ -996,7 +1001,17 @@ export class QueryQueue { executionResult = { error: (e.message || e).toString() // TODO error handling }; - this.logger('Error while querying', { + + // The queue removes an item when it cancels it - orphaned, stalled or explicit - and the + // driver then rejects the execution still in flight. Reporting that rejection as a query + // failure surfaces the cancellation in query history, which orphaned queries are kept out + // of. The lookup happens after the rejection, so a query which genuinely fails as its item + // is orphaned is reclassified too: harmless, nothing is waiting for either result. A + // timeout is excluded - it cancels a query which is still queued, so it stays an error. + const cancelled = !(e instanceof TimeoutError) && + (queryCancelled || await this.isQueryRemovedFromQueue(queueConnection, query, queryKeyHashed, queueId)); + + const logEvent = { queueId, queueSize, duration: ((new Date()).getTime() - startQueryTime), @@ -1009,8 +1024,24 @@ export class QueryQueue { newVersionEntry: query.query?.newVersionEntry, preAggregation: query.query?.preAggregation, addedToQueueTime: query.addedToQueueTime, - error: (e.stack || e).toString() - }); + }; + + if (cancelled) { + // `warning` rather than `error` on purpose: an `error` field is what marks a query as + // failed downstream, and a cancellation is not a failure. It still has to be one of the + // two, otherwise the default logger drops the event at the default `info` level. + this.logger('Cancelled query execution', { + ...logEvent, + warning: 'Query execution was rejected because the query had been cancelled', + cancellationError: (e.stack || e).toString() + }); + } else { + this.logger('Error while querying', { + ...logEvent, + error: (e.stack || e).toString() + }); + } + if (e instanceof TimeoutError) { const queryWithCancelHandle = await queueConnection.getQueryDef(queryKeyHashed, queueId); if (queryWithCancelHandle) { @@ -1063,6 +1094,32 @@ export class QueryQueue { } } + /** + * Whether the queue item is gone, which means a cancellation - orphaned, stalled or explicit - + * has removed it while its execution was still in flight. A queue lookup failure is reported as + * `false` so that a genuine query error is never hidden by a queue storage issue. + */ + protected async isQueryRemovedFromQueue( + queueConnection: QueueDriverConnectionInterface, + query: QueryDef, + queryKeyHashed: QueryKeyHash, + queueId: QueueId + ): Promise { + try { + return !(await queueConnection.getQueryDef(queryKeyHashed, queueId)); + } catch (e: any) { + this.logger('Error while checking query cancellation', { + queueId, + queryKey: query.queryKey, + queuePrefix: this.redisQueuePrefix, + requestId: query.requestId, + error: (e.stack || e).toString() + }); + + return false; + } + } + /** * Processing cancel query flow. */ diff --git a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts index 72814fd4b71d2..8bccf4a65d62c 100644 --- a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts +++ b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts @@ -42,6 +42,9 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => const processMessagePromises: Promise[] = []; const processCancelPromises: Promise[] = []; let cancelledQuery; + // Rejects of the in-flight `cancelable` queries, so that a cancellation can reject the + // running handler the way a driver rejects a query it has stopped. + let cancelableRejects: ((error: Error) => void)[] = []; let streamCallOrder: string[] = []; const tenantPrefix = crypto.randomBytes(6).toString('hex'); @@ -55,6 +58,17 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => await setCancelHandler(result); return delayFn(result, query.delay); }, + cancelable: async (query, setCancelHandler) => { + await setCancelHandler(query.result); + + return new Promise((resolve, reject) => { + const timer = setTimeout(() => resolve(query.result), query.delay); + cancelableRejects.push((error) => { + clearTimeout(timer); + reject(error); + }); + }); + }, }, streamHandler: async (query, stream) => { streamCount++; @@ -82,6 +96,10 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => delay: async (query) => { console.log(`cancel call: ${JSON.stringify(query)}`); cancelledQuery = query.queryKey; + }, + cancelable: async (query) => { + cancelledQuery = query.queryKey; + cancelableRejects.splice(0).forEach((reject) => reject(new Error('Query was cancelled'))); } }, continueWaitTimeout: 1, @@ -111,6 +129,7 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => streamCount = 0; streamHandlerDelay = 250; streamCallOrder = []; + cancelableRejects = []; }); afterAll(async () => { @@ -187,6 +206,38 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => expect(logger.mock.calls[3][0]).toEqual('Error while querying'); }); + test('cancelled query is not reported as an error', async () => { + cancelledQuery = null; + + const queryKey: QueryKey = ['select * from cancelled', []]; + // The client gives up on ContinueWaitError long before the handler would resolve + const pending = queue + .executeInQueue('cancelable', queryKey, { delay: 60 * 1000, result: '1' }, QueuePriority.Background) + .catch(e => e); + + // executionTimeout is 2s, so the cancellation has to reach a handler which is already + // running, otherwise the query fails with a timeout instead. The deadline stays well under + // it so that a worker which never picks the query up fails here rather than as a confusing + // assertion on the log events below. + const deadline = Date.now() + 750; + while (cancelableRejects.length === 0 && Date.now() < deadline) { + await pausePromise(10); + } + expect(cancelableRejects.length).toEqual(1); + + await queue.cancelQuery(queue.redisHash(queryKey), null); + expect(cancelledQuery).toEqual(queryKey); + expect(await pending).toBeInstanceOf(ContinueWaitError); + await awaitProcessing(); + + // The rejection the cancellation causes is a cancellation, not a query failure: reporting + // it as one would surface it in query history + const events = logger.mock.calls.map(([message]) => message); + expect(events).toContain('Cancelling query manual'); + expect(events).toContain('Cancelled query execution'); + expect(events).not.toContain('Error while querying'); + }); + test('stage reporting', async () => { const resultPromise = queue.executeInQueue('delay', '1', { delay: 200, result: '1' }, QueuePriority.Background, { stageQueryKey: '1',