feat: give back work to new replicas

Author: didinele

docs/roadmap/12-horizontal-scaling.md

@@ -45,13 +45,21 @@ untrusted Discord input.
 3. Greedily claim any index still free immediately above.
 4. Own the union of the fixed slices belonging to every index held.
 
+**Indices and shard ids are different number spaces**, and they collide confusingly at the low end — index `0`
+entitles its holder to shards `0-3`. Only _indices_ are ever leased in redis; shard ids are derived from them
+locally, by arithmetic. For a 14-shard bot there are four redis keys, not fourteen.
+
 ```
 shardCount=14, SHARDS_PER_REPLICA=4  ->  totalIndices=4
 
-index slices:  idx0 [0-3]  idx1 [4-7]  idx2 [8-11]  idx3 [12,13]
+redis key           what holding it means
+shardlease:<bot>:0  -> run shards 0,1,2,3
+shardlease:<bot>:1  -> run shards 4,5,6,7
+shardlease:<bot>:2  -> run shards 8,9,10,11
+shardlease:<bot>:3  -> run shards 12,13
 
-4 replicas:  one index each                    -> 4 / 4 / 4 / 2 shards
-3 replicas:  the third holds idx2 AND idx3     -> 4 / 4 / 6 shards
+4 replicas:  one index each                          -> 4 / 4 / 4 / 2 shards
+3 replicas:  the third holds indices 2 AND 3         -> 4 / 4 / 6 shards
 5 replicas:  the surplus finds nothing free and waits as a hot spare
 ```
 
@@ -100,7 +108,8 @@ Genuine imbalance has exactly two causes, and both are the cluster not being at
   replicas must cover four indices' worth of shards, one of them holds more. The alternative is leaving those
   shards uncovered for as long as the replica stays away, which is the trade this design refuses.
 - **A straggler.** A replica starting after its peers' settle window finds everything claimed and idles as a hot
-  spare while some peer holds two indices. See below.
+  spare while some peer holds two indices — until that peer notices it and hands one back, see
+  [Handing an index back](#handing-an-index-back).
 
 Perfect balance under a changing replica count would mean re-slicing `shardCount` across however many replicas are
 currently live — which makes every replica's assignment depend on every other's liveness, so one replica
@@ -114,11 +123,11 @@ than one replica temporarily carrying an extra index.
 replica holds the union of every index it claims, and it claims extras precisely when peers are missing — which is
 also when the surviving replicas are carrying the most load:
 
-| Cluster state                  | Shards on the heaviest replica |
-| ------------------------------ | ------------------------------ |
-| Fully provisioned              | `SHARDS_PER_REPLICA`           |
-| One peer missing               | `2 × SHARDS_PER_REPLICA`       |
-| Worst case (only one survivor) | the entire `shardCount`        |
+| Cluster state                  | Shards on the heaviest replica                                          |
+| ------------------------------ | ----------------------------------------------------------------------- |
+| Fully provisioned              | `SHARDS_PER_REPLICA`                                                    |
+| One peer missing               | `2 × SHARDS_PER_REPLICA`, until the peer returns and is handed one back |
+| Worst case (only one survivor) | the entire `shardCount`                                                 |
 
 Size for at least **twice** `SHARDS_PER_REPLICA` if a single replica loss should be absorbed without degrading,
 and treat `shardsOwned` in the boot log (and the `covering for missing replicas` message) as the signal that a
@@ -141,11 +150,33 @@ is the one case with nothing below it, so the lowest holder takes that one.
 
 This is only affordable because restarts RESUME. See below.
 
-### Known wart: stragglers
+### Handing an index back
+
+A replica that finds every index claimed idles as a **hot spare**, advertising itself in a redis sorted set
+(`shardspares:<botId>`, scored by when it last checked in). That advertisement is what lets the cluster recover
+its balance, and it exists because without it recovery simply never happened:
+
+> Four replicas, `C` dies. `B` is elected, restarts, and comes back holding indices 1 _and_ 2 — correct, and
+> carrying double load. `C`'s container then recovers, finds all four indices claimed, and idles forever. There is
+> no gap any more, so the watcher never fires. `B` runs 2× load next to an idle container until the next
+> `./compose up`. Restarting `B` doesn't help either: `B` and the spare just race for the freed indices and swap
+> roles.
+
+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.
+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.
+
+Suppressing the greedy step is deliberately all the negotiation there is — no replica tells another what to take.
+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.
 
-A replica starting well after its peers' settle window finds everything claimed and idles as a hot spare, leaving
-the cluster correct but unbalanced until the next restart. It is logged as a warning. Fixing it live needs
-cross-replica negotiation that is not worth it when `./compose` starts replicas together by construction.
+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".
 
 ## What actually had to change
 
@@ -217,7 +248,9 @@ exit with its sockets open instead leaves Discord holding a resumable session. T
 - `claimed replica slot, covering for missing replicas` means the cluster is short — one replica is carrying more
   than its target. Coverage is fine; capacity is not.
 - `no free replica index, idling as a hot spare` means more replicas are running than the shard count needs, or a
-  straggler missed its settle window.
+  straggler missed its settle window. Paired with `a hot spare is waiting...restarting to hand them over` on a
+  covering peer, that is the rebalance working; on its own for more than a watcher interval, nobody was covering.
+- `hot spare took over a freed replica index` closes that loop -- the spare is no longer idle.
 - `lost replica lease, restarting to re-derive shard assignment` means a renewal found somebody else holding the
   index. Rare and self-healing, but a repeated one means redis latency is eating the lease TTL.
 - `replica indices still unclaimed, restarting to take them over` means a peer died or was scaled away.

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

@@ -1,8 +1,10 @@
 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';
 
 const keys = new Map<string, string>();
+const sortedSets = new Map<string, Map<string, number>>();
 const info = vi.fn();
 const warn = vi.fn();
 
@@ -24,6 +26,25 @@ vi.mock('@chatsift/backend-core', () => ({
 			async eval() {
 				return 1;
 			},
+			async zAdd(key: string, { score, value }: { score: number; value: string }) {
+				const set = sortedSets.get(key) ?? new Map<string, number>();
+				set.set(value, score);
+				sortedSets.set(key, set);
+			},
+			async zRem(key: string, member: string) {
+				sortedSets.get(key)?.delete(member);
+			},
+			async zRemRangeByScore(key: string, min: number, max: number) {
+				const set = sortedSets.get(key);
+				for (const [member, score] of set ?? []) {
+					if (score >= min && score <= max) {
+						set!.delete(member);
+					}
+				}
+			},
+			async zCard(key: string) {
+				return sortedSets.get(key)?.size ?? 0;
+			},
 		},
 	}),
 }));
