fix(query-orchestrator): Reduce number of refresh key queries (#9809)

Cube generates a burst load of queries for refresh keys. We don't use a queue for cube store queries, but we use an in-memory cache inside QueryQueue. A cache miss leads to a problem where we start executing queries without debouncing.

Cube executes refresh keys in parallel, which DDoSes Cube Store. It's a quick fix to fix this behavior.

Author: ovr

packages/cubejs-backend-shared/src/promises.ts

@@ -265,7 +265,7 @@ export const retryWithTimeout = <T>(
   );
 
 /**
- * High order function that makes to debounce multi async calls to single call at one time
+ * Creates a debounced version of an asynchronous function.
  */
 export const asyncDebounce = <Ret, Arguments>(
   fn: (...args: Arguments[]) => Promise<Ret>,

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

@@ -2,7 +2,7 @@ import crypto from 'crypto';
 import csvWriter from 'csv-write-stream';
 import { LRUCache } from 'lru-cache';
 import { pipeline } from 'stream';
-import { getEnv, MaybeCancelablePromise, streamToArray } from '@cubejs-backend/shared';
+import { asyncDebounce, getEnv, MaybeCancelablePromise, streamToArray } from '@cubejs-backend/shared';
 import { CubeStoreCacheDriver, CubeStoreDriver } from '@cubejs-backend/cubestore-driver';
 import {
   BaseDriver,
@@ -34,6 +34,12 @@ export type QueryWithParams = [
   options?: QueryOptions
 ];
 
+export type LoadRefreshKeyOptions = {
+  requestId?: string;
+  skipRefreshKeyWaitForRenew?: boolean;
+  dataSource: string
+};
+
 export type Query = {
   requestId?: string;
   dataSource: string;
@@ -771,32 +777,31 @@ export class QueryCache {
   public loadRefreshKeys(
     cacheKeyQueries: QueryWithParams[],
     expireSecs: number,
-    options: {
-      requestId?: string;
-      skipRefreshKeyWaitForRenew?: boolean;
-      dataSource: string
-    }
+    options: LoadRefreshKeyOptions
   ) {
-    return cacheKeyQueries.map((q) => {
-      const [query, values, queryOptions]: QueryWithParams = Array.isArray(q) ? q : [q, [], {}];
-      return this.cacheQueryResult(
-        query,
-        values,
-        [query, values],
-        expireSecs,
-        {
-          renewalThreshold: this.options.refreshKeyRenewalThreshold || queryOptions?.renewalThreshold || 2 * 60,
-          renewalKey: q,
-          waitForRenew: !options.skipRefreshKeyWaitForRenew,
-          requestId: options.requestId,
-          dataSource: options.dataSource,
-          useInMemory: true,
-          external: queryOptions?.external,
-        },
-      );
-    });
+    return cacheKeyQueries.map((q) => this.loadRefreshKey(q, expireSecs, options));
   }
 
+  public loadRefreshKey = asyncDebounce(async (q: QueryWithParams, expireSecs: number, options: LoadRefreshKeyOptions) => {
+    const [query, values, queryOptions]: QueryWithParams = Array.isArray(q) ? q : [q, [], {}];
+
+    return this.cacheQueryResult(
+      query,
+      values,
+      [query, values],
+      expireSecs,
+      {
+        renewalThreshold: this.options.refreshKeyRenewalThreshold || queryOptions?.renewalThreshold || 2 * 60,
+        renewalKey: q,
+        waitForRenew: !options.skipRefreshKeyWaitForRenew,
+        requestId: options.requestId,
+        dataSource: options.dataSource,
+        useInMemory: true,
+        external: queryOptions?.external,
+      },
+    );
+  });
+
   public withLock = <T = any>(
     key: string,
     ttl: number,

packages/cubejs-testing/test/smoke-cubesql.test.ts

@@ -156,7 +156,7 @@ describe('SQL API', () => {
           Authorization: token,
         },
         body: JSON.stringify({
-          query: `SELECT orderDate FROM ECommerce LIMIT 0;`,
+          query: 'SELECT orderDate FROM ECommerce LIMIT 0;',
         }),
       });