fix: Orphaned queries in Redis queue during intensive load

Author: paveltiunov

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

@@ -164,6 +164,8 @@ class LocalQueueDriverConnection {
     if (!this.processingLocks[key]) {
       this.processingLocks[key] = processingId;
       lockAcquired = true;
+    } else {
+      return null;
     }
     let added = 0;
     if (Object.keys(this.active).length < this.concurrency && !this.active[key]) {

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

@@ -217,9 +217,14 @@ class QueryQueue {
     let processingLockAcquired;
     try {
       const processingId = await redisClient.getNextProcessingId();
-      [insertedCount, removedCount, activeKeys, queueSize, query, processingLockAcquired] =
-        await redisClient.retrieveForProcessing(queryKey, processingId);
-      const activated = activeKeys.indexOf(this.redisHash(queryKey)) !== -1;
+      const retrieveResult = await redisClient.retrieveForProcessing(queryKey, processingId);
+      if (retrieveResult) {
+        [insertedCount, removedCount, activeKeys, queueSize, query, processingLockAcquired] = retrieveResult;
+      }
+      const activated = activeKeys && activeKeys.indexOf(this.redisHash(queryKey)) !== -1;
+      if (!query) {
+        query = await redisClient.getQueryDef(this.redisHash(queryKey));
+      }
       if (query && insertedCount && activated && processingLockAcquired) {
         let executionResult;
         const startQueryTime = (new Date()).getTime();

packages/cubejs-query-orchestrator/orchestrator/RedisQueueDriver.js

@@ -134,17 +134,33 @@ class RedisQueueDriverConnection {
   }
 
   async retrieveForProcessing(queryKey, processingId) {
-    const [insertedCount, removedCount, activeKeys, queueSize, queryDef, processingLockAcquired] =
-      await this.redisClient.multi()
-        .zadd([this.activeRedisKey(), 'NX', processingId, this.redisHash(queryKey)])
-        .zremrangebyrank([this.activeRedisKey(), this.concurrency, -1])
-        .zrange([this.activeRedisKey(), 0, this.concurrency - 1])
-        .zcard(this.toProcessRedisKey())
-        .hget(([this.queriesDefKey(), this.redisHash(queryKey)]))
-        .set(this.queryProcessingLockKey(queryKey), processingId, 'NX')
-        .zadd([this.heartBeatRedisKey(), 'NX', new Date().getTime(), this.redisHash(queryKey)])
-        .execAsync();
-    return [insertedCount, removedCount, activeKeys, queueSize, JSON.parse(queryDef), processingLockAcquired];
+    try {
+      const lockKey = this.queryProcessingLockKey(queryKey);
+      await this.redisClient.watchAsync(lockKey);
+
+      const currentProcessId = await this.redisClient.getAsync(lockKey);
+
+      if (currentProcessId) {
+        return null;
+      }
+
+      const result =
+        await this.redisClient.multi()
+          .zadd([this.activeRedisKey(), 'NX', processingId, this.redisHash(queryKey)])
+          .zremrangebyrank([this.activeRedisKey(), this.concurrency, -1])
+          .zrange([this.activeRedisKey(), 0, this.concurrency - 1])
+          .zcard(this.toProcessRedisKey())
+          .hget(([this.queriesDefKey(), this.redisHash(queryKey)]))
+          .set(lockKey, processingId, 'NX')
+          .zadd([this.heartBeatRedisKey(), 'NX', new Date().getTime(), this.redisHash(queryKey)])
+          .execAsync();
+      if (result) {
+        result[4] = JSON.parse(result[4]);
+      }
+      return result;
+    } finally {
+      await this.redisClient.unwatchAsync();
+    }
   }
 
   async freeProcessingLock(queryKey, processingId, activated) {

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

@@ -198,6 +198,50 @@ const QueryQueueTest = (name, options) => {
       await queue.queueDriver.release(redisClient);
       await queue.queueDriver.release(redisClient2);
     });
+
+    test('activated but lock is not acquired', async () => {
+      const redisClient = await queue.queueDriver.createConnection();
+      const redisClient2 = await queue.queueDriver.createConnection();
+      const priority = 10;
+      const time = new Date().getTime();
+      const keyScore = time + (10000 - priority) * 1E14;
+
+      await queue.reconcileQueue();
+
+      await redisClient.addToQueue(
+        keyScore, 'activated1', time, 'handler', ['select'], priority, { stageQueryKey: 'race' }
+      );
+
+      await redisClient.addToQueue(
+        keyScore + 100, 'activated2', time + 100, 'handler2', ['select2'], priority, { stageQueryKey: 'race2' }
+      );
+
+      const processingId1 = await redisClient.getNextProcessingId();
+      const processingId2 = await redisClient.getNextProcessingId();
+      const processingId3 = await redisClient.getNextProcessingId();
+
+      const retrieve1 = await redisClient.retrieveForProcessing('activated1', processingId1);
+      console.log(retrieve1);
+      const retrieve2 = await redisClient2.retrieveForProcessing('activated2', processingId2);
+      console.log(retrieve2);
+      console.log(await redisClient.freeProcessingLock('activated1', processingId1, retrieve1 && retrieve1[2].indexOf('activated1') !== -1));
+      const retrieve3 = await redisClient.retrieveForProcessing('activated2', processingId3);
+      console.log(retrieve3);
+      console.log(await redisClient.freeProcessingLock('activated2', processingId3, retrieve3 && retrieve3[2].indexOf('activated2') !== -1));
+      console.log(retrieve2[2].indexOf('activated2') !== -1);
+      console.log(await redisClient2.freeProcessingLock('activated2', processingId2, retrieve2 && retrieve2[2].indexOf('activated2') !== -1));
+
+      const retrieve4 = await redisClient.retrieveForProcessing('activated2', await redisClient.getNextProcessingId());
+      console.log(retrieve4);
+      expect(retrieve4[0]).toBe(1);
+      expect(!!retrieve4[5]).toBe(true);
+
+      console.log(await redisClient.getQueryAndRemove('activated1'));
+      console.log(await redisClient.getQueryAndRemove('activated2'));
+
+      await queue.queueDriver.release(redisClient);
+      await queue.queueDriver.release(redisClient2);
+    });
   });
 };