[pull] master from cube-js:master

packages/cubejs-base-driver/src/queue-driver.interface.ts

@@ -1,8 +1,6 @@
 export type QueryDef = any;
 // Primary key of Queue item
 export type QueueId = string | number | bigint;
-// The lock token of a retrieval, always the item's queueId. Only the memory driver compares it.
-export type ProcessingId = string | number | bigint;
 export type QueryKey = (string | [string, any[]]) & {
   persistent?: true,
 };
@@ -13,21 +11,21 @@ export type GetActiveAndToProcessResponse = [active: QueryKeysTuple[], toProcess
 export type QueryStageStateResponse = [active: string[], toProcess: string[]] | [active: string[], toProcess: string[], defs: Record<string, QueryDef>];
 export type RetrieveForProcessingSuccess = [
   added: unknown,
-  // QueueId is required for Cube Store, other providers don't support it
+  // Identifies the retrieved generation of the queue item.
   queueId: QueueId | null,
   active: QueryKeyHash[],
   pending: number,
   def: QueryDef,
-  lockAquired: true
+  retrieved: true
 ];
 export type RetrieveForProcessingFail = [
   added: unknown,
-  // QueueId is required for Cube Store, other providers don't support it
+  // Null when no queue item was retrieved.
   queueId: QueueId | null,
   active: QueryKeyHash[],
   pending: number,
   def: null,
-  lockAquired: false
+  retrieved: false
 ];
 export type RetrieveForProcessingResponse = RetrieveForProcessingSuccess | RetrieveForProcessingFail | null;
 export type AddToQueueResponse = [
@@ -106,14 +104,13 @@ export interface QueueDriverConnectionInterface {
   getStalledQueries(): Promise<QueryKeysTuple[]>;
   getQueryStageState(onlyKeys: boolean): Promise<QueryStageStateResponse>;
   updateHeartBeat(hash: QueryKeyHash, queueId: QueueId | null): Promise<void>;
-  // Trying to acquire a lock for processing a queue item, this method can return null when
-  // multiple nodes tries to process the same query
-  retrieveForProcessing(hash: QueryKeyHash, processingId: ProcessingId): Promise<RetrieveForProcessingResponse>;
-  freeProcessingLock(hash: QueryKeyHash, processingId: ProcessingId, activated: unknown): Promise<void>;
-  optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, processingId: ProcessingId, queueId: QueueId | null): Promise<boolean>;
+  // Atomically moves a queue item to active. Returns null when another node is already
+  // processing the query or the concurrency budget is full.
+  retrieveForProcessing(hash: QueryKeyHash, queueId: QueueId): Promise<RetrieveForProcessingResponse>;
+  optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, queueId: QueueId): Promise<boolean>;
   cancelQuery(queryKey: QueryKey, queueId: QueueId | null): Promise<QueryDef | null>;
   getQueryAndRemove(hash: QueryKeyHash, queueId: QueueId | null): Promise<[QueryDef]>;
-  setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, processingId: ProcessingId, queueId: QueueId | null): Promise<unknown>;
+  setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, queueId: QueueId): Promise<unknown>;
   release(): void;
   //
   getQueriesToCancel(): Promise<QueryKeysTuple[]>

packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts

@@ -12,7 +12,6 @@ import {
   AddToQueueResponse,
   QueryKey,
   QueryKeyHash,
-  ProcessingId,
   QueueId,
   GetActiveAndToProcessResponse,
   QueryKeysTuple,
@@ -182,10 +181,6 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte
     return null;
   }
 
