chore: final reviews

Author: didinele

docs/roadmap/12-horizontal-scaling.md

@@ -164,8 +164,11 @@ its balance, and it exists because without it recovery simply never happened:
 
 So every transient replica loss used to cost balance permanently. The handoff closes that with two rules:
 
-1. **A covering replica sheds when a spare is waiting.** The watcher, on finding no gaps, checks whether it holds
-   more than one index while a spare is advertising — and if so restarts. Shutdown releases every index it holds.
+1. **An elected covering replica sheds when a spare is waiting.** The watcher, on finding no gaps, works out
+   which replicas are covering (from the lease owners) and elects the lowest-primary ones, **capped at the number
+   of spares actually waiting**. Shedding means restarting, and a restart releases _every_ index that replica
+   holds — so letting all covering replicas react to one spare would black out far more than that spare can take
+   back, and every restart past the first would be churn. Debounced over two checks, like the gap branch.
 2. **A replica stands down from greedy claiming while a spare is advertising.** This is the half that makes it
    stick: without it the shedding replica would grab its extra index straight back on the way up, swapping roles
    with the spare instead of rebalancing.
@@ -174,6 +177,22 @@ Suppressing the greedy step is deliberately all the negotiation there is — no
 Whoever is left claims what the other declined, on its next poll. If the spare dies mid-handoff its advertisement
 goes stale, the next boot covers as before, and a gap (if any) falls back to the watcher.
 
+**It converges, but not always in one step.** The rule is "the shedder drops to its primary, the spare takes the
+rest", which is exact when one peer died and returned (`B[1,2]` → `B[1] S[2]`) and coarser the more a single
+replica had absorbed:
+
+| Before                          | One spare returns  | Split                       |
+| ------------------------------- | ------------------ | --------------------------- |
+| `A[0]  B[1,2]`                  | `A[0] B[1] S[2]`   | even                        |
+| `A[0]  B[1,2,3]`                | `A[0] B[1] S[2,3]` | even for three replicas     |
+| `A[0,1,2,3]` (only one running) | `A[0] S[1,2,3]`    | 1/3 — `2/2` would be better |
+
+Each returning replica pulls the largest holder down to one index and takes what it drops, so `1/3` becomes
+`1/1/2` and then `1/1/1/1` as the rest come back. No step is ever worse than the state before it, coverage stays
+complete throughout, and the fully-staffed end state is one index each. Landing evenly in a _single_ step would
+mean a replica knowing how many peers are live and how many indices each holds — the cross-replica bookkeeping
+this design avoids everywhere else.
+
 The remaining rough edge is a **straggler**: a replica starting well after its peers' settle window still idles
 rather than triggering an immediate rebalance — it only gets an index once a covering peer notices it. That is one
 watcher interval, not "until the next deploy".

packages/private/bot-core/src/lib/__tests__/replica.test.ts

@@ -1,7 +1,14 @@
+import { Buffer } from 'node:buffer';
 import { setTimeout } from 'node:timers';
 import { setTimeout as sleepMs } from 'node:timers/promises';
 import { beforeEach, expect, test, vi } from 'vitest';
-import { claimReplicaSlot, computeTotalIndices, getReplicaIndex, shardIdsForIndices } from '../replica.js';
+import {
+	claimReplicaSlot,
+	computeTotalIndices,
+	electGiveBackOwners,
+	getReplicaIndex,
+	shardIdsForIndices,
+} from '../replica.js';
 
 const keys = new Map<string, string>();
 const sortedSets = new Map<string, Map<string, number>>();
