reduce pg pressure for dedupe tasks

Author: n1ru4l

packages/services/workflows/package.json

@@ -13,6 +13,7 @@
     "@hive/service-common": "workspace:*",
     "@hive/storage": "workspace:*",
     "@sentry/node": "7.120.2",
+    "bentocache": "1.1.0",
     "dotenv": "16.4.7",
     "graphile-worker": "0.16.6",
     "nodemailer": "7.0.11",

packages/services/workflows/src/kit.ts

@@ -1,3 +1,5 @@
+import { BentoCache, bentostore } from 'bentocache';
+import { memoryDriver } from 'bentocache/build/src/drivers/memory';
 import { makeWorkerUtils, WorkerUtils, type JobHelpers, type Task } from 'graphile-worker';
 import type { Pool } from 'pg';
 import { z } from 'zod';
@@ -71,6 +73,8 @@ export function implementTask<TPayload>(
  */
 export class TaskScheduler {
   tools: Promise<WorkerUtils>;
+  cache: BentoCache<{ store: ReturnType<typeof bentostore> }>;
+
   constructor(
     pgPool: Pool,
     private logger: Logger = new Logger(),
@@ -79,6 +83,17 @@ export class TaskScheduler {
       pgPool,
       logger: bridgeGraphileLogger(logger),
     });
+    this.cache = new BentoCache({
+      default: 'taskSchedule',
+      stores: {
+        taskSchedule: bentostore().useL1Layer(
+          memoryDriver({
+            maxItems: 10_000,
+            prefix: 'bentocache:graphile_worker_deduplication',
+          }),
+        ),
+      },
+    });
   }
 
   async scheduleTask<TPayload>(
@@ -111,21 +126,27 @@ export class TaskScheduler {
         typeof opts.dedupe.key === 'string' ? opts.dedupe.key : opts.dedupe.key(payload);
       const expiresAt = new Date(new Date().getTime() + opts.dedupe.ttl).toISOString();
 
-      const shouldSkip = await tools.withPgClient(async client => {
-        const result = await client.query(
-          `
-            INSERT INTO "graphile_worker_deduplication" ("task_name", "dedupe_key", "expires_at")
-            VALUES($1, $2, $3)
-            ON CONFLICT ("task_name", "dedupe_key")
-            DO
-              UPDATE SET "expires_at" = EXCLUDED.expires_at
-              WHERE "graphile_worker_deduplication"."expires_at" < NOW()
-            RETURNING xmax = 0 AS "inserted"
-          `,
-          [taskDefinition.name, dedupeKey, expiresAt],
-        );
-
-        return result.rows.length === 0;
+      const shouldSkip = await this.cache.getOrSet({
+        key: `${taskDefinition.name}:${dedupeKey}`,
+        ttl: opts.dedupe.ttl,
+        async factory() {
+          return await tools.withPgClient(async client => {
+            const result = await client.query(
+              `
+               INSERT INTO "graphile_worker_deduplication" ("task_name", "dedupe_key", "expires_at")
+               VALUES($1, $2, $3)
+               ON CONFLICT ("task_name", "dedupe_key")
+               DO
+                 UPDATE SET "expires_at" = EXCLUDED.expires_at
+                 WHERE "graphile_worker_deduplication"."expires_at" < NOW()
+               RETURNING xmax = 0 AS "inserted"
+             `,
+              [taskDefinition.name, dedupeKey, expiresAt],
+            );
+
+            return result.rows.length === 0;
+          });
+        },
       });
 
       if (shouldSkip) {