@@ -38,8 +59,16 @@ async function boot(shardCount: number, shardsPerReplica: number) {
 	return claimReplicaSlot({ botId: 'AMA', shardCount, shardsPerReplica, settleMs: 0 });
 }
 
+/**
+ * Mimics a replica parked in the hot-spare loop, which advertises itself on every poll.
+ */
+function advertiseSpare(atMs = Date.now()) {
+	sortedSets.set('shardspares:AMA', new Map([['a-waiting-spare', atMs]]));
+}
+
 beforeEach(() => {
 	keys.clear();
+	sortedSets.clear();
 	info.mockReset();
 	warn.mockReset();
 });
@@ -176,6 +205,27 @@ test('indices are claimed lowest-first so index 0 always has an owner', async ()
 	expect(slots.flatMap((slot) => slot.shardIds)).toContain(0);
 });
 
+test('held indices are always a contiguous run, never hopping over a live peer', async () => {
+	// `startWatching` picks the holder of `firstGap - 1` to close a gap, which is only a valid choice because a
+	// replica's run stops dead at the first index a peer holds. If the greedy loop ever `continue`d instead of
+	// breaking, a replica could hold [0, 3] and the watcher would start electing replicas that cannot actually
+	// reach the gap -- with every other test still green. This is the guard for that.
+	// Two replicas, five indices: indices 2-4 are free, but they sit above a live peer from A's point of view.
+	const slots = await Promise.all([boot(20, 4), boot(20, 4)]);
+	const [a, b] = slots.sort((left, right) => left.index - right.index);
+
+	// A stops dead at live B rather than hopping it to collect the free indices above.
+	expect(a!.heldIndices).toStrictEqual([0]);
+	// B is the one that reaches them, and takes them as an unbroken run.
+	expect(b!.heldIndices).toStrictEqual([1, 2, 3, 4]);
+
+	// Stated as the general invariant: whatever a replica holds is always start..start+n with no holes.
+	for (const slot of slots) {
+		const expected = Array.from({ length: slot.heldIndices.length }, (_, offset) => slot.heldIndices[0]! + offset);
+		expect(slot.heldIndices).toStrictEqual(expected);
+	}
+});
+
 test('a dead middle replica is taken over by the peer directly below it, not the lowest', async () => {
 	// Re-derivation claims the lowest free index then extends upward, stopping at the first index a live peer
 	// holds -- so a replica can never jump over a living peer. With A[0] B[1] C[2] D[3] and C gone, having A
@@ -236,3 +286,51 @@ test('a straggler that finds nothing free waits as a hot spare, then takes over
 	await expect(straggler).resolves.toMatchObject({ index: 2, shardIds: [8, 9, 10, 11] });
 	expect(warn.mock.calls[0]?.[1]).toBe('no free replica index, idling as a hot spare');
 });
+
+test('a covering replica stands down from greedy claiming while a spare is waiting', async () => {
+	advertiseSpare();
+
+	const slot = await boot(16, 4);
+
+	// Without the stand-down it would sweep up all four indices, leaving the spare nothing to take -- which is
+	// what made a recovered failure stay permanently lopsided.
+	expect(slot.heldIndices).toStrictEqual([0]);
+});
+
+test('a spare that stopped advertising does not keep a replica from covering', async () => {
+	// A spare that died mid-wait must not block its peers from covering the indices it will never claim.
+	advertiseSpare(Date.now() - 120_000);
+
+	const slot = await boot(16, 4);
+
+	expect(slot.heldIndices).toStrictEqual([0, 1, 2, 3]);
+});
+
+test('a returning replica gets an index back instead of idling forever', async () => {
+	// Three replicas for four indices: one of them is covering.
+	const slots = await Promise.all([boot(16, 4), boot(16, 4), boot(16, 4)]);
+	const covering = slots.find((slot) => slot.heldIndices.length > 1)!;
+	expect(covering.heldIndices).toStrictEqual([2, 3]);
+
+	// The fourth replica comes back. Nothing is free, so it advertises and waits.
+	const returning = claimReplicaSlot({
+		botId: 'AMA',
+		shardCount: 16,
+		shardsPerReplica: 4,
+		settleMs: 0,
+		hotSparePollMs: 10,
+	});
+	await sleepMs(40);
+
+	// The covering replica sheds: shutdown releases every index it holds, then it re-derives.
+	for (const index of covering.heldIndices) {
+		keys.delete(`shardlease:AMA:${index}`);
+	}
+
+	const [revived, tookOver] = await Promise.all([boot(16, 4), returning]);
+
+	// Balanced: one index each, and between them they still cover everything the covering replica had.
+	expect(revived.heldIndices).toHaveLength(1);
+	expect(tookOver.heldIndices).toHaveLength(1);
+	expect([...revived.heldIndices, ...tookOver.heldIndices].sort((left, right) => left - right)).toStrictEqual([2, 3]);
+});

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

@@ -10,11 +10,6 @@ import { createSessionStore } from './sessions.js';
 import { onShutdown } from './shutdown.js';
 
 export interface CreateBotGatewayOptions {
-	/**
-	 * Namespaces this bot's gateway sessions, replica leases and identify buckets in the shared redis. Same
-	 * widened key as `createBotClient`'s, since a custom ModMail instance (#216) is its own Discord application
-	 * with its own sessions and its own shard count.
-	 */
 	readonly botId: GuildListKey;
 	readonly intents: GatewayIntentBits;
 	readonly rest: REST;
@@ -46,8 +41,6 @@ export async function createBotGateway({
 	const { shardIds } = await claimReplicaSlot({
 		botId,
 		shardCount,
-		// Unset means this replica takes everything, so the claim resolves to a single index. Same code path,
-		// same redis keys, just a cluster of one.
 		shardsPerReplica: ENV.SHARDS_PER_REPLICA ?? shardCount,
 	});
 
@@ -60,9 +53,6 @@ export async function createBotGateway({
 		shardIds,
 		retrieveSessionInfo: async (shardId) => sessions.retrieveSessionInfo(shardId),
 		updateSessionInfo: async (shardId, sessionInfo) => sessions.updateSessionInfo(shardId, sessionInfo),
-		// Read off the manager rather than the `gatewayInformation` above, matching how the default builds
-		// `SimpleIdentifyThrottler` -- the manager caches its own fetch, so this costs nothing and stays correct
-		// if the value is ever re-read (e.g. after `updateShardCount`).
 		buildIdentifyThrottler: async (manager) =>
 			createRedisIdentifyThrottler(
 				botId,