feat(query-orchestrator): Debounce updates for table used/touch

cube-js · Nov 20, 2023 · 20ed9b9 · 20ed9b9
1 parent a5ff0f7
commit 20ed9b9
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 25 deletions.
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/MemoryQueue.ts b/packages/cubejs-query-orchestrator/src/orchestrator/MemoryQueue.ts
@@ -0,0 +1,75 @@
+import { QueryCache } from './QueryCache';
+
+export abstract class AbstractSetMemoryQueue {
+  protected readonly queue: Set<string> = new Set();
+
+  public constructor(
+    protected readonly capacity: number,
+    protected readonly concurrency: number,
+  ) {
+
+  }
+
+  protected execution: boolean = false;
+
+  public addToQueue(item: string) {
+    this.queue.add(item);
+
+    if (this.queue.size > 100) {
+      console.log('Too large capacity', this.queue.size);
+    }
+
+    this.run().catch(e => console.log(e));
+  }
+
+  public async run(): Promise<void> {
+    if (this.execution) {
+      return;
+    }
+
+    this.execution = true;
+
+    try {
+      let toExecute: string[] = [];
+
+      do {
+        for (const item of this.queue) {
+          toExecute.push(item);
+          this.queue.delete(item);
+
+          if (toExecute.length > this.concurrency) {
+            break;
+          }
+        }
+
+        console.log('toExecute', toExecute.length, {
+          toExecute
+        });
+
+        await Promise.all(toExecute.map(async (item) => this.execute(item)));
+        toExecute = [];
+      } while (toExecute.length > 0);
+    } finally {
+      this.execution = false;
+    }
+  }
+
+  abstract execute(item: string): Promise<void>;
+}
+
+export class TableTouchMemoryQueue extends AbstractSetMemoryQueue {
+  public constructor(
+    capacity: number,
+    concurrency: number,
+    protected readonly queryCache: QueryCache,
+    protected readonly touchTablePersistTime: number
+  ) {
+    super(capacity, concurrency);
+  }
+
+  public async execute(tableName: string): Promise<void> {
+    const key = this.queryCache.getKey('SQL_PRE_AGGREGATIONS_TABLES_TOUCH', tableName);
+    console.log('touch', key);
+    await this.queryCache.getCacheDriver().set(key, new Date().getTime(), this.touchTablePersistTime);
+  }
+}
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts
@@ -33,6 +33,7 @@ import { QueryQueue } from './QueryQueue';
 import { LargeStreamWarning } from './StreamObjectsCounter';
 import { CacheAndQueryDriverType } from './QueryOrchestrator';
 import { RedisPool } from './RedisPool';
+import { TableTouchMemoryQueue } from './MemoryQueue';
 
 /// Name of the inline table containing the lambda rows.
 export const LAMBDA_TABLE_PREFIX = 'lambda';
