redis-distributed-system

main.ts

import Redis from "ioredis";
import { pino } from "pino";
import { z } from "zod";
import { sleep } from "./lib/utils";

type StreamItem<T> = { id: string; data: T };
type StreamEntry = [itemId: string, fields: string[]];
type StreamRead = [streamKey: string, batch: StreamEntry[]];

const envSchema = z.object({
	REDIS_URL: z.string().min(1).default("redis://localhost:6380"),
	REDIS_STREAM_NAME: z.string().min(1).default("stream"),
	REDIS_STREAM_GROUP_NAME: z.string().min(1).default("workers"),
	REDIS_STREAM_MAXLEN: z.coerce.number().int().min(100).default(50_000),
	REDIS_LOCK_TTL_MS: z.coerce.number().int().min(1000).default(10_000),
	REDIS_WORKER_IDLE_MS: z.coerce.number().int().min(100).default(3_000),
	REDIS_WORKER_BATCH_SIZE: z.coerce.number().int().min(1).max(1000).default(32),
	REDIS_WORKER_CLAIM_MS: z.coerce.number().int().min(1000).default(30_000),
	REDIS_DEDUP_WINDOW_MS: z.coerce.number().int().min(1000).default(600_000),
	REDIS_LOG_LEVEL: z.enum(["debug", "info", "warn", "error"]).default("info"),
});

const env = envSchema.parse(process.env);

const client = new Redis(env.REDIS_URL);
const log = pino({ level: env.REDIS_LOG_LEVEL, transport: { target: "pino-pretty" } });

const KEYS = {
	lock: `${env.REDIS_STREAM_NAME}:lock`,
	dedup: `${env.REDIS_STREAM_NAME}:dedup`,
};

// [NOTE] @Agent compare token before mutating — after TTL expiry the lock may belong
// to another worker; blind DEL/PEXPIRE would steal or extend a lock we don't own
const LUA_RELEASE_SCRIPT = `
if redis.call('get', KEYS[1]) == ARGV[1] then
    return redis.call('del', KEYS[1])
end
return 0
`;

const LUA_RENEW_SCRIPT = `
if redis.call('get', KEYS[1]) == ARGV[1] then
    return redis.call('pexpire', KEYS[1], ARGV[2])
end
return 0
`;

const lock = {
	acquire: async (id: string): Promise<null | string> => {
		const token = `${id}:${crypto.randomUUID().slice(0, 8)}`;
		const acquired = await client.set(KEYS.lock, token, "PX", env.REDIS_LOCK_TTL_MS, "NX");
		if (!acquired) return null;
		log.info({ token }, "lock.acquire.completed");
		return token;
	},

	renew: async (token: string): Promise<boolean> => {
		const renewed =
			((await client.eval(
				LUA_RENEW_SCRIPT,
				1,
				KEYS.lock,
				token,
				env.REDIS_LOCK_TTL_MS,
			)) as number) === 1;
		if (!renewed) log.warn({ token }, "lock.renew.failed");
		return renewed;
	},

	release: async (token: string) => {
		const released = ((await client.eval(LUA_RELEASE_SCRIPT, 1, KEYS.lock, token)) as number) === 1;
		if (!released) log.warn({ token }, "lock.release.skipped");
	},
};

const dedup = {
	push: async (ids: string[]): Promise<string[]> => {
		if (ids.length === 0) return [];
		const now = Date.now();
		const pipeline = client.pipeline();
		for (const id of ids) pipeline.zadd(KEYS.dedup, "NX", now, id);
		pipeline.zremrangebyscore(KEYS.dedup, "-inf", now - env.REDIS_DEDUP_WINDOW_MS);
		const results = await pipeline.exec();
		return ids.filter((_, index) => results?.[index][1] === 1);
	},
};

function toObject(fields: string[]): Record<string, string> {
	const record: Record<string, string> = {};
	for (let i = 0; i < fields.length; i += 2) record[fields[i]] = fields[i + 1];
	return record;
}

