Add sharding support for electric

Author: ericallam

apps/webapp/app/env.server.ts

@@ -35,6 +35,9 @@ const EnvironmentSchema = z.object({
   API_ORIGIN: z.string().optional(),
   STREAM_ORIGIN: z.string().optional(),
   ELECTRIC_ORIGIN: z.string().default("http://localhost:3060"),
+  // A comma separated list of electric origins to shard into different electric instances by environmentId
+  // example: "http://localhost:3060,http://localhost:3061,http://localhost:3062"
+  ELECTRIC_ORIGIN_SHARDS: z.string().optional(),
   APP_ENV: z.string().default(process.env.NODE_ENV),
   SERVICE_NAME: z.string().default("trigger.dev webapp"),
   POSTHOG_PROJECT_KEY: z.string().default("phc_LFH7kJiGhdIlnO22hTAKgHpaKhpM8gkzWAFvHmf5vfS"),

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

@@ -6,6 +6,7 @@ import { randomUUID } from "node:crypto";
 import { createRedisClient, RedisClient, RedisWithClusterOptions } from "~/redis.server";
 import { longPollingFetch } from "~/utils/longPollingFetch";
 import { logger } from "./logger.server";
+import { jumpHash } from "@trigger.dev/core/v3/serverOnly";
 
 export interface CachedLimitProvider {
   getCachedLimit: (organizationId: string, defaultValue: number) => Promise<number | undefined>;
@@ -45,7 +46,7 @@ const RESERVED_COLUMNS = ["id", "taskIdentifier", "friendlyId", "status", "creat
 const RESERVED_SEARCH_PARAMS = ["createdAt", "tags", "skipColumns"];
 
 export type RealtimeClientOptions = {
-  electricOrigin: string;
+  electricOrigin: string | string[];
   redis: RedisWithClusterOptions;
   cachedLimitProvider: CachedLimitProvider;
   keyPrefix: string;
@@ -221,15 +222,26 @@ export class RealtimeClient {
     whereClause: string,
     clientVersion?: string
   ) {
-    const electricUrl = this.#constructRunsElectricUrl(url, whereClause, clientVersion);
+    const electricUrl = this.#constructRunsElectricUrl(
+      url,
+      environment,
+      whereClause,
+      clientVersion
+    );
 
     return this.#performElectricRequest(electricUrl, environment, undefined, clientVersion);
   }
 
-  #constructRunsElectricUrl(url: URL | string, whereClause: string, clientVersion?: string): URL {
+  #constructRunsElectricUrl(
+    url: URL | string,
+    environment: RealtimeEnvironment,
+    whereClause: string,
+    clientVersion?: string
+  ): URL {
     const $url = new URL(url.toString());
 
-    const electricUrl = new URL(`${this.options.electricOrigin}/v1/shape`);
+    const electricOrigin = this.#resolveElectricOrigin(environment.id);
+    const electricUrl = new URL(`${electricOrigin}/v1/shape`);
 
     // Copy over all the url search params to the electric url
     $url.searchParams.forEach((value, key) => {
@@ -428,6 +440,16 @@ export class RealtimeClient {
     return `${this.options.keyPrefix}:${environmentId}`;
   }
 
+  #resolveElectricOrigin(environmentId: string) {
+    if (typeof this.options.electricOrigin === "string") {
+      return this.options.electricOrigin;
+    }
+
+    const index = jumpHash(environmentId, this.options.electricOrigin.length);
+
+    return this.options.electricOrigin[index] ?? this.options.electricOrigin[0];
+  }
+
   #registerCommands() {
     this.redis.defineCommand("incrementAndCheckConcurrency", {
       numberOfKeys: 1,

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

@@ -4,8 +4,10 @@ import { RealtimeClient } from "./realtimeClient.server";
 import { getCachedLimit } from "./platform.v3.server";
 
 function initializeRealtimeClient() {
+  const electricOrigin = env.ELECTRIC_ORIGIN_SHARDS?.split(",") ?? env.ELECTRIC_ORIGIN;
+
   return new RealtimeClient({
-    electricOrigin: env.ELECTRIC_ORIGIN,
+    electricOrigin: electricOrigin,
     keyPrefix: "tr:realtime:concurrency",
     redis: {
       port: env.RATE_LIMIT_REDIS_PORT,