diff --git a/plugin-server/src/utils/db/db.ts b/plugin-server/src/utils/db/db.ts index 7d59501006cf1..c89c08520fffe 100644 --- a/plugin-server/src/utils/db/db.ts +++ b/plugin-server/src/utils/db/db.ts @@ -228,6 +228,22 @@ export class DB { }) } + public redisGetBuffer(key: string, tag: string): Promise { + return instrumentQuery('query.redisGetBuffer', tag, async () => { + const client = await this.redisPool.acquire() + const timeout = timeoutGuard('Getting redis key delayed. Waiting over 30 sec to get key.', { key }) + try { + return await tryTwice( + async () => await client.getBuffer(key), + `Waited 5 sec to get redis key: ${key}, retrying once!` + ) + } finally { + clearTimeout(timeout) + await this.redisPool.release(client) + } + }) + } + public redisSet( key: string, value: unknown, @@ -254,6 +270,49 @@ export class DB { }) } + public redisSetBuffer(key: string, value: Buffer, tag: string, ttlSeconds?: number): Promise { + return instrumentQuery('query.redisSetBuffer', tag, async () => { + const client = await this.redisPool.acquire() + const timeout = timeoutGuard('Setting redis key delayed. Waiting over 30 sec to set key', { key }) + try { + if (ttlSeconds) { + await client.setBuffer(key, value, 'EX', ttlSeconds) + } else { + await client.setBuffer(key, value) + } + } finally { + clearTimeout(timeout) + await this.redisPool.release(client) + } + }) + } + + public redisSetNX( + key: string, + value: unknown, + tag: string, + ttlSeconds?: number, + options: CacheOptions = {} + ): Promise<'OK' | null> { + const { jsonSerialize = true } = options + + return instrumentQuery('query.redisSetNX', tag, async () => { + const client = await this.redisPool.acquire() + const timeout = timeoutGuard('Setting redis key delayed. Waiting over 30 sec to set key (NX)', { key }) + try { + const serializedValue = jsonSerialize ? JSON.stringify(value) : (value as string) + if (ttlSeconds) { + return await client.set(key, serializedValue, 'EX', ttlSeconds, 'NX') + } else { + return await client.set(key, serializedValue, 'NX') + } + } finally { + clearTimeout(timeout) + await this.redisPool.release(client) + } + }) + } + public redisSetMulti(kv: Array<[string, unknown]>, ttlSeconds?: number, options: CacheOptions = {}): Promise { const { jsonSerialize = true } = options @@ -403,6 +462,45 @@ export class DB { }) } + public redisSAddAndSCard(key: string, value: Redis.ValueType, ttlSeconds?: number): Promise { + return instrumentQuery('query.redisSAddAndSCard', undefined, async () => { + const client = await this.redisPool.acquire() + const timeout = timeoutGuard('SADD+SCARD delayed. Waiting over 30 sec to perform SADD+SCARD', { + key, + value, + }) + try { + const multi = client.multi() + multi.sadd(key, value) + if (ttlSeconds) { + multi.expire(key, ttlSeconds) + } + multi.scard(key) + const results = await multi.exec() + const scardResult = ttlSeconds ? results[2] : results[1] + return scardResult[1] + } finally { + clearTimeout(timeout) + await this.redisPool.release(client) + } + }) + } + + public redisSCard(key: string): Promise { + return instrumentQuery('query.redisSCard', undefined, async () => { + const client = await this.redisPool.acquire() + const timeout = timeoutGuard('SCARD delayed. Waiting over 30 sec to perform SCARD', { + key, + }) + try { + return await client.scard(key) + } finally { + clearTimeout(timeout) + await this.redisPool.release(client) + } + }) + } + public redisPublish(channel: string, message: string): Promise { return instrumentQuery('query.redisPublish', undefined, async () => { const client = await this.redisPool.acquire() diff --git a/plugin-server/tests/main/db.test.ts b/plugin-server/tests/main/db.test.ts index 10e514d8323c9..cb9d1ad5828de 100644 --- a/plugin-server/tests/main/db.test.ts +++ b/plugin-server/tests/main/db.test.ts @@ -903,6 +903,87 @@ describe('DB', () => { expect(fetchedTeam).toEqual(null) }) }) + + describe('redis', () => { + describe('buffer operations', () => { + it('writes and reads buffers', async () => { + const buffer = Buffer.from('test') + await db.redisSetBuffer('test', buffer, 'testTag', 60) + const result = await db.redisGetBuffer('test', 'testTag') + expect(result).toEqual(buffer) + }) + }) + + describe('redisSetNX', () => { + it('it should only set a value if there is not already one present', async () => { + const set1 = await db.redisSetNX('test', 'first', 'testTag') + expect(set1).toEqual('OK') + const get1 = await db.redisGet('test', '', 'testTag') + expect(get1).toEqual('first') + + const set2 = await db.redisSetNX('test', 'second', 'testTag') + expect(set2).toEqual(null) + const get2 = await db.redisGet('test', '', 'testTag') + expect(get2).toEqual('first') + }) + + it('it should only set a value if there is not already one present, with a ttl', async () => { + const set1 = await db.redisSetNX('test', 'first', 'testTag', 60) + expect(set1).toEqual('OK') + const get1 = await db.redisGet('test', '', 'testTag') + expect(get1).toEqual('first') + + const set2 = await db.redisSetNX('test', 'second', 'testTag', 60) + expect(set2).toEqual(null) + const get2 = await db.redisGet('test', '', 'testTag') + expect(get2).toEqual('first') + }) + }) + + describe('redisSAddAndSCard', () => { + it('it should add a value to a set and return the number of elements in the set', async () => { + const add1 = await db.redisSAddAndSCard('test', 'A') + expect(add1).toEqual(1) + const add2 = await db.redisSAddAndSCard('test', 'A') + expect(add2).toEqual(1) + const add3 = await db.redisSAddAndSCard('test', 'B') + expect(add3).toEqual(2) + const add4 = await db.redisSAddAndSCard('test', 'B') + expect(add4).toEqual(2) + const add5 = await db.redisSAddAndSCard('test', 'A') + expect(add5).toEqual(2) + }) + + it('it should add a value to a set and return the number of elements in the set, with a TTL', async () => { + const add1 = await db.redisSAddAndSCard('test', 'A', 60) + expect(add1).toEqual(1) + const add2 = await db.redisSAddAndSCard('test', 'A', 60) + expect(add2).toEqual(1) + const add3 = await db.redisSAddAndSCard('test', 'B', 60) + expect(add3).toEqual(2) + const add4 = await db.redisSAddAndSCard('test', 'B', 60) + expect(add4).toEqual(2) + const add5 = await db.redisSAddAndSCard('test', 'A', 60) + expect(add5).toEqual(2) + }) + }) + + describe('redisSCard', () => { + it('it should return the number of elements in the set', async () => { + await db.redisSAddAndSCard('test', 'A') + const scard1 = await db.redisSCard('test') + expect(scard1).toEqual(1) + + await db.redisSAddAndSCard('test', 'B') + const scard2 = await db.redisSCard('test') + expect(scard2).toEqual(2) + + await db.redisSAddAndSCard('test', 'B') + const scard3 = await db.redisSCard('test') + expect(scard3).toEqual(2) + }) + }) + }) }) describe('PostgresRouter()', () => {