Skip to content

Commit ac67e5b

Browse files
committed
fix: Redis driver execAsync ignores watch directives
1 parent 6f9393b commit ac67e5b

File tree

3 files changed

+116
-44
lines changed

3 files changed

+116
-44
lines changed

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@ const redis = require('redis');
22
const { promisify } = require('util');
33

44
module.exports = function createRedisClient(url) {
5-
redis.Multi.prototype.execAsync = promisify(redis.Multi.prototype.exec);
5+
redis.Multi.prototype.execAsync = function execAsync() {
6+
return new Promise((resolve, reject) => this.exec((err, res) => (
7+
err ? reject(err) : resolve(res)
8+
)));
9+
};
610

711
const options = {
812
url,
@@ -30,6 +34,7 @@ module.exports = function createRedisClient(url) {
3034
'zrangebyscore',
3135
'keys',
3236
'watch',
37+
'unwatch',
3338
'incr',
3439
'decr',
3540
'lpush'

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

Lines changed: 55 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -75,21 +75,25 @@ class RedisQueueDriverConnection {
7575
}
7676

7777
async setResultAndRemoveQuery(queryKey, executionResult, processingId) {
78-
await this.redisClient.watchAsync(this.queryProcessingLockKey(queryKey));
79-
const currentProcessId = await this.redisClient.getAsync(this.queryProcessingLockKey(queryKey));
80-
if (processingId !== currentProcessId) {
81-
return false;
82-
}
78+
try {
79+
await this.redisClient.watchAsync(this.queryProcessingLockKey(queryKey));
80+
const currentProcessId = await this.redisClient.getAsync(this.queryProcessingLockKey(queryKey));
81+
if (processingId !== currentProcessId) {
82+
return false;
83+
}
8384

84-
return this.redisClient.multi()
85-
.lpush([this.resultListKey(queryKey), JSON.stringify(executionResult)])
86-
.zrem([this.activeRedisKey(), this.redisHash(queryKey)])
87-
.zrem([this.heartBeatRedisKey(), this.redisHash(queryKey)])
88-
.zrem([this.toProcessRedisKey(), this.redisHash(queryKey)])
89-
.zrem([this.recentRedisKey(), this.redisHash(queryKey)])
90-
.hdel([this.queriesDefKey(), this.redisHash(queryKey)])
91-
.del(this.queryProcessingLockKey(queryKey))
92-
.execAsync();
85+
return this.redisClient.multi()
86+
.lpush([this.resultListKey(queryKey), JSON.stringify(executionResult)])
87+
.zrem([this.activeRedisKey(), this.redisHash(queryKey)])
88+
.zrem([this.heartBeatRedisKey(), this.redisHash(queryKey)])
89+
.zrem([this.toProcessRedisKey(), this.redisHash(queryKey)])
90+
.zrem([this.recentRedisKey(), this.redisHash(queryKey)])
91+
.hdel([this.queriesDefKey(), this.redisHash(queryKey)])
92+
.del(this.queryProcessingLockKey(queryKey))
93+
.execAsync();
94+
} finally {
95+
await this.redisClient.unwatchAsync();
96+
}
9397
}
9498

9599
getOrphanedQueries() {
@@ -144,43 +148,51 @@ class RedisQueueDriverConnection {
144148
}
145149

146150
async freeProcessingLock(queryKey, processingId, activated) {
147-
const lockKey = this.queryProcessingLockKey(queryKey);
148-
await this.redisClient.watchAsync(lockKey);
149-
const currentProcessId = await this.redisClient.getAsync(lockKey);
150-
if (currentProcessId === processingId) {
151-
let removeCommand = this.redisClient.multi()
152-
.del(lockKey);
153-
if (activated) {
154-
removeCommand = removeCommand.zrem([this.activeRedisKey(), this.redisHash(queryKey)]);
151+
try {
152+
const lockKey = this.queryProcessingLockKey(queryKey);
153+
await this.redisClient.watchAsync(lockKey);
154+
const currentProcessId = await this.redisClient.getAsync(lockKey);
155+
if (currentProcessId === processingId) {
156+
let removeCommand = this.redisClient.multi()
157+
.del(lockKey);
158+
if (activated) {
159+
removeCommand = removeCommand.zrem([this.activeRedisKey(), this.redisHash(queryKey)]);
160+
}
161+
await removeCommand
162+
.execAsync();
155163
}
156-
await removeCommand
157-
.execAsync();
164+
} finally {
165+
await this.redisClient.unwatchAsync();
158166
}
159167
}
160168

161169
async optimisticQueryUpdate(queryKey, toUpdate, processingId) {
162-
let query = await this.getQueryDef(queryKey);
163-
for (let i = 0; i < 10; i++) {
164-
if (query) {
165-
// eslint-disable-next-line no-await-in-loop
166-
await this.redisClient.watchAsync(this.queryProcessingLockKey(queryKey));
167-
const currentProcessId = await this.redisClient.getAsync(this.queryProcessingLockKey(queryKey));
168-
if (currentProcessId !== processingId) {
169-
return false;
170-
}
171-
let [beforeUpdate] = await this.redisClient
172-
.multi()
173-
.hget([this.queriesDefKey(), this.redisHash(queryKey)])
174-
.hset([this.queriesDefKey(), this.redisHash(queryKey), JSON.stringify({ ...query, ...toUpdate })])
175-
.execAsync();
176-
beforeUpdate = JSON.parse(beforeUpdate);
177-
if (JSON.stringify(query) === JSON.stringify(beforeUpdate)) {
178-
return true;
170+
try {
171+
let query = await this.getQueryDef(queryKey);
172+
for (let i = 0; i < 10; i++) {
173+
if (query) {
174+
// eslint-disable-next-line no-await-in-loop
175+
await this.redisClient.watchAsync(this.queryProcessingLockKey(queryKey));
176+
const currentProcessId = await this.redisClient.getAsync(this.queryProcessingLockKey(queryKey));
177+
if (currentProcessId !== processingId) {
178+
return false;
179+
}
180+
let [beforeUpdate] = await this.redisClient
181+
.multi()
182+
.hget([this.queriesDefKey(), this.redisHash(queryKey)])
183+
.hset([this.queriesDefKey(), this.redisHash(queryKey), JSON.stringify({ ...query, ...toUpdate })])
184+
.execAsync();
185+
beforeUpdate = JSON.parse(beforeUpdate);
186+
if (JSON.stringify(query) === JSON.stringify(beforeUpdate)) {
187+
return true;
188+
}
189+
query = beforeUpdate;
179190
}
180-
query = beforeUpdate;
181191
}
192+
throw new Error(`Can't update ${queryKey} with ${JSON.stringify(toUpdate)}`);
193+
} finally {
194+
await this.redisClient.unwatchAsync();
182195
}
183-
throw new Error(`Can't update ${queryKey} with ${JSON.stringify(toUpdate)}`);
184196
}
185197

186198
release() {

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,61 @@ const QueryQueueTest = (name, options) => {
143143
const result = await queue.executeInQueue('foo', query, query);
144144
expect(result).toBe('select * from bar');
145145
});
146+
147+
test('queue driver lock obtain race condition', async () => {
148+
const redisClient = await queue.queueDriver.createConnection();
149+
const redisClient2 = await queue.queueDriver.createConnection();
150+
const priority = 10;
151+
const time = new Date().getTime();
152+
const keyScore = time + (10000 - priority) * 1E14;
153+
154+
// console.log(await redisClient.getQueryAndRemove('race'));
155+
// console.log(await redisClient.getQueryAndRemove('race1'));
156+
157+
if (redisClient.redisClient) {
158+
await redisClient2.redisClient.setAsync(redisClient.queryProcessingLockKey('race'), '100');
159+
await redisClient.redisClient.watchAsync(redisClient.queryProcessingLockKey('race'));
160+
await redisClient2.redisClient.setAsync(redisClient.queryProcessingLockKey('race'), Math.random());
161+
162+
const res = await redisClient.redisClient.multi()
163+
.set(redisClient.queryProcessingLockKey('race'), '100')
164+
.set(redisClient.queryProcessingLockKey('race1'), '100')
165+
.execAsync();
166+
167+
expect(res).toBe(null);
168+
await redisClient.redisClient.delAsync(redisClient.queryProcessingLockKey('race'));
169+
await redisClient.redisClient.delAsync(redisClient.queryProcessingLockKey('race1'));
170+
}
171+
172+
await queue.reconcileQueue();
173+
174+
await redisClient.addToQueue(
175+
keyScore, 'race', time, 'handler', ['select'], priority, { stageQueryKey: 'race' }
176+
);
177+
178+
await redisClient.addToQueue(
179+
keyScore + 100, 'race2', time + 100, 'handler2', ['select2'], priority, { stageQueryKey: 'race2' }
180+
);
181+
182+
const processingId1 = await redisClient.getNextProcessingId();
183+
const processingId4 = await redisClient.getNextProcessingId();
184+
185+
await redisClient.freeProcessingLock('race', processingId1, true);
186+
await redisClient.freeProcessingLock('race2', processingId4, true);
187+
188+
await redisClient2.retrieveForProcessing('race2', await redisClient.getNextProcessingId());
189+
190+
const processingId = await redisClient.getNextProcessingId();
191+
const retrieve6 = await redisClient.retrieveForProcessing('race', processingId);
192+
console.log(retrieve6);
193+
expect(!!retrieve6[5]).toBe(true);
194+
195+
console.log(await redisClient.getQueryAndRemove('race'));
196+
console.log(await redisClient.getQueryAndRemove('race2'));
197+
198+
await queue.queueDriver.release(redisClient);
199+
await queue.queueDriver.release(redisClient2);
200+
});
146201
});
147202
};
148203

0 commit comments

Comments
 (0)