-  public async freeProcessingLock(_hash: QueryKeyHash, _processingId: string, _activated: unknown): Promise<void> {
-    // nothing to do
-  }
-
   public async getActiveQueries(): Promise<QueryKeysTuple[]> {
     const rows = await this.driver.query<CubeStoreListResponse>('QUEUE ACTIVE ?', [
       this.options.redisQueuePrefix
@@ -336,7 +331,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte
     return null;
   }
 
-  public async optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, _processingId: ProcessingId, queueId: QueueId): Promise<boolean> {
+  public async optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, queueId: QueueId): Promise<boolean> {
     await this.driver.query('QUEUE MERGE_EXTRA ? ?', [
       // queryKeyHash as compatibility fallback
       queueId || this.prefixKey(hash),
@@ -372,7 +367,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte
     ];
   }
 
-  public async retrieveForProcessing(hash: QueryKeyHash, _processingId: string): Promise<RetrieveForProcessingResponse> {
+  public async retrieveForProcessing(hash: QueryKeyHash, _queueId: QueueId): Promise<RetrieveForProcessingResponse> {
     const rows = await this.driver.query<CubeStoreRetrieveResponse>('QUEUE RETRIEVE EXTENDED CONCURRENCY ? ?', [
       this.options.concurrency,
       this.prefixKey(hash),
@@ -404,7 +399,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte
     return null;
   }
 
-  public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, _processingId: ProcessingId, queueId: QueueId): Promise<boolean> {
+  public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, queueId: QueueId): Promise<boolean> {
     const rows = await this.driver.query('QUEUE ACK ? ?', [
       // queryKeyHash as compatibility fallback
       queueId || this.prefixKey(hash),

packages/cubejs-query-orchestrator/DEVELOPMENT.md

@@ -182,12 +182,13 @@ sequenceDiagram
 ## Background execution: `processQuery` → `executeQuery`
 
 `processQuery` retrieves the item and nothing else. What it hands to `sendProcessMessageFn` is a
-`RetrievedQuery` — `{ queryKeyHash, queueId, processingId, queueSize, query }`, plain data on
+`RetrievedQuery` — `{ queryKeyHash, queueId, queueSize, query }`, plain data on
 purpose, so a custom implementation can serialize it and let another process run
 `executeQuery`. The default implementation calls `executeQuery` in-process.
 
-`processingId` is the lock token of the retrieval and always carries the `queueId`. Only the
-memory driver compares it, Cube Store ignores it and keys every command off the `queueId`.
+`queueId` identifies the specific generation of a queue item. The memory driver compares it
+with the active entry before updating or acknowledging a query, while Cube Store keys those
+commands directly off the `queueId`.
 
 `sendProcessMessageFn` must resolve once the hand-off is done, **not** once the query is
 executed: reconcile awaits it, and `executeQuery` ends with `reconcileQueue`, which is
@@ -200,8 +201,8 @@ Two consequences of retrieving before the hand-off:
   `@<processUid>` suffix matches, so `sendProcessMessageFn` is always called on the owning
   process for them — it just must not route them elsewhere.
 - A retrieved item is already active. If a custom hand-off loses the message, the item is only
-  recovered by the stalled-heartbeat / `TO_CANCEL` path; `freeProcessingLock` is a no-op on
-  Cube Store, so the retrieval cannot be cheaply undone.
+  recovered by the stalled-heartbeat / `TO_CANCEL` path; a successful retrieval is not undone
+  when the hand-off fails.
 
 A stream query is dispatched while `executeInQueue` is still running, so `waitForQueryStream`
 subscribes to `streamStarted` *before* the dispatch — a handler which starts fast would
@@ -222,10 +223,10 @@ sequenceDiagram
     QueryQueue->>QueueDriver: retrieveForProcessing
     QueueDriver->>CubeStore: QUEUE RETRIEVE EXTENDED CONCURRENCY ?n ?path
     CubeStore-->>QueueDriver: RetrieveResponse
-    QueueDriver-->>QueryQueue: [added, queueId, activeKeys, queueSize, def, lockAcquired]
+    QueueDriver-->>QueryQueue: [added, queueId, activeKeys, queueSize, def, retrieved]
     Note over QueueDriver,CubeStore: The retrieval is atomic in Cube Store:<br/>only one node moves the item to active
 
-    alt def && added && activeKeys includes our key && lockAcquired
+    alt def && added && activeKeys includes our key && retrieved
         QueryQueue-)Background: sendProcessMessageFn(RetrievedQuery)
         Note over QueryQueue,Background: Detached from here on: the hand-off returns,<br/>the execution keeps running
 
@@ -254,8 +255,7 @@ sequenceDiagram
         Background->>Background: reconcileQueue
         Note over Background: The freed concurrency slot is<br/>immediately given to the next query
     else the retrieval did not succeed
-        QueryQueue->>QueueDriver: freeProcessingLock
-        Note over QueryQueue,QueueDriver: Another node is running it, or the<br/>concurrency budget is full. No-op for Cube Store
+        Note over QueryQueue,QueueDriver: Another node is running it, or the<br/>concurrency budget is full. Queue state is unchanged
     end
 ```
 

packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts

@@ -4,7 +4,6 @@ import {
   QueryKey,
   QueryKeyHash,
   QueueId,
-  ProcessingId,
   QueryDef,
   AddToQueueQuery,
   AddToQueueOptions,
@@ -54,8 +53,6 @@ export class LocalQueueDriverConnectionState {
   public active: Record<QueryKeyHash, QueueItem> = {};
 
   public heartBeat: Record<QueryKeyHash, QueueItem> = {};
-
-  public processingLocks: Record<QueryKeyHash, any> = {};
 }
 
 export class LocalQueueDriverConnection implements QueueDriverConnectionInterface {
@@ -223,7 +220,6 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac
     delete this.state.toProcess[queryKeyHash];
     delete this.state.recent[queryKeyHash];
     delete this.state.queryDef[queryKeyHash];
-    delete this.state.processingLocks[queryKeyHash];
 
     return [query];
   }
@@ -233,8 +229,8 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac
     return query;
   }
 
-  public async setResultAndRemoveQuery(queryKeyHash: QueryKeyHash, executionResult: any, processingId: ProcessingId, _queueId?: QueueId | null): Promise<boolean> {
-    if (this.state.processingLocks[queryKeyHash] !== processingId) {
+  public async setResultAndRemoveQuery(queryKeyHash: QueryKeyHash, executionResult: any, queueId: QueueId): Promise<boolean> {
+    if (this.state.active[queryKeyHash]?.queueId !== queueId) {
       return false;
     }
 
@@ -245,7 +241,6 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac
     delete this.state.toProcess[queryKeyHash];
     delete this.state.recent[queryKeyHash];
     delete this.state.queryDef[queryKeyHash];
-    delete this.state.processingLocks[queryKeyHash];
 
     promise.resolved = true;
     if (promise.resolve) {
@@ -277,48 +272,44 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac
     }
   }
 
-  public async retrieveForProcessing(queryKeyHash: QueryKeyHash, processingId: ProcessingId): Promise<RetrieveForProcessingResponse> {
-    let lockAcquired = false;
-
-    if (!this.state.processingLocks[queryKeyHash]) {
-      this.state.processingLocks[queryKeyHash] = processingId;
-      lockAcquired = true;
-    } else {
-      return null;
+  public async retrieveForProcessing(queryKeyHash: QueryKeyHash, queueId: QueueId): Promise<RetrieveForProcessingResponse> {
+    const query = this.state.queryDef[queryKeyHash];
+    const activeKeys = this.queueArray(this.state.active) as QueryKeyHash[];
+
+    if (
+      !query ||
+      query.queueId !== queueId ||
+      this.state.toProcess[queryKeyHash]?.queueId !== queueId ||
+      this.state.active[queryKeyHash] ||
+      activeKeys.length >= this.concurrency
+    ) {
+      return [
+        0,
+        null,
+        activeKeys,
+        Object.keys(this.state.toProcess).length,
+        null,
+        false
+      ];
     }
 
-    let added = 0;
-
-    if (Object.keys(this.state.active).length < this.concurrency && !this.state.active[queryKeyHash]) {
-      this.state.active[queryKeyHash] = { key: queryKeyHash, order: Number(processingId), queueId: processingId };
-      delete this.state.toProcess[queryKeyHash];
-
-      added = 1;
-    }
+    this.state.active[queryKeyHash] = { key: queryKeyHash, order: Number(queueId), queueId };
+    delete this.state.toProcess[queryKeyHash];
 
-    this.state.heartBeat[queryKeyHash] = { key: queryKeyHash, order: new Date().getTime(), queueId: processingId };
+    this.state.heartBeat[queryKeyHash] = { key: queryKeyHash, order: new Date().getTime(), queueId };
 
     return [
-      added,
-      this.state.queryDef[queryKeyHash]?.queueId ?? null,
+      1,
+      query.queueId,
       this.queueArray(this.state.active) as QueryKeyHash[],
       Object.keys(this.state.toProcess).length,
-      this.state.queryDef[queryKeyHash],
-      lockAcquired
+      query,
+      true
     ];
   }
 
-  public async freeProcessingLock(queryKeyHash: QueryKeyHash, processingId: ProcessingId, activated: any): Promise<void> {
-    if (this.state.processingLocks[queryKeyHash] === processingId) {
-      delete this.state.processingLocks[queryKeyHash];
-      if (activated) {
-        delete this.state.active[queryKeyHash];
-      }
-    }
-  }
-
-  public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: any, processingId: ProcessingId, _queueId?: QueueId | null): Promise<boolean> {
-    if (this.state.processingLocks[queryKeyHash] !== processingId) {
+  public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: any, queueId: QueueId): Promise<boolean> {
+    if (this.state.active[queryKeyHash]?.queueId !== queueId || !this.state.queryDef[queryKeyHash]) {
       return false;
     }