@@ -23,6 +30,11 @@ vi.mock('@chatsift/backend-core', () => ({
 			async exists(key: string) {
 				return keys.has(key) ? 1 : 0;
 			},
+			async get(key: string) {
+				// Buffers, like the shared client's type mapping produces.
+				const value = keys.get(key);
+				return value === undefined ? null : Buffer.from(value);
+			},
 			async eval() {
 				return 1;
 			},
@@ -59,6 +71,13 @@ async function boot(shardCount: number, shardsPerReplica: number) {
 	return claimReplicaSlot({ botId: 'AMA', shardCount, shardsPerReplica, settleMs: 0 });
 }
 
+/**
+ * Boots a replica that is expected to find nothing free, so it parks in the hot-spare loop and advertises.
+ */
+async function parkSpare() {
+	return claimReplicaSlot({ botId: 'AMA', shardCount: 16, shardsPerReplica: 4, settleMs: 0, hotSparePollMs: 5 });
+}
+
 /**
  * Mimics a replica parked in the hot-spare loop, which advertises itself on every poll.
  */
@@ -334,3 +353,65 @@ test('a returning replica gets an index back instead of idling forever', async (
 	expect(tookOver.heldIndices).toHaveLength(1);
 	expect([...revived.heldIndices, ...tookOver.heldIndices].sort((left, right) => left - right)).toStrictEqual([2, 3]);
 });
+
+test('a replica can hold three or four indices when peers are missing', async () => {
+	const [a, b] = await Promise.all([boot(16, 4), boot(16, 4)]);
+	expect(b.heldIndices).toStrictEqual([1, 2, 3]);
+	expect(a.heldIndices).toStrictEqual([0]);
+
+	keys.clear();
+	const lone = await boot(16, 4);
+	expect(lone.heldIndices).toStrictEqual([0, 1, 2, 3]);
+});
+
+test('each returning spare pulls the biggest holder down and never loses coverage', async () => {
+	// The handoff is "the shedder drops to its primary, the spare takes the rest" -- exactly right when one
+	// peer died and came back, and progressively coarser the more indices one replica had absorbed. What must
+	// hold at every step is that coverage stays complete and the spare stops being idle.
+	let held = [await boot(16, 4)];
+	expect(held[0]!.heldIndices).toStrictEqual([0, 1, 2, 3]);
+
+	// Replicas trickle back one at a time; each one forces the largest holder to shed.
+	for (const expectedMax of [3, 2, 1]) {
+		const spare = parkSpare();
+		await sleepMs(30);
+
+		const covering = held.find((slot) => slot.heldIndices.length > 1)!;
+		for (const index of covering.heldIndices) {
+			keys.delete(`shardlease:AMA:${index}`);
+		}
+
+		const [rebooted, tookOver] = await Promise.all([boot(16, 4), spare]);
+		held = [...held.filter((slot) => slot !== covering), rebooted, tookOver];
+
+		// Coverage is whole again, and nobody is idle.
+		expect(held.flatMap((slot) => slot.heldIndices).sort((left, right) => left - right)).toStrictEqual([0, 1, 2, 3]);
+		expect(Math.max(...held.map((slot) => slot.heldIndices.length))).toBe(expectedMax);
+	}
+
+	// Fully staffed: one index each.
+	expect(held.map((slot) => slot.heldIndices.length)).toStrictEqual([1, 1, 1, 1]);
+});
+
+test('only as many replicas give back as there are spares to take the work', () => {
+	// Two replicas each covering two indices, one spare. A restart sheds *every* index a replica holds, so if
+	// both reacted the cluster would go fully dark to hand back a single index -- and the second restart would
+	// achieve nothing, since one spare can only absorb one.
+	const owners = ['r0', 'r0', 'r1', 'r1'];
+
+	expect(electGiveBackOwners(owners, 1)).toStrictEqual(['r0']);
+	expect(electGiveBackOwners(owners, 2)).toStrictEqual(['r0', 'r1']);
+	// More spares than covering replicas is fine -- nobody extra is elected.
+	expect(electGiveBackOwners(owners, 5)).toStrictEqual(['r0', 'r1']);
+});
+
+test('nobody gives back when no spare is waiting, or when nobody is covering', () => {
+	expect(electGiveBackOwners(['r0', 'r0', 'r1', 'r1'], 0)).toStrictEqual([]);
+	expect(electGiveBackOwners(['r0', 'r1', 'r2', 'r3'], 2)).toStrictEqual([]);
+});
+
+test('give-back is elected lowest primary first, and ignores unclaimed indices', () => {
+	// Deterministic ordering matters: every replica runs this independently and must reach the same answer.
+	expect(electGiveBackOwners(['r2', 'r2', 'r0', 'r0'], 1)).toStrictEqual(['r2']);
+	expect(electGiveBackOwners([null, 'r1', 'r1', null], 1)).toStrictEqual(['r1']);
+});

packages/private/bot-core/src/lib/replica.ts

@@ -44,12 +44,39 @@ async function withdrawSpare(botId: GuildListKey, token: string): Promise<void>
 }
 
 /**
- * Whether any replica is currently idle and waiting for an index, pruning advertisements that have gone stale.
+ * How many replicas are currently idle and waiting for an index, pruning advertisements that have gone stale.
  */
-async function hasWaitingSpare(botId: GuildListKey): Promise<boolean> {
+async function countWaitingSpares(botId: GuildListKey): Promise<number> {
 	const { redis } = getContext();
 	await redis.zRemRangeByScore(sparesKey(botId), 0, Date.now() - SPARE_STALE_MS);
-	return (await redis.zCard(sparesKey(botId))) > 0;
+	return redis.zCard(sparesKey(botId));
+}
+
+async function hasWaitingSpare(botId: GuildListKey): Promise<boolean> {
+	return (await countWaitingSpares(botId)) > 0;
+}
+
+/**
+ * Which replicas should give an index back this round, lowest primary first, capped at the number of spares
+ * actually waiting to take one.
+ */
+export function electGiveBackOwners(owners: (string | null)[], spareCount: number): string[] {
+	if (spareCount <= 0) {
+		return [];
+	}
+
+	const indicesByOwner = new Map<string, number[]>();
+	for (const [index, owner] of owners.entries()) {
+		if (owner) {
+			indicesByOwner.set(owner, [...(indicesByOwner.get(owner) ?? []), index]);
+		}
+	}
+
+	return [...indicesByOwner.entries()]
+		.filter(([, indices]) => indices.length > 1)
+		.sort(([, left], [, right]) => left[0]! - right[0]!)
+		.slice(0, spareCount)
+		.map(([owner]) => owner);
 }
 
 /**
@@ -203,6 +230,10 @@ export async function claimReplicaSlot({
 		// advertised, not silently: a peer covering two indices watches for this and hands one back.
 		logger.warn({ botId, totalIndices, shardCount, shardsPerReplica }, 'no free replica index, idling as a hot spare');
 
+		// Registered here rather than after a slot is claimed: a spare `SIGTERM`ed mid-idle would otherwise leave
+		// its advertisement to expire, and a covering peer could restart to hand off to a replica already gone.
+		onShutdown('replica-spare', async () => withdrawSpare(botId, token));
+
 		while (primary === null) {
 			await advertiseSpare(botId, token);
 			await sleep(hotSparePollMs);
@@ -257,14 +288,14 @@ export async function claimReplicaSlot({
 		shardIds.length > shardsPerReplica ? 'claimed replica slot, covering for missing replicas' : 'claimed replica slot',
 	);
 
-	startWatching(botId, heldIndices, totalIndices);
+	startWatching(botId, heldIndices, totalIndices, token);
 	onShutdown('replica-lease', async () => releaseAll(botId, heldIndices, token));
 
 	return { index: primary, heldIndices, shardIds };
 }
 
 /**
- * Re-asserts what indeces this replica is responsible for
+ * Re-asserts what indices this replica is responsible for
  */
 function startRenewing(botId: GuildListKey, heldIndices: number[], token: string): void {
 	let lastRenewedAt = Date.now();
@@ -311,18 +342,22 @@ function startRenewing(botId: GuildListKey, heldIndices: number[], token: string
  * `claimReplicaSlot` never claims past a live peer, so nobody else could fill it anyway. A gap at index 0
  * has nothing below it, so the lowest remaining holder takes that one.
  */
-function startWatching(botId: GuildListKey, heldIndices: number[], totalIndices: number): void {
+function startWatching(botId: GuildListKey, heldIndices: number[], totalIndices: number, token: string): void {
 	let consecutiveGaps = 0;
+	let consecutiveGiveBacks = 0;
 
 	setInterval(async () => {
 		const { logger, redis } = getContext();
 
 		try {
-			const holders = await Promise.all(
-				Array.from({ length: totalIndices }, async (_, index) => redis.exists(leaseKey(botId, index))),
-			);
-
-			const gaps = holders.flatMap((held, index) => (held ? [] : [index]));
+			// Owners rather than mere existence: the give-back branch has to see who holds what across the whole
+			// cluster, not just whether an index is taken. Values come back as buffers under the shared client's
+			// type mapping, so normalise before comparing to our own token.
+			const owners = (
+				await Promise.all(Array.from({ length: totalIndices }, async (_, index) => redis.get(leaseKey(botId, index))))
+			).map((owner) => owner?.toString() ?? null);
+
+			const gaps = owners.flatMap((owner, index) => (owner ? [] : [index]));
 			if (gaps.length === 0) {
 				consecutiveGaps = 0;
 
@@ -331,19 +366,32 @@ function startWatching(botId: GuildListKey, heldIndices: number[], totalIndices:
 				// being a gap the moment it absorbs it. Nothing else would ever notice, so the cluster would run
 				// permanently lopsided until the next deploy. Restarting sheds the extras (shutdown releases every
 				// held index) and the greedy step stands down on the way back up, leaving them for the spare.
-				if (heldIndices.length > 1 && (await hasWaitingSpare(botId))) {
-					logger.info(
-						{ botId, heldIndices },
-						'a hot spare is waiting and this replica is covering extra indices, restarting to hand them over',
-					);
-					process.kill(process.pid, 'SIGTERM');
+
+				// Elected rather than "whoever is covering", and debounced like the gap branch below: a restart
+				// sheds every index this replica holds, so several covering replicas all reacting to one spare
+				// would black out far more than that spare can take back.
+				const givingBack = electGiveBackOwners(owners, await countWaitingSpares(botId));
+				if (!givingBack.includes(token)) {
+					consecutiveGiveBacks = 0;
+					return;
+				}
+
+				consecutiveGiveBacks += 1;
+				if (consecutiveGiveBacks < 2) {
+					return;
 				}
 
+				logger.info(
+					{ botId, heldIndices },
+					'a hot spare is waiting and this replica is covering extra indices, restarting to hand them over',
+				);
+				process.kill(process.pid, 'SIGTERM');
 				return;
 			}
 
+			consecutiveGiveBacks = 0;
 			const firstGap = gaps[0]!;
-			const lowestHeld = Math.min(...holders.flatMap((held, index) => (held ? [index] : [])));
+			const lowestHeld = Math.min(...owners.flatMap((owner, index) => (owner ? [index] : [])));
 			const canFillIt = firstGap === 0 ? lowestHeld === heldIndices[0] : heldIndices.includes(firstGap - 1);
 			if (!canFillIt) {
 				return;
@@ -364,7 +412,7 @@ function startWatching(botId: GuildListKey, heldIndices: number[], totalIndices:
 }
 
 /**
- * Release all our indeces
+ * Release all our indices
  */
 async function releaseAll(botId: GuildListKey, heldIndices: number[], token: string): Promise<void> {
 	const { redis } = getContext();

packages/private/bot-core/src/lib/shutdown.ts

@@ -55,8 +55,7 @@ async function shutdown(signal: string): Promise<void> {
 	await sleep(DRAIN_GRACE_MS);
 	await runRegisteredSteps();
 
-	// TODO: We don't kill the gateway to prevent it from settings state to null/none. Should investigate
-	// if discord.js is making the right assumption here, and if not, open a PR.
+	// TODO: discord.js PR
 	try {
 		await Promise.all([redis.quit(), db.end({ timeout: 2 })]);
 	} catch (error) {