From 19933201b688be4ed44bb4b1c246f4a2cb4281f4 Mon Sep 17 00:00:00 2001
From: nick-funk <code@nickfunk.com>
Date: Tue, 10 Sep 2024 11:46:27 -0600
Subject: [PATCH] optimize mget performance on counts v2 endpoint

---
 .../server/app/handlers/api/story/count.ts    | 87 +++++++++++--------
 1 file changed, 52 insertions(+), 35 deletions(-)

diff --git a/server/src/core/server/app/handlers/api/story/count.ts b/server/src/core/server/app/handlers/api/story/count.ts
index d9077c5a99..098cb2a26b 100644
--- a/server/src/core/server/app/handlers/api/story/count.ts
+++ b/server/src/core/server/app/handlers/api/story/count.ts
@@ -186,6 +186,12 @@ async function calculateStoryCount(
 
 export type CountV2Options = Pick<AppOptions, "mongo" | "redis">;
 
+interface CountResult {
+  storyID: string;
+  redisCount?: number;
+  count: number;
+}
+
 export const countsV2Handler =
   ({ mongo, redis }: CountV2Options): RequestHandler<TenantCoralRequest> =>
   async (req, res, next) => {
@@ -198,42 +204,53 @@ export const countsV2Handler =
 
       const { storyIDs }: CountsV2Body = validate(CountsV2BodySchema, req.body);
 
-      const storyCounts = await Promise.all(
-        storyIDs.map(async (storyID) => {
-          // check that cache is available
-
-          // try to look in Redis cache here first by key
-          const key = `${tenant.id}:${storyID}:count`;
-          const redisCount = await redis.get(key);
-
-          if (redisCount) {
-            logger.debug("found story count for counts v2 in redis cache", {
-              storyID,
-              redisCount,
-            });
-            return { storyID, redisCount };
-          } else {
-            const count = await retrieveStoryCommentCounts(
-              mongo,
-              tenant.id,
-              storyID
-            );
-
-            // add count to Redis cache here then return
-            await redis.set(key, count);
-            logger.debug("set story count for counts v2 in redis cache", {
-              storyID,
-              count,
-            });
-
-            return {
-              storyID,
-              count,
-            };
-          }
-        })
+      const redisCounts = await redis.mget(
+        ...storyIDs.map((id) => `${tenant.id}:${id}:count`)
       );
-      res.send(JSON.stringify(storyCounts));
+
+      const countResults = new Map<string, CountResult>();
+      const missingIDs: string[] = [];
+
+      for (let i = 0; i < storyIDs.length; i++) {
+        const storyID = storyIDs[i];
+        const redisCount = redisCounts[i];
+
+        if (redisCount !== null && redisCount !== undefined) {
+          try {
+            const count = parseInt(redisCount, 10);
+            countResults.set(storyID, { storyID, redisCount: count, count });
+          } catch {
+            missingIDs.push(storyID);
+          }
+        } else {
+          missingIDs.push(storyID);
+        }
+      }
+
+      for (const missingID of missingIDs) {
+        const count = await retrieveStoryCommentCounts(
+          mongo,
+          tenant.id,
+          missingID
+        );
+
+        const key = `${tenant.id}:${missingID}:count`;
+        await redis.set(key, count);
+        logger.debug("set story count for counts v2 in redis cache", {
+          storyID: missingID,
+          count,
+        });
+
+        countResults.set(missingID, { storyID: missingID, count });
+      }
+
+      const results: Array<CountResult | null> = [];
+      for (const storyID of storyIDs) {
+        const value = countResults.get(storyID) ?? null;
+        results.push(value);
+      }
+
+      res.send(JSON.stringify(results));
     } catch (err) {
       return next(err);
     }