feat(cubestore-driver): Queue - protect race conditions (#6598)

Author: ovr

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

@@ -1,4 +1,8 @@
 export type QueryDef = unknown;
+// Primary key of Queue item
+export type QueueId = string | number | bigint;
+// This was used as lock for Redis, deprecated.
+export type ProcessingId = string | number;
 export type QueryKey = (string | [string, any[]]) & {
   persistent?: true,
 };
@@ -7,19 +11,21 @@ export interface QueryKeyHash extends String {
 }
 
 export type GetActiveAndToProcessResponse = [active: string[], toProcess: string[]];
-export type AddToQueueResponse = [added: number, _b: any, _c: any, queueSize: number, addedToQueueTime: number];
+export type AddToQueueResponse = [added: number, queueId: QueueId | null, queueSize: number, addedToQueueTime: number];
 export type QueryStageStateResponse = [active: string[], toProcess: string[]] | [active: string[], toProcess: string[], defs: Record<string, QueryDef>];
 export type RetrieveForProcessingSuccess = [
-  added: any /** todo(ovr): Remove, useless */,
-  removed: any /** todo(ovr): Remove, useless */,
+  added: unknown,
+  // QueueId is required for Cube Store, other providers doesn't support it
+  queueId: QueueId | null,
   active: QueryKeyHash[],
   pending: number,
   def: QueryDef,
   lockAquired: true
 ];
 export type RetrieveForProcessingFail = [
-  added: any /** todo(ovr): Remove, useless */,
-  removed: any /** todo(ovr): Remove, useless */,
+  added: unknown,
+  // QueueId is required for Cube Store, other providers doesn't support it
+  queueId: QueueId | null,
   active: QueryKeyHash[],
   pending: number,
   def: null,
@@ -48,11 +54,9 @@ export interface QueueDriverOptions {
   processUid?: string;
 }
 
-export type ProcessingId = string | number;
-
 export interface QueueDriverConnectionInterface {
   redisHash(queryKey: QueryKey): QueryKeyHash;
-  getResultBlocking(queryKey: QueryKey): Promise<unknown>;
+  getResultBlocking(queryKey: QueryKeyHash, queueId: QueueId): Promise<unknown>;
   getResult(queryKey: QueryKey): Promise<any>;
   /**
    * Adds specified by the queryKey query to the queue, returns tuple
@@ -82,10 +86,10 @@ export interface QueueDriverConnectionInterface {
   // 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): Promise<boolean>;
+  optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, processingId: ProcessingId, queueId: QueueId | null): Promise<boolean>;
   cancelQuery(queryKey: QueryKey): Promise<QueryDef | null>;
   getQueryAndRemove(hash: QueryKeyHash): Promise<[QueryDef]>;
-  setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, processingId: ProcessingId): Promise<unknown>;
+  setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, processingId: ProcessingId, queueId: QueueId | null): Promise<unknown>;
   release(): void;
   //
   getQueriesToCancel(): Promise<string[]>

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

@@ -76,7 +76,8 @@ export class CubeStoreDriver extends BaseDriver implements DriverInterface {
 
   public async query<R = any>(query: string, values: any[], options?: QueryOptions): Promise<R[]> {
     const { inlineTables, ...queryTracingObj } = options ?? {};
-    return this.connection.query(formatSql(query, values || []), inlineTables ?? [], { ...queryTracingObj, instance: getEnv('instanceId') });
+    const sql = formatSql(query, values || []);
+    return this.connection.query(sql, inlineTables ?? [], { ...queryTracingObj, instance: getEnv('instanceId') });
   }
 
   public async release() {

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

@@ -12,6 +12,7 @@ import {
   QueryKey,
   QueryKeyHash,
   ProcessingId,
+  QueueId,
 } from '@cubejs-backend/base-driver';
 import { getProcessUid } from '@cubejs-backend/shared';
 
@@ -76,8 +77,7 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
     if (rows && rows.length) {
       return [
         rows[0].added === 'true' ? 1 : 0,
-        null,
-        null,
+        rows[0].id ? parseInt(rows[0].id, 10) : null,
         parseInt(rows[0].pending, 10),
         data.addedToQueueTime
       ];
@@ -243,9 +243,10 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
     return null;
   }
 
-  public async optimisticQueryUpdate(queryKey: any, toUpdate: any, _processingId: any): Promise<boolean> {
+  public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: unknown, _processingId: ProcessingId, queueId: QueueId): Promise<boolean> {
     await this.driver.query('QUEUE MERGE_EXTRA ? ?', [
-      this.prefixKey(queryKey),
+      // queryKeyHash as compatibility fallback
+      queueId || this.prefixKey(queryKeyHash),
       JSON.stringify(toUpdate)
     ]);
 
@@ -257,7 +258,7 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
   }
 
   public async retrieveForProcessing(queryKeyHashed: QueryKeyHash, _processingId: string): Promise<RetrieveForProcessingResponse> {
-    const rows = await this.driver.query<{ active: string | null, pending: string, payload: string, extra: string | null }>('QUEUE RETRIEVE EXTENDED CONCURRENCY ? ?', [
+    const rows = await this.driver.query<{ id: string /* cube store convert int64 to string */, active: string | null, pending: string, payload: string, extra: string | null }>('QUEUE RETRIEVE EXTENDED CONCURRENCY ? ?', [
       this.options.concurrency,
       this.prefixKey(queryKeyHashed),
     ]);
@@ -269,7 +270,12 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
         const def = this.decodeQueryDefFromRow(rows[0], 'retrieveForProcessing');
 
         return [
-          1, null, active, pending, def, true
+          1,
+          rows[0].id ? parseInt(rows[0].id, 10) : null,
+          active,
+          pending,
+          def,
+          true
         ];
       } else {
         return [
@@ -281,14 +287,11 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
     return null;
   }
 
-  public async getResultBlocking(queryKey: string): Promise<QueryDef | null> {
-    return this.getResultBlockingByHash(this.redisHash(queryKey));
-  }
-
-  public async getResultBlockingByHash(queryKeyHash: QueryKeyHash): Promise<QueryDef | null> {
+  public async getResultBlocking(queryKeyHash: QueryKeyHash, queueId: QueueId): Promise<QueryDef | null> {
     const rows = await this.driver.query('QUEUE RESULT_BLOCKING ? ?', [
       this.options.continueWaitTimeout * 1000,
-      this.prefixKey(queryKeyHash),
+      // queryKeyHash as compatibility fallback
+      queueId || this.prefixKey(queryKeyHash),
     ]);
     if (rows && rows.length) {
       return this.decodeQueryDefFromRow(rows[0], 'getResultBlocking');
@@ -297,9 +300,10 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
     return null;
   }
 
-  public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, _processingId: ProcessingId): Promise<boolean> {
+  public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, _processingId: ProcessingId, queueId: QueueId): Promise<boolean> {
     const rows = await this.driver.query('QUEUE ACK ? ? ', [
-      this.prefixKey(hash),
+      // queryKeyHash as compatibility fallback
+      queueId || this.prefixKey(hash),
       executionResult ? JSON.stringify(executionResult) : executionResult
     ]);
     if (rows && rows.length === 1) {

packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriver.js

@@ -48,11 +48,7 @@ export class LocalQueueDriverConnection {
     return this.resultPromises[resultListKey];
   }
 
-  async getResultBlocking(queryKey) {
-    return this.getResultBlockingByHash(this.redisHash(queryKey));
-  }
-
-  async getResultBlockingByHash(queryKeyHash) {
+  async getResultBlocking(queryKeyHash) {
     // Double redisHash apply is being used here
     const resultListKey = this.resultListKey(queryKeyHash);
     if (!this.queryDef[queryKeyHash] && !this.resultPromises[resultListKey]) {
@@ -144,7 +140,13 @@ export class LocalQueueDriverConnection {
       });
     }
 
-    return [added, null, null, Object.keys(this.toProcess).length, queryQueueObj.addedToQueueTime]; // TODO nulls
+    return [
+      added,
+      // this driver doesnt support queue id
+      null,
+      Object.keys(this.toProcess).length,
+      queryQueueObj.addedToQueueTime
+    ];
   }
 
   getToProcessQueries() {
@@ -268,8 +270,14 @@ export class LocalQueueDriverConnection {
     }
 
     return [
-      added, null, this.queueArray(this.active), Object.keys(this.toProcess).length, this.queryDef[key], lockAcquired
-    ]; // TODO nulls
+      added,
+      // this driver doesnt support queue id
+      null,
+      this.queueArray(this.active),
+      Object.keys(this.toProcess).length,
+      this.queryDef[key],
+      lockAcquired
+    ];
   }
 
   freeProcessingLock(queryKey, processingId, activated) {

packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.js

@@ -167,95 +167,6 @@ export class QueryQueue {
     return stream;
   }
 
-  /**
-   * Depends on the `queryHandler` value either runs `executeQueryInQueue`
-   * or `executeStreamInQueue` method.
-   *
-   * @param {string} queryHandler For the regular query is eq to 'query'.
-   * @param {*} queryKey
-   * @param {*} query
-   * @param {number=} priority
-   * @param {*=} options
-   * @returns {*}
-   *
-   * @throw {ContinueWaitError}
-   */
-  async executeInQueue(
-    queryHandler,
-    queryKey,
-    query,
-    priority,
-    options,
-  ) {
-    return this.executeQueryInQueue(
-      queryHandler,
-      queryKey,
-      query,
-      priority,
-      options,
-    );
-  }
-
-  /**
-   * Push persistent query to the queue and call `QueryQueue.reconcileQueue()` method.
-   *
-   * @param {string} queryHandler
-   * @param {*} queryKey
-   * @param {*} query
-   * @param {number=} priority
-   * @param {*=} options
-   * @returns {Promise<void>}
-   */
-  async executeStreamInQueue(
-    queryHandler,
-    queryKey,
-    query,
-    priority,
-    options,
-  ) {
-    options = options || {};
-    const queueConnection = await this.queueDriver.createConnection();
-    try {
-      priority = priority || 0;
-      if (!(priority >= -10000 && priority <= 10000)) {
-        throw new Error(
-          'Priority should be between -10000 and 10000'
-        );
-      }
-      const time = new Date().getTime();
-      const keyScore = time + (10000 - priority) * 1E14;
-
-      options.orphanedTimeout = query.orphanedTimeout;
-      const orphanedTimeout = 'orphanedTimeout' in query
-        ? query.orphanedTimeout
-        : this.orphanedTimeout;
-      const orphanedTime = time + (orphanedTimeout * 1000);
-
-      const [added, _b, _c, queueSize, addedToQueueTime] = await queueConnection.addToQueue(
-        keyScore, queryKey, orphanedTime, queryHandler, query, priority, options
-      );
-      if (added > 0) {
-        this.logger('Added to queue (persistent)', {
-          priority,
-          queueSize,
-          queryKey,
-          queuePrefix: this.redisQueuePrefix,
-          requestId: options.requestId,
-          metadata: query.metadata,
-          preAggregationId: query.preAggregation?.preAggregationId,
-          newVersionEntry: query.newVersionEntry,
-          forceBuild: query.forceBuild,
-          preAggregation: query.preAggregation,
-          addedToQueueTime,
-          persistent: queryKey.persistent,
-        });
-      }
-      this.reconcileQueue();
-    } finally {
-      this.queueDriver.release(queueConnection);
-    }
-  }
-
   /**
    * Push query to the queue and call `QueryQueue.reconcileQueue()` method if
    * `options.skipQueue` is set to `false`, execute query skipping queue
@@ -270,7 +181,7 @@ export class QueryQueue {
    *
    * @throw {ContinueWaitError}
    */
-  async executeQueryInQueue(
+  async executeInQueue(
     queryHandler,
     queryKey,
     query,
@@ -334,7 +245,7 @@ export class QueryQueue {
       const orphanedTimeout = 'orphanedTimeout' in query ? query.orphanedTimeout : this.orphanedTimeout;
       const orphanedTime = time + (orphanedTimeout * 1000);
 
-      const [added, _b, _c, queueSize, addedToQueueTime] = await queueConnection.addToQueue(
+      const [added, queueId, queueSize, addedToQueueTime] = await queueConnection.addToQueue(
         keyScore, queryKey, orphanedTime, queryHandler, query, priority, options
       );
 
@@ -400,7 +311,7 @@ export class QueryQueue {
       } else {
         // Result here won't be fetched for a jobed build query (initialized by
         // the /cubejs-system/v1/pre-aggregations/jobs endpoint).
-        result = !query.isJob && await queueConnection.getResultBlocking(queryKey);
+        result = !query.isJob && await queueConnection.getResultBlocking(queryKeyHash, queueId);
       }
 
       // We don't want to throw the ContinueWaitError for a jobed build query.
@@ -763,8 +674,8 @@ export class QueryQueue {
   }
 
   /**
-   * Processing query specified by the `queryKey`. This method incapsulate most
-   * of the logic related with the queues updates, heartbeating, etc.
+   * Processing query specified by the `queryKey`. This method encapsulate most
+   * of the logic related with the queues updates, heartbeat, etc.
    *
    * @param {QueryKeyHash} queryKeyHashed
    * @return {Promise<{ result: undefined | Object, error: string | undefined }>}
@@ -773,17 +684,18 @@ export class QueryQueue {
     const queueConnection = await this.queueDriver.createConnection();
 
     let insertedCount;
-    let _removedCount;
+    let queueId;
     let activeKeys;
     let queueSize;
     let query;
     let processingLockAcquired;
+
     try {
       const processingId = await queueConnection.getNextProcessingId();
       const retrieveResult = await queueConnection.retrieveForProcessing(queryKeyHashed, processingId);
 
       if (retrieveResult) {
-        [insertedCount /** todo(ovr): Remove */, _removedCount/** todo(ovr): Remove */, activeKeys, queueSize, query, processingLockAcquired] = retrieveResult;
+        [insertedCount, queueId, activeKeys, queueSize, query, processingLockAcquired] = retrieveResult;
       }
 
       const activated = activeKeys && activeKeys.indexOf(queryKeyHashed) !== -1;
@@ -808,7 +720,7 @@ export class QueryQueue {
           preAggregation: query.query?.preAggregation,
           addedToQueueTime: query.addedToQueueTime,
         });
-        await queueConnection.optimisticQueryUpdate(queryKeyHashed, { startQueryTime }, processingId);
+        await queueConnection.optimisticQueryUpdate(queryKeyHashed, { startQueryTime }, processingId, queueId);
 
         const heartBeatTimer = setInterval(
           () => queueConnection.updateHeartBeat(queryKeyHashed),
@@ -840,7 +752,7 @@ export class QueryQueue {
                     query.query,
                     async (cancelHandler) => {
                       try {
-                        return queueConnection.optimisticQueryUpdate(queryKeyHashed, { cancelHandler }, processingId);
+                        return queueConnection.optimisticQueryUpdate(queryKeyHashed, { cancelHandler }, processingId, queueId);
                       } catch (e) {
                         this.logger('Error while query update', {
                           queryKey: query.queryKey,
@@ -916,7 +828,7 @@ export class QueryQueue {
 
         clearInterval(heartBeatTimer);
 
-        if (!(await queueConnection.setResultAndRemoveQuery(queryKeyHashed, executionResult, processingId))) {
+        if (!(await queueConnection.setResultAndRemoveQuery(queryKeyHashed, executionResult, processingId, queueId))) {
           this.logger('Orphaned execution result', {
             processingId,
             warn: 'Result for query was not set due to processing lock wasn\'t acquired',
@@ -937,7 +849,7 @@ export class QueryQueue {
         // if (query?.queryHandler === 'stream') {
         //   const [active] = await queueConnection.getQueryStageState(true);
         //   if (active && active.length > 0) {
-        //     await Promise.race(active.map(keyHash => queueConnection.getResultBlockingByHash(keyHash)));
+        //     await Promise.race(active.map(keyHash => queueConnection.getResultBlocking(keyHash)));
         //     await this.reconcileQueue();
         //   }
         // }