Optimization using auto-pipeling @ upstash-redis and Promise.all()

Author: SomajitDey

api/helper.js

@@ -1,5 +1,10 @@
+/*
+Refs:
+https://upstash.com/docs/redis/sdks/ts/pipelining/pipeline-transaction
+https://upstash.com/docs/redis/sdks/ts/pipelining/auto-pipeline
+*/
 import Crypto from 'node:crypto';
-import { createClient } from '@vercel/kv';
+import { Redis } from '@upstash/redis';
 
 const secret = process.env.SECRET;
 const sigLen = parseInt(process.env.SIG_LEN);
@@ -13,14 +18,18 @@ const dbKeyPrefix = {
                 cache: "cache:"
             }
 // Redis client for user database
-const redisData = createClient({
+const redisData = new Redis({
   url: process.env.UPSTASH_REDIS_REST_URL,
   token: process.env.UPSTASH_REDIS_REST_TOKEN,
+  latencyLogging: false,
+  enableAutoPipelining: true
 })
 // Redis client for ratelimiter database
-const redisRateLimit = createClient({
+const redisRateLimit = new Redis({
   url: process.env.KV_REST_API_URL,
   token: process.env.KV_REST_API_TOKEN,
+  latencyLogging: false,
+  enableAutoPipelining: true
 })
 
 function hash(str){
@@ -61,7 +70,13 @@ export function genPublicKey(privateOrPublicKey){
 export function cacheSet(privateKey, obj){
     const publicKey = genPublicKey(privateKey);
     const dbKey = dbKeyPrefix.cache + publicKey;
-    return redisRateLimit.hset(dbKey, obj).then(redisRateLimit.expire(dbKey, cacheTtl));
+    // Promise.all below enables both commands to be executed in a single http request (using same pipeline)
+    // As Redis is single-threaded, the commands are executed in order
+    // See https://upstash.com/docs/redis/sdks/ts/pipelining/auto-pipeline
+    return Promise.all([
+      redisRateLimit.hset(dbKey, obj),
+      redisRateLimit.expire(dbKey, cacheTtl)
+    ])
 }
 
 export function cacheGet(publicKey, key){
@@ -84,15 +99,20 @@ export function genKeyPair(seed = Crypto.randomUUID()){
 
 export async function publicProduce(publicKey, data){
     const dbKey = dbKeyPrefix.manyToOne + publicKey;
-    return redisData.rpush(dbKey, data).then(redisData.expire(dbKey, ttl));
+    return Promise.all([
+      redisData.rpush(dbKey, data),
+      redisData.expire(dbKey, ttl)
+    ])
 }
 
 export async function privateConsume(privateKey){
     const publicKey = genPublicKey(privateKey);
     const dbKey = dbKeyPrefix.manyToOne + publicKey;
-    const llen = await redisData.llen(dbKey);
-    if (!llen) return [];
-    return redisData.lpop(dbKey, llen);
+    const atomicTransaction = redisData.multi();
+    atomicTransaction.lrange(dbKey, 0, -1);
+    atomicTransaction.del(dbKey);
+    return atomicTransaction.exec()
+      .then((values) => values[0]);
 }
 
 export async function privateProduce(privateKey, data){
@@ -117,9 +137,11 @@ export async function privateStats(privateKey){
     const publicKey = genPublicKey(privateKey);
     const dbKeyConsume = dbKeyPrefix.manyToOne + publicKey;
     const dbKeyPublish = dbKeyPrefix.oneToMany + publicKey;
-    const countConsume = await redisData.llen(dbKeyConsume);
-    const ttlConsume = await redisData.ttl(dbKeyConsume);
-    const ttlPublish = await redisData.ttl(dbKeyPublish);
+    const [ countConsume, ttlConsume, ttlPublish ] = await Promise.all([
+      redisData.llen(dbKeyConsume),
+      redisData.ttl(dbKeyConsume),
+      redisData.ttl(dbKeyPublish)
+    ])
     return {
       consume: {
         count: countConsume,
@@ -141,23 +163,29 @@ export async function oneToOneProduce(privateKey, key, data){
     const dbKey = dbKeyPrefix.oneToOne + publicKey;
     let field = {};
     field[key] = data;
-    return redisData.hset(dbKey, field).then(redisData.expire(dbKey, ttl));
+    return Promise.all([
+      redisData.hset(dbKey, field),
+      redisData.expire(dbKey, ttl)
+    ])
 }
 
 export async function oneToOneConsume(publicKey, key){
     const dbKey = dbKeyPrefix.oneToOne + publicKey;
     const field = key;
-    return redisData.hget(dbKey, field).then(redisData.hdel(dbKey, field));
+    const atomicTransaction = redisData.multi();
+    atomicTransaction.hget(dbKey, field);
+    atomicTransaction.hdel(dbKey, field);
+    return atomicTransaction.exec()
+      .then((values) => values[0]);
 }
 
 export async function oneToOneTTL(privateKey, key){
     const publicKey = genPublicKey(privateKey);
     const dbKey = dbKeyPrefix.oneToOne + publicKey;
     const field = key;
-    const bool = await redisData.hexists(dbKey, field);
-    if (bool) {
-        return {ttl: await redisData.ttl(dbKey)};
-    } else {
-        return {ttl: 0};
-    }
+    const [ bool, ttl ] = await Promise.all([
+      redisData.hexists(dbKey, field),
+      redisData.ttl(dbKey)
+    ])
+    return {ttl: bool ? ttl : 0};
 }

api/index.js

@@ -72,7 +72,7 @@ fastify.post('/public/:publicKey', async (request, reply) => {
         // Try posting the data to webhook, if any, with timeout. On fail, store/bin data for later retrieval.
         try {
             const webhook = await helper.cacheGet(publicKey, 'hook');
-            if (webhook == null) throw new Error('No webhook');
+            if (!webhook) throw new Error('No webhook');
 
             await fetch(webhook, {
                 method: "POST",
@@ -86,9 +86,10 @@ fastify.post('/public/:publicKey', async (request, reply) => {
 
             webhookUsed = webhook;
         } catch (err) {
-            await helper.publicProduce(publicKey, data);
-            // Delete webhook from cache if webhook is of no use!
-            if (err.message !== 'No webhook') await helper.cacheDel(publicKey, 'hook');
+            await Promise.all([
+              helper.publicProduce(publicKey, data),
+              helper.cacheDel(publicKey, 'hook')
+            ])
         }
 
         if (redirectOnOk == null) {
@@ -113,15 +114,30 @@ fastify.get('/private/:privateKey', async (request, reply) => {
     const { privateKey } = request.params;
     const webhook = request.query.hook;
     const statsPresent = 'stats' in request.query;
+
+    let webhookHandler, statsHandler, dataHandler;
+    webhookHandler = statsHandler = dataHandler = () => null; // default
+
     try {
         if (helper.validate(privateKey) !== 'private') throw 401;
         if (webhook == null) {
-            await helper.cacheDel(privateKey, 'hook');
+            webhookHandler = () => helper.cacheDel(privateKey, 'hook');
+        } else {
+            webhookHandler = () => helper.cacheSet(privateKey, {hook:webhook});
+        }
+        if (statsPresent) {
+          statsHandler = () => helper.privateStats(privateKey);
         } else {
-            await helper.cacheSet(privateKey, {hook:webhook});
+          dataHandler = () => helper.privateConsume(privateKey);
         }
-        if (statsPresent) return helper.privateStats(privateKey);
-        const dataArray = await helper.privateConsume(privateKey);
+        
+        const [ _, statsObj, dataArray ] = await Promise.all([
+          webhookHandler(),
+          statsHandler(),
+          dataHandler()
+        ])
+        
+        if (statsPresent) return statsObj;
         if (!dataArray.length) throw 404;
         reply.send(dataArray);
     } catch (err) {

api/test.js

@@ -25,6 +25,6 @@ console.log(await helper.publicConsume(key.public));
 
 await helper.oneToOneProduce(key.private, 'some Key', 'data "for one to one at some key');
 console.log(await helper.oneToOneConsume(key.public, 'some Key'));
-console.log(await helper.oneToOneIsConsumed(key.private, 'some Key'));
+console.log(await helper.oneToOneTTL(key.private, 'some Key'));
 
 console.log('End of synchronous execution. Anything logged after this is from async only!')

middleware.js

@@ -28,12 +28,8 @@ const ratelimit = new Ratelimit({
 
 export default async function middleware(request) {
   // You could alternatively limit based on user ID or similar
-  const ip = ipAddress(request) || '127.0.0.1'
-  const { success, pending, reset } = await ratelimit.limit(
-    ip
-  )
-
-  await pending;
+  const ip = ipAddress(request) || '127.0.0.1';
+  const { success, reset } = await ratelimit.limit(ip);
 
   return success ? next() : Response.json(
     { message: `Try after ${(reset - Date.now())/1000} seconds`, error: "Too Many Requests", statusCode: 429 },