@@ -653,7 +654,7 @@ export class PreAggregationLoader {
     if (versionEntryByContentVersion && !this.forceBuild) {
       const targetTableName = this.targetTableName(versionEntryByContentVersion);
       // No need to block here
-      this.updateLastTouch(targetTableName);
+      this.preAggregations.updateLastTouch(targetTableName);
       return {
         targetTableName,
         refreshKeyValues: [],
@@ -672,7 +673,7 @@ export class PreAggregationLoader {
       if (versionEntryByStructureVersion) {
         const targetTableName = this.targetTableName(versionEntryByStructureVersion);
         // No need to block here
-        this.updateLastTouch(targetTableName);
+        this.preAggregations.updateLastTouch(targetTableName);
         return {
           targetTableName,
           refreshKeyValues: [],
@@ -712,7 +713,8 @@ export class PreAggregationLoader {
         throw new Error(`Pre-aggregation table is not found for ${this.preAggregation.tableName} after it was successfully created`);
       }
       const targetTableName = this.targetTableName(lastVersion);
-      this.updateLastTouch(targetTableName);
+      // No need to block here
+      this.preAggregations.updateLastTouch(targetTableName);
       return {
         targetTableName,
         refreshKeyValues: [],
@@ -743,7 +745,7 @@ export class PreAggregationLoader {
             });
           });
         const targetTableName = this.targetTableName(newVersionEntry);
-        this.updateLastTouch(targetTableName);
+        this.preAggregations.updateLastTouch(targetTableName);
         return {
           targetTableName,
           refreshKeyValues: [],
@@ -791,7 +793,7 @@ export class PreAggregationLoader {
       return mostRecentResult();
     }
     const targetTableName = this.targetTableName(versionEntry);
-    this.updateLastTouch(targetTableName);
+    this.preAggregations.updateLastTouch(targetTableName);
     return {
       targetTableName,
       refreshKeyValues: [],
@@ -800,14 +802,6 @@ export class PreAggregationLoader {
     };
   }
 
-  private updateLastTouch(tableName: string) {
-    this.preAggregations.updateLastTouch(tableName).catch(e => {
-      this.logger('Error on pre-aggregation touch', {
-        error: (e.stack || e), preAggregation: this.preAggregation, requestId: this.requestId,
-      });
-    });
-  }
-
   protected contentVersion(invalidationKeys) {
     const versionArray = [this.preAggregation.structureVersionLoadSql || this.preAggregation.loadSql];
     if (this.preAggregation.indexesSql && this.preAggregation.indexesSql.length) {
@@ -895,7 +889,7 @@ export class PreAggregationLoader {
   }
 
   public refresh(newVersionEntry: VersionEntry, invalidationKeys: InvalidationKeys, client) {
-    this.updateLastTouch(this.targetTableName(newVersionEntry));
+    this.preAggregations.updateLastTouch(this.targetTableName(newVersionEntry));
     let refreshStrategy = this.refreshStoreInSourceStrategy;
     if (this.preAggregation.external) {
       const readOnly =
@@ -1971,6 +1965,9 @@ export class PreAggregations {
 
   private readonly getQueueEventsBus: any;
 
+  private readonly touchQueue: TableTouchMemoryQueue;
+
+
   public constructor(
     private readonly redisPrefix: string,
     private readonly driverFactory: DriverFactoryByDataSource,
@@ -1979,10 +1976,10 @@ export class PreAggregations {
     options,
   ) {
     this.options = options || {};
-
+    this.touchTablePersistTime = options.touchTablePersistTime || getEnv('touchPreAggregationTimeout');
+    this.touchQueue = new TableTouchMemoryQueue(16_768, 5, this.queryCache, this.touchTablePersistTime);
     this.externalDriverFactory = options.externalDriverFactory;
     this.structureVersionPersistTime = options.structureVersionPersistTime || 60 * 60 * 24 * 30;
-    this.touchTablePersistTime = options.touchTablePersistTime || getEnv('touchPreAggregationTimeout');
     this.dropPreAggregationsWithoutTouch = options.dropPreAggregationsWithoutTouch || getEnv('dropPreAggregationsWithoutTouch');
     this.usedTablePersistTime = options.usedTablePersistTime || getEnv('dbQueryTimeout');
     this.externalRefresh = options.externalRefresh;
@@ -1994,17 +1991,13 @@ export class PreAggregations {
     return this.queryCache.getKey('SQL_PRE_AGGREGATIONS_TABLES_USED', tableName);
   }
 
-  protected tablesTouchRedisKey(tableName) {
-    // TODO add dataSource?
-    return this.queryCache.getKey('SQL_PRE_AGGREGATIONS_TABLES_TOUCH', tableName);
-  }
-
   protected refreshEndReachedKey() {
     // TODO add dataSource?
     return this.queryCache.getKey('SQL_PRE_AGGREGATIONS_REFRESH_END_REACHED', '');
   }
 
   public async addTableUsed(tableName) {
+    console.log('tableUsed', tableName);
     return this.queryCache.getCacheDriver().set(this.tablesUsedRedisKey(tableName), true, this.usedTablePersistTime);
   }
 
@@ -2013,13 +2006,16 @@ export class PreAggregations {
       .map(k => k.replace(this.tablesUsedRedisKey(''), ''));
   }
 
-  public async updateLastTouch(tableName) {
-    return this.queryCache.getCacheDriver().set(this.tablesTouchRedisKey(tableName), new Date().getTime(), this.touchTablePersistTime);
+  public updateLastTouch(tableName: string) {
+    this.touchQueue.addToQueue(tableName);
   }
 
   public async tablesTouched() {
-    return (await this.queryCache.getCacheDriver().keysStartingWith(this.tablesTouchRedisKey('')))
-      .map(k => k.replace(this.tablesTouchRedisKey(''), ''));
+    // TODO add dataSource?
+    const key = this.queryCache.getKey('SQL_PRE_AGGREGATIONS_TABLES_TOUCH', '');
+
+    return (await this.queryCache.getCacheDriver().keysStartingWith(key))
+      .map(k => k.replace(key, ''));
   }
 
   public async updateRefreshEndReached() {