WIL realtime concurrency tracking

ericallam · ericallam · commit 259c9053853a · 2024-10-14T17:17:15.000+01:00
diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug/route.tsx
@@ -11,7 +11,7 @@ import { useTypedMatchesData } from "~/hooks/useTypedMatchData";
 import { useUser } from "~/hooks/useUser";
 import { OrganizationsPresenter } from "~/presenters/OrganizationsPresenter.server";
 import { getImpersonationId } from "~/services/impersonation.server";
-import { getCurrentPlan, getUsage } from "~/services/platform.v3.server";
+import { getCachedUsage, getCurrentPlan, getUsage } from "~/services/platform.v3.server";
 import { requireUserId } from "~/services/session.server";
 import { telemetry } from "~/services/telemetry.server";
 import { organizationPath } from "~/utils/pathBuilder";
@@ -30,6 +30,24 @@ export function useCurrentPlan(matches?: UIMatch[]) {
   return data?.currentPlan;
 }
 
+export const shouldRevalidate: ShouldRevalidateFunction = (params) => {
+  const { currentParams, nextParams } = params;
+
+  const current = ParamsSchema.safeParse(currentParams);
+  const next = ParamsSchema.safeParse(nextParams);
+
+  if (current.success && next.success) {
+    if (current.data.organizationSlug !== next.data.organizationSlug) {
+      return true;
+    }
+    if (current.data.projectParam !== next.data.projectParam) {
+      return true;
+    }
+  }
+
+  return params.currentUrl.pathname !== params.nextUrl.pathname;
+};
+
 export const loader = async ({ request, params }: LoaderFunctionArgs) => {
   const userId = await requireUserId(request);
   const impersonationId = await getImpersonationId(request);
@@ -51,11 +69,14 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
   const firstDayOfMonth = new Date();
   firstDayOfMonth.setUTCDate(1);
   firstDayOfMonth.setUTCHours(0, 0, 0, 0);
-  const tomorrow = new Date();
-  tomorrow.setUTCDate(tomorrow.getDate() + 1);
+  const firstDayOfNextMonth = new Date();
+  firstDayOfNextMonth.setUTCMonth(firstDayOfNextMonth.getUTCMonth() + 1);
+  firstDayOfNextMonth.setUTCDate(1);
+  firstDayOfNextMonth.setUTCHours(0, 0, 0, 0);
+
   const [plan, usage] = await Promise.all([
     getCurrentPlan(organization.id),
-    getUsage(organization.id, { from: firstDayOfMonth, to: tomorrow }),
+    getCachedUsage(organization.id, { from: firstDayOfMonth, to: firstDayOfNextMonth }),
   ]);
 
   let hasExceededFreeTier = false;
@@ -104,23 +125,3 @@ export function ErrorBoundary() {
     <RouteErrorDisplay button={{ title: "Home", to: "/" }} />
   );
 }
-
-export const shouldRevalidate: ShouldRevalidateFunction = ({
-  defaultShouldRevalidate,
-  currentParams,
-  nextParams,
-}) => {
-  const current = ParamsSchema.safeParse(currentParams);
-  const next = ParamsSchema.safeParse(nextParams);
-
-  if (current.success && next.success) {
-    if (current.data.organizationSlug !== next.data.organizationSlug) {
-      return true;
-    }
-    if (current.data.projectParam !== next.data.projectParam) {
-      return true;
-    }
-  }
-
-  return defaultShouldRevalidate;
-};
diff --git a/apps/webapp/app/routes/realtime.v1.batches.$batchId.ts b/apps/webapp/app/routes/realtime.v1.batches.$batchId.ts
@@ -1,12 +1,11 @@
 import { type ActionFunctionArgs, json } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { $replica } from "~/db.server";
-import { env } from "~/env.server";
 import { permittedToReadBatch } from "~/services/accessControl.server";
 import { authenticateApiRequest } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
+import { realtimeClient } from "~/services/realtimeClient.server";
 import { makeApiCors } from "~/utils/apiCors";