async function producer<T>(id: string, source: () => Promise<StreamItem<T>[]>) {
	let stopping = false;
	process.on("SIGINT", () => {
		stopping = true;
	});

	while (!stopping) {
		const token = await lock.acquire(id);
		if (!token) {
			await sleep(env.REDIS_WORKER_IDLE_MS);
			continue;
		}
		try {
			while (!stopping) {
				const batch = await source();
				if (batch.length === 0) break;
				const pipeline = client.pipeline();
				const fresh = new Set(await dedup.push(batch.map((item) => item.id)));
				for (const item of batch) {
					if (!fresh.has(item.id)) continue;
					pipeline.xadd(
						env.REDIS_STREAM_NAME,
						"MAXLEN",
						"~",
						String(env.REDIS_STREAM_MAXLEN),
						"*",
						"data",
						JSON.stringify(item.data),
					);
				}
				await pipeline.exec();
				log.info({ id, fetched: batch.length, fresh: fresh.size }, "producer.batch.completed");
				if (!(await lock.renew(token))) break;
			}
		} finally {
			await lock.release(token);
		}
	}

	log.info({ id }, "producer.shutdown.completed");
	process.exit(0);
}

async function consumer<T>(id: string, work: (item: T) => Promise<void>) {
	const worker = log.child({ id });

	// [NOTE] @Agent SIGINT only flips the flag — the in-flight batch finishes and acks
	// before we unregister, so no receipt is left pending and nothing needs XAUTOCLAIM
	let stopping = false;
	process.on("SIGINT", () => {
		stopping = true;
		worker.info("consumer.shutdown.started");
	});

	async function consume(batch: StreamEntry[]) {
		if (batch.length === 0) return;
		for (const [itemId, fields] of batch) {
			const item = JSON.parse(toObject(fields).data) as T;
			worker.debug({ itemId }, "consumer.item.processing");
			await work(item);
		}
		await client.xack(
			env.REDIS_STREAM_NAME,
			env.REDIS_STREAM_GROUP_NAME,
			...batch.map(([itemId]) => itemId),
		);
	}

	let lastClaim = 0;
	while (!stopping) {
		if (Date.now() - lastClaim > env.REDIS_WORKER_CLAIM_MS) {
			const [, claimed, deleted] = (await client.xautoclaim(
				env.REDIS_STREAM_NAME,
				env.REDIS_STREAM_GROUP_NAME,
				id,
				env.REDIS_WORKER_CLAIM_MS,
				"0",
				"COUNT",
				env.REDIS_WORKER_BATCH_SIZE,
			)) as [string, StreamEntry[], string[]];
			if (claimed.length) worker.warn({ count: claimed.length }, "consumer.claim.rescued");
			if (deleted.length) worker.warn({ count: deleted.length }, "consumer.claim.trimmed");
			await consume(claimed);
			lastClaim = Date.now();
		}

		const results = (await client.xreadgroup(
			"GROUP",
			env.REDIS_STREAM_GROUP_NAME,
			id,
			"COUNT",
			env.REDIS_WORKER_BATCH_SIZE,
			"BLOCK",
			2000,
			"STREAMS",
			env.REDIS_STREAM_NAME,
			">",
		)) as StreamRead[] | null;
		if (!results) continue;
		await consume(results[0][1]);
	}

	await client.xgroup("DELCONSUMER", env.REDIS_STREAM_NAME, env.REDIS_STREAM_GROUP_NAME, id);
	worker.info("consumer.shutdown.completed");
	process.exit(0);
}

async function stats() {
	const groups = (await client.xinfo("GROUPS", env.REDIS_STREAM_NAME)) as string[][];
	const group = groups
		.map((fields) => toObject(fields))
		.find((entry) => entry.name === env.REDIS_STREAM_GROUP_NAME);
	const perMinute = (await client.xrange(env.REDIS_STREAM_NAME, String(Date.now() - 60_000), "+"))
		.length;
	return {
		lag: Number(group?.lag ?? 0),
		pending: Number(group?.pending ?? 0),
		consumers: Number(group?.consumers ?? 0),
		perMinute,
	};
}

async function setup() {
	try {
		await client.xgroup(
			"CREATE",
			env.REDIS_STREAM_NAME,
			env.REDIS_STREAM_GROUP_NAME,
			"0",
			"MKSTREAM",
		);
	} catch (error) {
		// [NOTE] @Agent BUSYGROUP = group already exists — expected on every boot but the first
		if (!(error instanceof Error && error.message.includes("BUSYGROUP"))) throw error;
	}
}

export { consumer, lock, producer, setup, stats };
export type { StreamItem };