Implement test for realtime client using testcontainers

also updated electric to latest version

Author: ericallam

.vscode/extensions.json

@@ -1,8 +1,4 @@
 {
-	"recommendations": [
-        "denoland.vscode-deno"
-    ],
-	"unwantedRecommendations": [
-		
-	]
+  "recommendations": ["bierner.comment-tagged-templates"],
+  "unwantedRecommendations": []
 }

apps/webapp/app/routes/_app.orgs.$organizationSlug/route.tsx

@@ -45,9 +45,12 @@ export const shouldRevalidate: ShouldRevalidateFunction = (params) => {
     }
   }
 
+  // This prevents revalidation when there are search params changes
+  // IMPORTANT: If the loader function depends on search params, this should be updated
   return params.currentUrl.pathname !== params.nextUrl.pathname;
 };
 
+// IMPORTANT: Make sure to update shouldRevalidate if this loader depends on search params
 export const loader = async ({ request, params }: LoaderFunctionArgs) => {
   const userId = await requireUserId(request);
   const impersonationId = await getImpersonationId(request);
@@ -69,6 +72,9 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
   const firstDayOfMonth = new Date();
   firstDayOfMonth.setUTCDate(1);
   firstDayOfMonth.setUTCHours(0, 0, 0, 0);
+
+  // Using the 1st day of next month means we get the usage for the current month
+  // and the cache key for getCachedUsage is stable over the month
   const firstDayOfNextMonth = new Date();
   firstDayOfNextMonth.setUTCMonth(firstDayOfNextMonth.getUTCMonth() + 1);
   firstDayOfNextMonth.setUTCDate(1);

apps/webapp/app/routes/realtime.v1.batches.$batchId.ts

@@ -4,7 +4,7 @@ import { $replica } from "~/db.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 { realtimeClient } from "~/services/realtimeClientGlobal.server";
 import { makeApiCors } from "~/utils/apiCors";
 
 const ParamsSchema = z.object({

apps/webapp/app/routes/realtime.v1.runs.$runId.ts

@@ -4,7 +4,7 @@ import { $replica } from "~/db.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 { realtimeClient } from "~/services/realtimeClientGlobal.server";
 import { makeApiCors } from "~/utils/apiCors";
 
 const ParamsSchema = z.object({

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

@@ -38,7 +38,7 @@ function initializePlatformCache() {
   const memory = new MemoryStore({ persistentMap: new Map() });
   const redisCacheStore = new RedisCacheStore({
     connection: {
-      keyPrefix: `cache:platform:v3:`,
+      keyPrefix: "tr:cache:platform:v3",
       port: env.REDIS_PORT,
       host: env.REDIS_HOST,
       username: env.REDIS_USERNAME,

apps/webapp/app/services/realtimeClient.server.ts

@@ -1,38 +1,47 @@
+import { json } from "@remix-run/server-runtime";
 import Redis, { Callback, Result, type RedisOptions } from "ioredis";
+import { randomUUID } from "node:crypto";
 import { longPollingFetch } from "~/utils/longPollingFetch";
-import { getCachedLimit } from "./platform.v3.server";
 import { logger } from "./logger.server";
-import { AuthenticatedEnvironment } from "./apiAuth.server";
-import { json } from "@remix-run/server-runtime";
-import { env } from "~/env.server";
-import { singleton } from "~/utils/singleton";
+
+export interface CachedLimitProvider {
+  getCachedLimit: (organizationId: string, defaultValue: number) => Promise<number | undefined>;
+}
 
 export type RealtimeClientOptions = {
   electricOrigin: string;
   redis: RedisOptions;
+  cachedLimitProvider: CachedLimitProvider;
   keyPrefix: string;
-  expiryTime?: number;
+  expiryTimeInSeconds?: number;
+};
+
+export type RealtimeEnvironment = {
+  id: string;
+  organizationId: string;
 };
 
 export class RealtimeClient {
   private redis: Redis;
-  private expiryTime: number;
+  private expiryTimeInSeconds: number;
+  private cachedLimitProvider: CachedLimitProvider;
 
   constructor(private options: RealtimeClientOptions) {
     this.redis = new Redis(options.redis);
-    this.expiryTime = options.expiryTime ?? 3600; // default to 1 hour
+    this.expiryTimeInSeconds = options.expiryTimeInSeconds ?? 60 * 5; // default to 5 minutes
+    this.cachedLimitProvider = options.cachedLimitProvider;
     this.#registerCommands();
   }
 
   async streamRunsWhere(
     url: URL | string,
-    authenticatedEnv: AuthenticatedEnvironment,
+    environment: RealtimeEnvironment,
     whereClause: string,
     responseWrapper?: (response: Response) => Promise<Response>
   ) {
     const electricUrl = this.#constructElectricUrl(url, whereClause);
 
-    return this.#performElectricRequest(electricUrl, authenticatedEnv, responseWrapper);
+    return this.#performElectricRequest(electricUrl, environment, responseWrapper);
   }
 
   #constructElectricUrl(url: URL | string, whereClause: string): URL {
@@ -51,7 +60,7 @@ export class RealtimeClient {
 
   async #performElectricRequest(
     url: URL,
-    authenticatedEnv: AuthenticatedEnvironment,
+    environment: RealtimeEnvironment,
     responseWrapper: (response: Response) => Promise<Response> = (r) => Promise.resolve(r)
   ) {
     const shapeId = extractShapeId(url);
@@ -67,40 +76,40 @@ export class RealtimeClient {
       return longPollingFetch(url.toString());
     }
 
+    const requestId = randomUUID();
+
     // We now need to wrap the longPollingFetch in a concurrency tracker
-    const concurrencyLimitResult = await getCachedLimit(
-      authenticatedEnv.organizationId,
-      "realtimeConcurrentConnections",
+    const concurrencyLimit = await this.cachedLimitProvider.getCachedLimit(
+      environment.organizationId,
       100_000
     );
 
-    if (!concurrencyLimitResult.val) {
+    if (!concurrencyLimit) {
       logger.error("Failed to get concurrency limit", {
-        organizationId: authenticatedEnv.organizationId,
-        concurrencyLimitResult,
+        organizationId: environment.organizationId,
       });
 
       return responseWrapper(json({ error: "Failed to get concurrency limit" }, { status: 500 }));
     }
 
-    const concurrencyLimit = concurrencyLimitResult.val;
-
     logger.debug("[realtimeClient] increment and check", {
       concurrencyLimit,
       shapeId,
-      authenticatedEnv: {
-        id: authenticatedEnv.id,
-        organizationId: authenticatedEnv.organizationId,
+      requestId,
+      environment: {
+        id: environment.id,
+        organizationId: environment.organizationId,
       },
     });
 
-    const canProceed = await this.#incrementAndCheck(
-      authenticatedEnv.id,
-      shapeId,
-      concurrencyLimit
-    );
+    const canProceed = await this.#incrementAndCheck(environment.id, requestId, concurrencyLimit);
 
     if (!canProceed) {
+      logger.debug("[realtimeClient] too many concurrent requests", {
+        requestId,
+        environmentId: environment.id,
+      });
+
       return responseWrapper(json({ error: "Too many concurrent requests" }, { status: 429 }));
     }
 
@@ -109,41 +118,42 @@ export class RealtimeClient {
       const response = await longPollingFetch(url.toString());
 
       // Decrement the counter after the long polling request is complete
-      await this.#decrementConcurrency(authenticatedEnv.id, shapeId);
+      await this.#decrementConcurrency(environment.id, requestId);
 
       return response;
     } catch (error) {
       // Decrement the counter if the request fails
-      await this.#decrementConcurrency(authenticatedEnv.id, shapeId);
+      await this.#decrementConcurrency(environment.id, requestId);
 
       throw error;
     }
   }
 
-  async #incrementAndCheck(environmentId: string, shapeId: string, limit: number) {
+  async #incrementAndCheck(environmentId: string, requestId: string, limit: number) {
     const key = this.#getKey(environmentId);
-    const now = Date.now().toString();
+    const now = Date.now();
 
     const result = await this.redis.incrementAndCheckConcurrency(
       key,
-      now,
-      shapeId,
-      this.expiryTime.toString(),
+      now.toString(),
+      requestId,
+      this.expiryTimeInSeconds.toString(), // expiry time
+      (now - this.expiryTimeInSeconds * 1000).toString(), // cutoff time
       limit.toString()
     );
 
     return result === 1;
   }
 
-  async #decrementConcurrency(environmentId: string, shapeId: string) {
+  async #decrementConcurrency(environmentId: string, requestId: string) {
     logger.debug("[realtimeClient] decrement", {
-      shapeId,
+      requestId,
       environmentId,
     });
 
     const key = this.#getKey(environmentId);
 
-    await this.redis.zrem(key, shapeId);
+    await this.redis.zrem(key, requestId);
   }
 
   #getKey(environmentId: string): string {
@@ -153,13 +163,17 @@ export class RealtimeClient {
   #registerCommands() {
     this.redis.defineCommand("incrementAndCheckConcurrency", {
       numberOfKeys: 1,
-      lua: `
+      lua: /* lua */ `
         local concurrencyKey = KEYS[1]
 
-        local timestamp = ARGV[1]
+        local timestamp = tonumber(ARGV[1])
         local requestId = ARGV[2]
-        local expiryTime = ARGV[3]
-        local limit = tonumber(ARGV[4])
+        local expiryTime = tonumber(ARGV[3])
+        local cutoffTime = tonumber(ARGV[4])
+        local limit = tonumber(ARGV[5])
+
+        -- Remove expired entries
+        redis.call('ZREMRANGEBYSCORE', concurrencyKey, '-inf', cutoffTime)
 
         -- Add the new request to the sorted set
         redis.call('ZADD', concurrencyKey, timestamp, requestId)
@@ -199,25 +213,9 @@ declare module "ioredis" {
       timestamp: string,
       requestId: string,
       expiryTime: string,
+      cutoffTime: string,
       limit: string,
       callback?: Callback<number>
     ): Result<number, Context>;
   }
 }
-
-function initializeRealtimeClient() {
-  return new RealtimeClient({
-    electricOrigin: env.ELECTRIC_ORIGIN,
-    keyPrefix: `tr:realtime:concurrency`,
-    redis: {
-      port: env.REDIS_PORT,
-      host: env.REDIS_HOST,
-      username: env.REDIS_USERNAME,
-      password: env.REDIS_PASSWORD,
-      enableAutoPipelining: true,
-      ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
-    },
-  });
-}
-
-export const realtimeClient = singleton("realtimeClient", initializeRealtimeClient);

apps/webapp/app/services/realtimeClientGlobal.server.ts

@@ -0,0 +1,32 @@
+import { env } from "~/env.server";
+import { singleton } from "~/utils/singleton";
+import { RealtimeClient } from "./realtimeClient.server";
+import { getCachedLimit } from "./platform.v3.server";
+
+function initializeRealtimeClient() {
+  return new RealtimeClient({
+    electricOrigin: env.ELECTRIC_ORIGIN,
+    keyPrefix: "tr:realtime:concurrency",
+    redis: {
+      port: env.REDIS_PORT,
+      host: env.REDIS_HOST,
+      username: env.REDIS_USERNAME,
+      password: env.REDIS_PASSWORD,
+      enableAutoPipelining: true,
+      ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+    },
+    cachedLimitProvider: {
+      async getCachedLimit(organizationId, defaultValue) {
+        const result = await getCachedLimit(
+          organizationId,
+          "realtimeConcurrentConnections",
+          defaultValue
+        );
+
+        return result.val;
+      },
+    },
+  });
+}
+
+export const realtimeClient = singleton("realtimeClient", initializeRealtimeClient);

apps/webapp/package.json

@@ -185,6 +185,7 @@
     "zod-validation-error": "^1.5.0"
   },
   "devDependencies": {
+    "@internal/testcontainers": "workspace:*",
     "@remix-run/dev": "2.1.0",
     "@remix-run/eslint-config": "2.1.0",
     "@remix-run/testing": "^2.1.0",
@@ -243,9 +244,10 @@
     "ts-node": "^10.7.0",
     "tsconfig-paths": "^3.14.1",
     "typescript": "^5.1.6",
+    "vite-tsconfig-paths": "^4.0.5",
     "vitest": "^1.4.0"
   },
   "engines": {
     "node": ">=16.0.0"
   }
-}
+}