-import { longPollingFetch } from "~/utils/longPollingFetch";
 
 const ParamsSchema = z.object({
   batchId: z.string(),
@@ -53,17 +52,11 @@ export async function loader({ request, params }: ActionFunctionArgs) {
       return apiCors(json({ error: "Batch Run not found" }, { status: 404 }));
     }
 
-    const url = new URL(request.url);
-    const originUrl = new URL(`${env.ELECTRIC_ORIGIN}/v1/shape/public."TaskRun"`);
-    url.searchParams.forEach((value, key) => {
-      originUrl.searchParams.set(key, value);
-    });
-
-    originUrl.searchParams.set("where", `"batchId"='${batchRun.id}'`);
-
-    const finalUrl = originUrl.toString();
-
-    return longPollingFetch(finalUrl);
+    return realtimeClient.streamRunsWhere(
+      request.url,
+      authenticationResult.environment,
+      `"batchId"='${batchRun.id}'`
+    );
   } catch (error) {
     if (error instanceof Response) {
       // Error responses from longPollingFetch
diff --git a/apps/webapp/app/routes/realtime.v1.runs.$runId.ts b/apps/webapp/app/routes/realtime.v1.runs.$runId.ts
@@ -1,12 +1,11 @@
 import { type ActionFunctionArgs, json } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { $replica } from "~/db.server";
-import { env } from "~/env.server";
 import { permittedToReadRun } from "~/services/accessControl.server";
 import { authenticateApiRequest } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
+import { realtimeClient } from "~/services/realtimeClient.server";
 import { makeApiCors } from "~/utils/apiCors";
-import { longPollingFetch } from "~/utils/longPollingFetch";
 
 const ParamsSchema = z.object({
   runId: z.string(),
@@ -53,17 +52,12 @@ export async function loader({ request, params }: ActionFunctionArgs) {
       return apiCors(json({ error: "Task Run not found" }, { status: 404 }));
     }
 
-    const url = new URL(request.url);
-    const originUrl = new URL(`${env.ELECTRIC_ORIGIN}/v1/shape/public."TaskRun"`);
-    url.searchParams.forEach((value, key) => {
-      originUrl.searchParams.set(key, value);
-    });
-
-    originUrl.searchParams.set("where", `"id"='${run.id}'`);
-
-    const finalUrl = originUrl.toString();
-
-    return longPollingFetch(finalUrl);
+    return realtimeClient.streamRunsWhere(
+      request.url,
+      authenticationResult.environment,
+      `"id"='${run.id}'`,
+      apiCors
+    );
   } catch (error) {
     if (error instanceof Response) {
       // Error responses from longPollingFetch
diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts
@@ -1,12 +1,12 @@
-import { DefaultStatefulContext, Namespace, Cache as UnkeyCache, createCache } from "@unkey/cache";
+import { createCache, DefaultStatefulContext, Namespace, Cache as UnkeyCache } from "@unkey/cache";
 import { MemoryStore } from "@unkey/cache/stores";
 import { Ratelimit } from "@upstash/ratelimit";
 import { Request as ExpressRequest, Response as ExpressResponse, NextFunction } from "express";
 import { RedisOptions } from "ioredis";
 import { createHash } from "node:crypto";
 import { z } from "zod";
 import { env } from "~/env.server";
-import { authenticateApiKey, authenticateAuthorizationHeader } from "./apiAuth.server";
+import { authenticateAuthorizationHeader } from "./apiAuth.server";
 import { logger } from "./logger.server";
 import { createRedisRateLimitClient, Duration, RateLimiter } from "./rateLimiter.server";
 import { RedisCacheStore } from "./unkey/redisCacheStore.server";
@@ -153,6 +153,7 @@ export function authorizationRateLimitMiddleware({
     },
   });
 
+  // This cache holds the rate limit configuration for each org, so we don't have to fetch it every request
   const cache = createCache({
     limiter: new Namespace<RateLimiterConfig>(ctx, {
       stores: [memory, redisCacheStore],
diff --git a/apps/webapp/app/services/platform.v3.server.ts b/apps/webapp/app/services/platform.v3.server.ts
@@ -1,16 +1,75 @@
-import { BillingClient, Limits, SetPlanBody, UsageSeriesParams } from "@trigger.dev/platform/v3";
 import { Organization, Project } from "@trigger.dev/database";
+import {
+  BillingClient,
+  Limits,
+  SetPlanBody,
+  UsageSeriesParams,
+  UsageResult,
+} from "@trigger.dev/platform/v3";
+import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
+import { MemoryStore } from "@unkey/cache/stores";
 import { redirect } from "remix-typedjson";
 import { $replica } from "~/db.server";
 import { env } from "~/env.server";
 import { redirectWithErrorMessage, redirectWithSuccessMessage } from "~/models/message.server";
 import { createEnvironment } from "~/models/organization.server";
 import { logger } from "~/services/logger.server";
 import { newProjectPath, organizationBillingPath } from "~/utils/pathBuilder";
+import { singleton } from "~/utils/singleton";
+import { RedisCacheStore } from "./unkey/redisCacheStore.server";
+
+function initializeClient() {
+  if (isCloud() && process.env.BILLING_API_URL && process.env.BILLING_API_KEY) {
+    const client = new BillingClient({
+      url: process.env.BILLING_API_URL,
+      apiKey: process.env.BILLING_API_KEY,
+    });
+    console.log(`🤑 Billing client initialized: ${process.env.BILLING_API_URL}`);
+    return client;
+  } else {
+    console.log(`🤑 Billing client not initialized`);
+  }
+}
+
+const client = singleton("billingClient", initializeClient);
+
+function initializePlatformCache() {
+  const ctx = new DefaultStatefulContext();
+  const memory = new MemoryStore({ persistentMap: new Map() });
+  const redisCacheStore = new RedisCacheStore({
+    connection: {
+      keyPrefix: `cache:platform:v3:`,
+      port: env.REDIS_PORT,
+      host: env.REDIS_HOST,
+      username: env.REDIS_USERNAME,
+      password: env.REDIS_PASSWORD,
+      enableAutoPipelining: true,
+      ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+    },
+  });
+
+  // This cache holds the limits fetched from the platform service
+  const cache = createCache({
+    limits: new Namespace<number>(ctx, {
+      stores: [memory, redisCacheStore],
+      fresh: 60_000 * 5, // 5 minutes
+      stale: 60_000 * 10, // 10 minutes
+    }),
+    usage: new Namespace<UsageResult>(ctx, {
+      stores: [memory, redisCacheStore],
+      fresh: 60_000 * 5, // 5 minutes
+      stale: 60_000 * 10, // 10 minutes
+    }),
+  });
+
+  return cache;
+}
+
+const platformCache = singleton("platformCache", initializePlatformCache);
 
 export async function getCurrentPlan(orgId: string) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.currentPlan(orgId);
 
@@ -60,8 +119,8 @@ export async function getCurrentPlan(orgId: string) {
 }
 
 export async function getLimits(orgId: string) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.currentPlan(orgId);
     if (!result.success) {
@@ -87,9 +146,15 @@ export async function getLimit(orgId: string, limit: keyof Limits, fallback: num
   return fallback;
 }
 
+export async function getCachedLimit(orgId: string, limit: keyof Limits, fallback: number) {
+  return platformCache.limits.swr(`${orgId}:${limit}`, async () => {
+    return getLimit(orgId, limit, fallback);
+  });
+}
+
 export async function customerPortalUrl(orgId: string, orgSlug: string) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     return client.createPortalSession(orgId, {
       returnUrl: `${env.APP_ORIGIN}${organizationBillingPath({ slug: orgSlug })}`,
@@ -101,8 +166,8 @@ export async function customerPortalUrl(orgId: string, orgSlug: string) {
 }
 
 export async function getPlans() {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.plans();
     if (!result.success) {
@@ -122,7 +187,6 @@ export async function setPlan(
   callerPath: string,
   plan: SetPlanBody
 ) {
-  const client = getClient();
   if (!client) {
     throw redirectWithErrorMessage(callerPath, request, "Error setting plan");
   }
@@ -178,8 +242,8 @@ export async function setPlan(
 }
 
 export async function getUsage(organizationId: string, { from, to }: { from: Date; to: Date }) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.usage(organizationId, { from, to });
     if (!result.success) {
@@ -193,9 +257,27 @@ export async function getUsage(organizationId: string, { from, to }: { from: Dat
   }
 }
 
+export async function getCachedUsage(
+  organizationId: string,
+  { from, to }: { from: Date; to: Date }
+) {
+  if (!client) return undefined;
+
+  const result = await platformCache.usage.swr(
+    `${organizationId}:${from.toISOString()}:${to.toISOString()}`,
+    async () => {
+      const usageResponse = await getUsage(organizationId, { from, to });
+
+      return usageResponse;
+    }
+  );
+
+  return result.val;
+}
+
 export async function getUsageSeries(organizationId: string, params: UsageSeriesParams) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.usageSeries(organizationId, params);
     if (!result.success) {
@@ -214,8 +296,8 @@ export async function reportInvocationUsage(
   costInCents: number,
   additionalData?: Record<string, any>
 ) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.reportInvocationUsage({
       organizationId,
@@ -234,8 +316,8 @@ export async function reportInvocationUsage(
 }
 
 export async function reportComputeUsage(request: Request) {
-  const client = getClient();
   if (!client) return undefined;
+
   return fetch(`${process.env.BILLING_API_URL}/api/v1/usage/ingest/compute`, {
     method: "POST",
     headers: request.headers,
@@ -244,8 +326,8 @@ export async function reportComputeUsage(request: Request) {
 }
 
 export async function getEntitlement(organizationId: string) {
-  const client = getClient();
   if (!client) return undefined;
+
   try {
     const result = await client.getEntitlement(organizationId);
     if (!result.success) {
@@ -275,19 +357,6 @@ export async function projectCreated(organization: Organization, project: Projec
   }
 }
 
-function getClient() {
-  if (isCloud() && process.env.BILLING_API_URL && process.env.BILLING_API_KEY) {
-    const client = new BillingClient({
-      url: process.env.BILLING_API_URL,
-      apiKey: process.env.BILLING_API_KEY,
-    });
-    console.log(`Billing client initialized: ${process.env.BILLING_API_URL}`);
-    return client;
-  } else {
-    console.log(`Billing client not initialized`);
-  }
-}
-
 function isCloud(): boolean {
   const acceptableHosts = [
     "https://cloud.trigger.dev",
diff --git a/apps/webapp/app/services/realtimeClient.server.ts b/apps/webapp/app/services/realtimeClient.server.ts
diff --git a/references/nextjs-realtime/next.config.mjs b/references/nextjs-realtime/next.config.mjs
