chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,614 @@
|
||||
import { json } from "@remix-run/server-runtime";
|
||||
import { tryCatch } from "@trigger.dev/core/utils";
|
||||
import { safeParseNaturalLanguageDurationAgo } from "@trigger.dev/core/v3/isomorphic";
|
||||
import type { Callback, Result } from "ioredis";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import type { RedisClient, RedisWithClusterOptions } from "~/redis.server";
|
||||
import { createRedisClient } from "~/redis.server";
|
||||
import { longPollingFetch } from "~/utils/longPollingFetch";
|
||||
import { logger } from "./logger.server";
|
||||
import { jumpHash } from "@trigger.dev/core/v3/serverOnly";
|
||||
import type { Cache } from "@unkey/cache";
|
||||
import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
|
||||
import { createLRUMemoryStore } from "@internal/cache";
|
||||
import { RedisCacheStore } from "./unkey/redisCacheStore.server";
|
||||
import { env } from "~/env.server";
|
||||
import type { API_VERSIONS } from "~/api/versions";
|
||||
import { CURRENT_API_VERSION } from "~/api/versions";
|
||||
|
||||
export interface CachedLimitProvider {
|
||||
getCachedLimit: (organizationId: string, defaultValue: number) => Promise<number | undefined>;
|
||||
}
|
||||
|
||||
const DEFAULT_ELECTRIC_COLUMNS = [
|
||||
"id",
|
||||
"taskIdentifier",
|
||||
"createdAt",
|
||||
"updatedAt",
|
||||
"startedAt",
|
||||
"delayUntil",
|
||||
"queuedAt",
|
||||
"expiredAt",
|
||||
"completedAt",
|
||||
"friendlyId",
|
||||
"number",
|
||||
"isTest",
|
||||
"status",
|
||||
"usageDurationMs",
|
||||
"costInCents",
|
||||
"baseCostInCents",
|
||||
"ttl",
|
||||
"payload",
|
||||
"payloadType",
|
||||
"metadata",
|
||||
"metadataType",
|
||||
"output",
|
||||
"outputType",
|
||||
"runTags",
|
||||
"error",
|
||||
"realtimeStreams",
|
||||
];
|
||||
|
||||
const RESERVED_COLUMNS = ["id", "taskIdentifier", "friendlyId", "status", "createdAt"];
|
||||
const RESERVED_SEARCH_PARAMS = ["createdAt", "tags", "skipColumns"];
|
||||
|
||||
export type RealtimeClientOptions = {
|
||||
electricOrigin: string | string[];
|
||||
redis: RedisWithClusterOptions;
|
||||
cachedLimitProvider: CachedLimitProvider;
|
||||
keyPrefix: string;
|
||||
expiryTimeInSeconds?: number;
|
||||
};
|
||||
|
||||
export type RealtimeEnvironment = {
|
||||
id: string;
|
||||
organizationId: string;
|
||||
};
|
||||
|
||||
export type RealtimeRunsParams = {
|
||||
tags?: string[];
|
||||
createdAt?: string;
|
||||
};
|
||||
|
||||
export type RealtimeRequestOptions = {
|
||||
skipColumns?: string[];
|
||||
};
|
||||
|
||||
export class RealtimeClient {
|
||||
private redis: RedisClient;
|
||||
private expiryTimeInSeconds: number;
|
||||
private cachedLimitProvider: CachedLimitProvider;
|
||||
private cache: Cache<{ createdAtFilter: string }>;
|
||||
|
||||
constructor(private options: RealtimeClientOptions) {
|
||||
this.redis = createRedisClient("trigger:realtime", options.redis);
|
||||
this.expiryTimeInSeconds = options.expiryTimeInSeconds ?? 60 * 5; // default to 5 minutes
|
||||
this.cachedLimitProvider = options.cachedLimitProvider;
|
||||
this.#registerCommands();
|
||||
|
||||
const ctx = new DefaultStatefulContext();
|
||||
const memory = createLRUMemoryStore(1000);
|
||||
const redisCacheStore = new RedisCacheStore({
|
||||
connection: {
|
||||
keyPrefix: "tr:cache:realtime",
|
||||
port: options.redis.port,
|
||||
host: options.redis.host,
|
||||
username: options.redis.username,
|
||||
password: options.redis.password,
|
||||
tlsDisabled: options.redis.tlsDisabled,
|
||||
clusterMode: options.redis.clusterMode,
|
||||
},
|
||||
});
|
||||
|
||||
// This cache holds the limits fetched from the platform service
|
||||
const cache = createCache({
|
||||
createdAtFilter: new Namespace<string>(ctx, {
|
||||
stores: [memory, redisCacheStore],
|
||||
fresh: 60_000 * 60 * 24 * 7, // 1 week
|
||||
stale: 60_000 * 60 * 24 * 14, // 2 weeks
|
||||
}),
|
||||
});
|
||||
|
||||
this.cache = cache;
|
||||
}
|
||||
|
||||
async streamRun(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
runId: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
return this.#streamRunsWhere(
|
||||
url,
|
||||
environment,
|
||||
`id='${runId}'`,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
}
|
||||
|
||||
async streamBatch(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
batchId: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
const whereClauses: string[] = [
|
||||
`"runtimeEnvironmentId"='${environment.id}'`,
|
||||
`"batchId"='${batchId}'`,
|
||||
];
|
||||
|
||||
const whereClause = whereClauses.join(" AND ");
|
||||
|
||||
return this.#streamRunsWhere(
|
||||
url,
|
||||
environment,
|
||||
whereClause,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
}
|
||||
|
||||
async streamRuns(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
params: RealtimeRunsParams,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
const whereClauses: string[] = [`"runtimeEnvironmentId"='${environment.id}'`];
|
||||
|
||||
if (params.tags) {
|
||||
whereClauses.push(`"runTags" @> ARRAY[${params.tags.map((t) => `'${t}'`).join(",")}]`);
|
||||
}
|
||||
|
||||
const createdAtFilter = await this.#calculateCreatedAtFilter(url, params.createdAt);
|
||||
|
||||
if (createdAtFilter) {
|
||||
whereClauses.push(`"createdAt" > '${createdAtFilter.toISOString()}'`);
|
||||
}
|
||||
|
||||
const whereClause = whereClauses.join(" AND ");
|
||||
|
||||
const response = await this.#streamRunsWhere(
|
||||
url,
|
||||
environment,
|
||||
whereClause,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
|
||||
if (createdAtFilter) {
|
||||
const [setCreatedAtFilterError] = await tryCatch(
|
||||
this.#setCreatedAtFilterFromResponse(response, createdAtFilter)
|
||||
);
|
||||
|
||||
if (setCreatedAtFilterError) {
|
||||
logger.error("[realtimeClient] Failed to set createdAt filter", {
|
||||
error: setCreatedAtFilterError,
|
||||
createdAtFilter,
|
||||
responseHeaders: Object.fromEntries(response.headers.entries()),
|
||||
responseStatus: response.status,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
async #calculateCreatedAtFilter(url: URL | string, createdAt?: string) {
|
||||
const duration = createdAt ?? "24h";
|
||||
const $url = new URL(url.toString());
|
||||
const shapeId = extractShapeId($url);
|
||||
|
||||
if (!shapeId) {
|
||||
// This means we need to calculate the createdAt filter and store it in redis after we get back the response
|
||||
const createdAtFilter = safeParseNaturalLanguageDurationAgo(duration);
|
||||
|
||||
// Validate that the createdAt filter is in the past, and not more than the maximum age in the past.
|
||||
// if it's more than the maximum age in the past, just return the maximum age in the past Date
|
||||
if (
|
||||
createdAtFilter &&
|
||||
createdAtFilter < new Date(Date.now() - env.REALTIME_MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS)
|
||||
) {
|
||||
return new Date(Date.now() - env.REALTIME_MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS);
|
||||
}
|
||||
|
||||
return createdAtFilter;
|
||||
} else {
|
||||
// We need to get the createdAt filter value from redis, if there is none we need to return undefined
|
||||
const [createdAtFilterError, createdAtFilter] = await tryCatch(
|
||||
this.#getCreatedAtFilter(shapeId)
|
||||
);
|
||||
|
||||
if (createdAtFilterError) {
|
||||
logger.error("[realtimeClient] Failed to get createdAt filter", {
|
||||
shapeId,
|
||||
error: createdAtFilterError,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
return createdAtFilter;
|
||||
}
|
||||
}
|
||||
|
||||
async #getCreatedAtFilter(shapeId: string) {
|
||||
const createdAtFilterCacheResult = await this.cache.createdAtFilter.get(shapeId);
|
||||
|
||||
if (createdAtFilterCacheResult.err) {
|
||||
logger.error("[realtimeClient] Failed to get createdAt filter", {
|
||||
shapeId,
|
||||
error: createdAtFilterCacheResult.err,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (!createdAtFilterCacheResult.val) {
|
||||
return;
|
||||
}
|
||||
|
||||
return new Date(createdAtFilterCacheResult.val);
|
||||
}
|
||||
|
||||
async #setCreatedAtFilterFromResponse(response: Response, createdAtFilter: Date) {
|
||||
const shapeId = extractShapeIdFromResponse(response);
|
||||
|
||||
if (!shapeId) {
|
||||
return;
|
||||
}
|
||||
|
||||
await this.cache.createdAtFilter.set(shapeId, createdAtFilter.toISOString());
|
||||
}
|
||||
|
||||
async #streamRunsWhere(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
whereClause: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
const electricUrl = this.#constructRunsElectricUrl(
|
||||
url,
|
||||
environment,
|
||||
whereClause,
|
||||
requestOptions,
|
||||
clientVersion
|
||||
);
|
||||
|
||||
return this.#performElectricRequest(
|
||||
electricUrl,
|
||||
environment,
|
||||
apiVersion,
|
||||
signal,
|
||||
clientVersion
|
||||
);
|
||||
}
|
||||
|
||||
#constructRunsElectricUrl(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
whereClause: string,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string
|
||||
): URL {
|
||||
const $url = new URL(url.toString());
|
||||
|
||||
const electricOrigin = this.#resolveElectricOrigin($url, whereClause, 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) => {
|
||||
if (RESERVED_SEARCH_PARAMS.includes(key)) {
|
||||
return;
|
||||
}
|
||||
|
||||
electricUrl.searchParams.set(key, value);
|
||||
});
|
||||
|
||||
electricUrl.searchParams.set("where", whereClause);
|
||||
electricUrl.searchParams.set("table", 'public."TaskRun"');
|
||||
|
||||
if (!clientVersion) {
|
||||
// If the client version is not provided, that means we're using an older client
|
||||
// This means the client will be sending shape_id instead of handle
|
||||
electricUrl.searchParams.set("handle", electricUrl.searchParams.get("shape_id") ?? "");
|
||||
}
|
||||
|
||||
let skipColumns = getSkipColumns($url.searchParams, requestOptions);
|
||||
|
||||
if (skipColumns.length > 0) {
|
||||
skipColumns = skipColumns.filter((c) => c !== "" && !RESERVED_COLUMNS.includes(c));
|
||||
|
||||
electricUrl.searchParams.set(
|
||||
"columns",
|
||||
DEFAULT_ELECTRIC_COLUMNS.filter((c) => !skipColumns.includes(c))
|
||||
.map((c) => `"${c}"`)
|
||||
.join(",")
|
||||
);
|
||||
} else {
|
||||
electricUrl.searchParams.set(
|
||||
"columns",
|
||||
DEFAULT_ELECTRIC_COLUMNS.map((c) => `"${c}"`).join(",")
|
||||
);
|
||||
}
|
||||
|
||||
return electricUrl;
|
||||
}
|
||||
|
||||
async #performElectricRequest(
|
||||
url: URL,
|
||||
environment: RealtimeEnvironment,
|
||||
apiVersion: API_VERSIONS,
|
||||
signal?: AbortSignal,
|
||||
clientVersion?: string
|
||||
) {
|
||||
const shapeId = extractShapeId(url);
|
||||
|
||||
logger.debug("[realtimeClient] request", {
|
||||
url: url.toString(),
|
||||
});
|
||||
|
||||
const rewriteResponseHeaders: Record<string, string> = clientVersion
|
||||
? {}
|
||||
: { "electric-handle": "electric-shape-id", "electric-offset": "electric-chunk-last-offset" };
|
||||
|
||||
if (!shapeId) {
|
||||
// If the shapeId is not present, we're just getting the initial value
|
||||
return this.#doLongPollingFetch(url, apiVersion, signal, rewriteResponseHeaders);
|
||||
}
|
||||
|
||||
const isLive = isLiveRequestUrl(url);
|
||||
|
||||
if (!isLive) {
|
||||
return this.#doLongPollingFetch(url, apiVersion, signal, rewriteResponseHeaders);
|
||||
}
|
||||
|
||||
const requestId = randomUUID();
|
||||
|
||||
// We now need to wrap the longPollingFetch in a concurrency tracker
|
||||
const concurrencyLimit = await this.cachedLimitProvider.getCachedLimit(
|
||||
environment.organizationId,
|
||||
100_000
|
||||
);
|
||||
|
||||
if (!concurrencyLimit) {
|
||||
logger.error("Failed to get concurrency limit", {
|
||||
organizationId: environment.organizationId,
|
||||
});
|
||||
|
||||
return json({ error: "Failed to get concurrency limit" }, { status: 500 });
|
||||
}
|
||||
|
||||
logger.debug("[realtimeClient] increment and check", {
|
||||
concurrencyLimit,
|
||||
shapeId,
|
||||
requestId,
|
||||
environment: {
|
||||
id: environment.id,
|
||||
organizationId: environment.organizationId,
|
||||
},
|
||||
});
|
||||
|
||||
const canProceed = await this.#incrementAndCheck(environment.id, requestId, concurrencyLimit);
|
||||
|
||||
if (!canProceed) {
|
||||
logger.debug("[realtimeClient] too many concurrent requests", {
|
||||
requestId,
|
||||
environmentId: environment.id,
|
||||
});
|
||||
|
||||
return json({ error: "Too many concurrent requests" }, { status: 429 });
|
||||
}
|
||||
|
||||
try {
|
||||
// ... (rest of your existing code for the long polling request)
|
||||
const response = await this.#doLongPollingFetch(
|
||||
url,
|
||||
apiVersion,
|
||||
signal,
|
||||
rewriteResponseHeaders
|
||||
);
|
||||
|
||||
// If this is the initial request, the response.headers['electric-handle'] will be the shapeId
|
||||
// And we may need to set the "createdAt" filter timestamp keyed by the shapeId
|
||||
// Then in the next request, we will get the createdAt timestamp value via the shapeId and use it to filter the results
|
||||
|
||||
// Decrement the counter after the long polling request is complete
|
||||
await this.#decrementConcurrency(environment.id, requestId);
|
||||
|
||||
return response;
|
||||
} catch (error) {
|
||||
// Decrement the counter if the request fails
|
||||
await this.#decrementConcurrency(environment.id, requestId);
|
||||
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async #doLongPollingFetch(
|
||||
url: URL,
|
||||
apiVersion: API_VERSIONS,
|
||||
signal?: AbortSignal,
|
||||
rewriteResponseHeaders?: Record<string, string>
|
||||
) {
|
||||
if (apiVersion === CURRENT_API_VERSION) {
|
||||
return longPollingFetch(url.toString(), { signal }, rewriteResponseHeaders);
|
||||
}
|
||||
|
||||
const response = await longPollingFetch(url.toString(), { signal }, rewriteResponseHeaders);
|
||||
|
||||
return this.#rewriteResponseForNoneApiVersion(response);
|
||||
}
|
||||
|
||||
async #rewriteResponseForNoneApiVersion(response: Response) {
|
||||
// Get the raw response body
|
||||
const responseBody = await response.text();
|
||||
|
||||
// Rewrite the response body
|
||||
const rewrittenResponseBody = this.#rewriteResponseBodyForNoneApiVersion(responseBody);
|
||||
|
||||
// Return the rewritten response
|
||||
return new Response(rewrittenResponseBody, {
|
||||
status: response.status,
|
||||
headers: response.headers,
|
||||
});
|
||||
}
|
||||
|
||||
// Rewrites "status":"DEQUEUED" to "status":"EXECUTING"
|
||||
#rewriteResponseBodyForNoneApiVersion(responseBody: string) {
|
||||
return responseBody.replace(/"status":"DEQUEUED"/g, '"status":"EXECUTING"');
|
||||
}
|
||||
|
||||
async #incrementAndCheck(environmentId: string, requestId: string, limit: number) {
|
||||
const key = this.#getKey(environmentId);
|
||||
const now = Date.now();
|
||||
|
||||
const result = await this.redis.incrementAndCheckConcurrency(
|
||||
key,
|
||||
now.toString(),
|
||||
requestId,
|
||||
this.expiryTimeInSeconds.toString(), // expiry time
|
||||
(now - this.expiryTimeInSeconds * 1000).toString(), // cutoff time
|
||||
limit.toString()
|
||||
);
|
||||
|
||||
return result === 1;
|
||||
}
|
||||
|
||||
async #decrementConcurrency(environmentId: string, requestId: string) {
|
||||
logger.debug("[realtimeClient] decrement", {
|
||||
requestId,
|
||||
environmentId,
|
||||
});
|
||||
|
||||
const key = this.#getKey(environmentId);
|
||||
|
||||
await this.redis.zrem(key, requestId);
|
||||
}
|
||||
|
||||
#getKey(environmentId: string): string {
|
||||
return `${this.options.keyPrefix}:${environmentId}`;
|
||||
}
|
||||
|
||||
#resolveElectricOrigin(url: URL, whereClause: string, environmentId: string) {
|
||||
if (typeof this.options.electricOrigin === "string") {
|
||||
return this.options.electricOrigin;
|
||||
}
|
||||
|
||||
const shardKey = this.#getShardKey(whereClause, environmentId);
|
||||
|
||||
const index = jumpHash(shardKey, this.options.electricOrigin.length);
|
||||
|
||||
const origin = this.options.electricOrigin[index] ?? this.options.electricOrigin[0];
|
||||
|
||||
logger.debug("[realtimeClient] resolveElectricOrigin", {
|
||||
whereClause,
|
||||
environmentId,
|
||||
shardKey,
|
||||
index,
|
||||
electricOrigin: origin,
|
||||
});
|
||||
|
||||
return origin;
|
||||
}
|
||||
|
||||
#getShardKey(whereClause: string, environmentId: string) {
|
||||
return [environmentId, whereClause].join(":");
|
||||
}
|
||||
|
||||
#registerCommands() {
|
||||
this.redis.defineCommand("incrementAndCheckConcurrency", {
|
||||
numberOfKeys: 1,
|
||||
lua: /* lua */ `
|
||||
local concurrencyKey = KEYS[1]
|
||||
|
||||
local timestamp = tonumber(ARGV[1])
|
||||
local requestId = ARGV[2]
|
||||
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)
|
||||
|
||||
-- Set the expiry time on the key
|
||||
redis.call('EXPIRE', concurrencyKey, expiryTime)
|
||||
|
||||
-- Get the total number of concurrent requests
|
||||
local totalRequests = redis.call('ZCARD', concurrencyKey)
|
||||
|
||||
-- Check if the limit has been exceeded
|
||||
if totalRequests > limit then
|
||||
-- Remove the request we just added
|
||||
redis.call('ZREM', concurrencyKey, requestId)
|
||||
return 0
|
||||
end
|
||||
|
||||
-- Return 1 to indicate success
|
||||
return 1
|
||||
`,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function extractShapeId(url: URL) {
|
||||
return url.searchParams.get("handle") ?? url.searchParams.get("shape_id");
|
||||
}
|
||||
|
||||
function extractShapeIdFromResponse(response: Response) {
|
||||
return response.headers.get("electric-handle");
|
||||
}
|
||||
|
||||
function isLiveRequestUrl(url: URL) {
|
||||
return url.searchParams.has("live") && url.searchParams.get("live") === "true";
|
||||
}
|
||||
|
||||
declare module "ioredis" {
|
||||
interface RedisCommander<Context> {
|
||||
incrementAndCheckConcurrency(
|
||||
key: string,
|
||||
timestamp: string,
|
||||
requestId: string,
|
||||
expiryTime: string,
|
||||
cutoffTime: string,
|
||||
limit: string,
|
||||
callback?: Callback<number>
|
||||
): Result<number, Context>;
|
||||
}
|
||||
}
|
||||
|
||||
function getSkipColumns(searchParams: URLSearchParams, requestOptions?: RealtimeRequestOptions) {
|
||||
if (requestOptions?.skipColumns) {
|
||||
return requestOptions.skipColumns;
|
||||
}
|
||||
|
||||
const skipColumnsRaw = searchParams.get("skipColumns");
|
||||
|
||||
if (skipColumnsRaw) {
|
||||
return skipColumnsRaw.split(",").map((c) => c.trim());
|
||||
}
|
||||
|
||||
return [];
|
||||
}
|
||||
Reference in New Issue
Block a user