add cache with best effort invalidation

Author: nicktrn

apps/webapp/app/services/platform.v3.server.ts

@@ -21,6 +21,7 @@ import { logger } from "~/services/logger.server";
 import { newProjectPath, organizationBillingPath } from "~/utils/pathBuilder";
 import { singleton } from "~/utils/singleton";
 import { RedisCacheStore } from "./unkey/redisCacheStore.server";
+import { engine } from "~/v3/runEngine.server";
 import { existsSync, readFileSync } from "node:fs";
 import { z } from "zod";
 import { MachinePresetName } from "@trigger.dev/core/v3";
@@ -308,6 +309,8 @@ export async function setPlan(
       }
       case "free_connected": {
         if (result.accepted) {
+          // Invalidate billing cache since plan changed
+          engine.invalidateBillingCache(organization.id);
           return redirect(newProjectPath(organization, "You're on the Free plan."));
         } else {
           return redirectWithErrorMessage(
@@ -321,13 +324,17 @@ export async function setPlan(
         return redirect(result.checkoutUrl);
       }
       case "updated_subscription": {
+        // Invalidate billing cache since subscription changed
+        engine.invalidateBillingCache(organization.id);
         return redirectWithSuccessMessage(
           callerPath,
           request,
           "Subscription updated successfully."
         );
       }
       case "canceled_subscription": {
+        // Invalidate billing cache since subscription was canceled
+        engine.invalidateBillingCache(organization.id);
         return redirectWithSuccessMessage(callerPath, request, "Subscription canceled.");
       }
     }

internal-packages/run-engine/src/engine/index.ts

@@ -298,6 +298,7 @@ export class RunEngine {
       runAttemptSystem: this.runAttemptSystem,
       machines: this.options.machines,
       billing: this.options.billing,
+      redisOptions: this.options.cache?.redis ?? this.options.runLock.redis,
     });
   }
 
@@ -1348,4 +1349,12 @@ export class RunEngine {
         orgId: run.organizationId!,
       }));
   }
+
+  /**
+   * Invalidates the billing cache for an organization when their plan changes
+   * Runs in background and handles all errors internally
+   */
+  invalidateBillingCache(orgId: string): void {
+    this.dequeueSystem.invalidateBillingCache(orgId);
+  }
 }

internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -1,5 +1,14 @@
+import {
+  createCache,
+  DefaultStatefulContext,
+  MemoryStore,
+  Namespace,
+  RedisCacheStore,
+  UnkeyCache,
+} from "@internal/cache";
+import type { RedisOptions } from "@internal/redis";
 import { startSpan } from "@internal/tracing";
-import { assertExhaustive, tryCatch } from "@trigger.dev/core";
+import { assertExhaustive } from "@trigger.dev/core";
 import { DequeuedMessage, RetryOptions } from "@trigger.dev/core/v3";
 import { getMaxDuration } from "@trigger.dev/core/v3/isomorphic";
 import { PrismaClientOrTransaction } from "@trigger.dev/database";
@@ -11,24 +20,66 @@ import { RunEngineOptions } from "../types.js";
 import { ExecutionSnapshotSystem, getLatestExecutionSnapshot } from "./executionSnapshotSystem.js";
 import { RunAttemptSystem } from "./runAttemptSystem.js";
 import { SystemResources } from "./systems.js";
+import { ServiceValidationError } from "../errors.js";
 
 export type DequeueSystemOptions = {
   resources: SystemResources;
   machines: RunEngineOptions["machines"];
   executionSnapshotSystem: ExecutionSnapshotSystem;
   runAttemptSystem: RunAttemptSystem;
   billing?: RunEngineOptions["billing"];
+  redisOptions: RedisOptions;
 };
 
+// Cache TTLs for billing information - shorter than other caches since billing can change
+const BILLING_FRESH_TTL = 60000 * 5; // 5 minutes
+const BILLING_STALE_TTL = 60000 * 10; // 10 minutes
+
 export class DequeueSystem {
   private readonly $: SystemResources;
   private readonly executionSnapshotSystem: ExecutionSnapshotSystem;
   private readonly runAttemptSystem: RunAttemptSystem;
+  private readonly billingCache: UnkeyCache<{
+    billing: { isPaying: boolean };
+  }>;
 
   constructor(private readonly options: DequeueSystemOptions) {
     this.$ = options.resources;
     this.executionSnapshotSystem = options.executionSnapshotSystem;
     this.runAttemptSystem = options.runAttemptSystem;
+
+    // Initialize billing cache
+    const ctx = new DefaultStatefulContext();
+    const memory = new MemoryStore({ persistentMap: new Map() });
+    const redisCacheStore = new RedisCacheStore({
+      name: "dequeue-system",
+      connection: {
+        ...options.redisOptions,
+        keyPrefix: "engine:dequeue-system:cache:",
+      },
+      useModernCacheKeyBuilder: true,
+    });
+
+    this.billingCache = createCache({
+      billing: new Namespace<{ isPaying: boolean }>(ctx, {
+        stores: [memory, redisCacheStore],
+        fresh: BILLING_FRESH_TTL,
+        stale: BILLING_STALE_TTL,
+      }),
+    });
+  }
+
+  /**
+   * Invalidates the billing cache for an organization when their plan changes
+   * Runs in background and handles all errors internally
+   */
+  invalidateBillingCache(orgId: string): void {
+    this.billingCache.billing.remove(orgId).catch((error) => {
+      this.$.logger.warn("Failed to invalidate billing cache", {
+        orgId,
+        error: error.message,
+      });
+    });
   }
 
   /**
@@ -382,22 +433,7 @@ export class DequeueSystem {
               const nextAttemptNumber = currentAttemptNumber + 1;
 
               // Get billing information if available
-              let isPaying = false;
-              if (this.options.billing?.getCurrentPlan) {
-                const [error, planResult] = await tryCatch(
-                  this.options.billing.getCurrentPlan(orgId)
-                );
-
-                if (error) {
-                  this.$.logger.error("Failed to get billing information", {
-                    orgId,
-                    runId,
-                    error: error.message,
-                  });
-                } else {
-                  isPaying = planResult.isPaying;
-                }
-              }
+              const billing = await this.#getBillingInfo({ orgId, runId });
 
               const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot(
                 prisma,
@@ -467,11 +503,7 @@ export class DequeueSystem {
                 project: {
                   id: lockedTaskRun.projectId,
                 },
-                billing: {
-                  currentPlan: {
-                    isPaying,
-                  },
-                },
+                billing,
               } satisfies DequeuedMessage;
             }
           );
@@ -636,4 +668,37 @@ export class DequeueSystem {
       });
     });
   }
+
+  async #getBillingInfo({
+    orgId,
+    runId,
+  }: {
+    orgId: string;
+    runId: string;
+  }): Promise<{ currentPlan: { isPaying: boolean } }> {
+    if (!this.options.billing?.getCurrentPlan) {
+      return { currentPlan: { isPaying: false } };
+    }
+
+    const result = await this.billingCache.billing.swr(orgId, async () => {
+      // This is safe because options can't change at runtime
+      const planResult = await this.options.billing!.getCurrentPlan(orgId);
+
+      return { isPaying: planResult.isPaying };
+    });
+
+    if (result.err) {
+      throw result.err;
+    }
+
+    if (!result.val) {
+      throw new ServiceValidationError(
+        `Could not resolve billing information for organization ${orgId}`,
+        undefined,
+        { orgId, runId }
+      );
+    }
+
+    return { currentPlan: { isPaying: result.val.isPaying } };
+  }
 }