feat: draining members and the index of their destinations

README.md

@@ -96,8 +96,9 @@ With `keyPrefix: 'myservice'`:
 | Key | Type | Owner | Purpose |
 |---|---|---|---|
 | `myservice:members` | set | pod | the set of memberIds known to be live |
-| `myservice:member:<memberId>` | hash with `address`, `load` | pod | live pod registration and load metric, TTL refreshed by heartbeat |
+| `myservice:member:<memberId>` | hash with `address`, `load` and, while draining, `draining` | pod | live pod registration and load metric, TTL refreshed by heartbeat |
 | `myservice:destination:<id>` | set of memberIds | coordinator + pod | pods currently serving this destination |
+| `myservice:owned:<memberId>` | set of destination ids | coordinator + pod | the destinations of a pod, with the `memberIndex` option |
 | `myservice:lock:<lockId>` | hash with `podId`, `destinationId`, metadata | pod | lockId routing for transaction-bound calls |
 
 ## Pod side: `Member`
@@ -132,6 +133,10 @@ await member.unregisterLock(lockId)
 const peers = await member.listPeerLoad()
 
 // Graceful shutdown:
+// Before stopping, stop receiving destinations while the ones held are moved or finished: the registry never binds a
+// draining pod, nor reassigns orphans to it, and listPeerLoad() reports the flag
+await member.setDraining(true)
+
 clearInterval(heartbeat)
 await member.deregister()
 await member.close()
@@ -169,14 +174,21 @@ if (lockRouting) {
 }
 
 // Other primitives:
-await registry.listLiveMembers()                  // [{ memberId, address, load }, ...]
+await registry.listLiveMembers()                  // [{ memberId, address, load, draining }, ...]
 await registry.pickMember({ destinationId: destId })
-await registry.addPodToDestination(destId, memberId)
+await registry.addPodToDestination(destId, memberId)  // false when the pod is draining: nothing is bound
 await registry.hasBinding(destId)
 await registry.deregisterDestination(destId)
+await registry.listOwned(memberId)                // the destinations of a pod, with the member index
 await registry.close()
 ```
 
+### Member index
+
+With `memberIndex: true`, on the `Registry` and on every `Member`, each binding also records the destination in the set of the pod, `myservice:owned:<memberId>`, in the same script or transaction: claims and reassignments add it to the new pod and remove it from the pods unbound, pruning a dead pod removes it, `addPodToDestination`, `deregisterDestination`, `member.addToDestination()` and `member.removeFromDestination()` do the same. `registry.listOwned(memberId)` and `member.listOwned()` read the set, so that the destinations of a pod are listed without scanning them all, i.e. to move them before the pod stops.
+
+The index is exact only if every writer of the destinations maintains it, this library and any script of the application alike: enable it everywhere at once, and fill the sets of the destinations bound before, with a scan, before relying on it.
+
 ## Resolution and failover
 
 `resolveDestination` reads the destination's pod set, filters by liveness, and applies the allocation strategy. The four cases:
@@ -191,6 +203,8 @@ await registry.close()
 
 All writes use `SADD` / `SREM` (atomic). Concurrent first-touch by two coordinators can produce a destination with two pods from the start. That's a valid steady state, not a corrupted one.
 
+A pod which called `member.setDraining(true)` is never picked, and the binding scripts refuse it even when it was picked from a list read before it started draining: `resolveDestination` then picks another pod, `addPodToDestination` returns `false`, and `pickAndRegister` answers 503 (`refusedMessage`, result `refused`), so that the caller retries on another pod. The flag is written again with the record of the pod when it expired. The destinations the pod already holds keep resolving to it until it deregisters or moves them.
+
 ## Allocation strategies
 
 Pluggable. Built-in: `round-robin` (default), `least-loaded`, `random`. Custom strategies implement:

src/helpers/pick-and-register.ts

@@ -4,12 +4,14 @@ import type { FastifyReplyFromHooks } from '@fastify/reply-from'
 import '@fastify/reply-from'
 import type { Registry } from '../registry.ts'
 
-export type PickAndRegisterResult = 'spawned' | 'unavailable' | 'upstream_error'
+export type PickAndRegisterResult = 'spawned' | 'unavailable' | 'upstream_error' | 'refused'
 
 export interface PickAndRegisterOptions {
   registerIdFrom: (responseBody: any) => string
   expectedStatus?: number
   unavailableMessage?: string
+  /** The message of the 503 sent when the pod started draining after it was picked, so the binding was refused. */
+  refusedMessage?: string
   onResult?: (result: PickAndRegisterResult) => void
 }
 
@@ -21,6 +23,7 @@ export function pickAndRegister (
     registerIdFrom,
     expectedStatus = 201,
     unavailableMessage = 'No pods available',
+    refusedMessage = 'The pod started draining, retry',
     onResult
   } = opts
 
@@ -48,7 +51,15 @@ export function pickAndRegister (
         if (res.statusCode === expectedStatus) {
           const id = registerIdFrom(body)
           registry.addPodToDestination(id, member.memberId).then(
-            () => {
+            (bound) => {
+              // The pod started draining between the pick and the binding: the resource it created is not routable,
+              // the caller retries on another pod
+              if (!bound) {
+                onResult?.('refused')
+                replyOut.code(503).send({ error: refusedMessage })
+                return
+              }
+
               onResult?.('spawned')
               replyOut.send(body)
             },

src/member.ts

@@ -12,6 +12,8 @@ export interface MemberOptions {
   keyPrefix?: string
   ttl?: number
   getLoad?: () => number
+  /** Maintains the set of the destinations of this member, see the `memberIndex` option of `Registry`. */
+  memberIndex?: boolean
 }
 
 export class Member {
@@ -22,6 +24,8 @@ export class Member {
   #keyPrefix: string
   #ttl: number
   #getLoad: () => number
+  #memberIndex: boolean
+  #draining: boolean
 
   constructor (opts: MemberOptions) {
     const { redis, owned } = connect(opts.redis, opts.redisOptions)
@@ -32,6 +36,8 @@ export class Member {
     this.#keyPrefix = opts.keyPrefix ?? 'coordinator'
     this.#ttl = opts.ttl ?? 30
     this.#getLoad = opts.getLoad ?? (() => 0)
+    this.#memberIndex = opts.memberIndex ?? false
+    this.#draining = false
   }
 
   get memberId (): string {
@@ -54,6 +60,10 @@ export class Member {
     return `${this.#keyPrefix}:destination:${destId}`
   }
 
+  #ownedKey (): string {
+    return `${this.#keyPrefix}:owned:${this.#memberId}`
+  }
+
   #lockKey (lockId: string): string {
     return `${this.#keyPrefix}:lock:${lockId}`
   }
@@ -71,9 +81,11 @@ export class Member {
   async #writeRecord (): Promise<void> {
     const pipeline = this.#redis.pipeline()
     pipeline.sadd(this.#membersKey(), this.#memberId)
+    // A record written again after it expired keeps the member draining
     pipeline.hset(this.#memberKey(), {
       address: this.#address,
-      load: String(this.#getLoad())
+      load: String(this.#getLoad()),
+      ...(this.#draining ? { draining: '1' } : {})
     })
     pipeline.expire(this.#memberKey(), this.#ttl)
     await pipeline.exec()
@@ -86,12 +98,42 @@ export class Member {
     await pipeline.exec()
   }
 
+  // A draining member keeps the destinations it holds, but is bound to no new one, nor reassigned any, until it
+  // resumes: the registry checks the flag atomically when it binds. The heartbeats keep the flag, and write it again
+  // with a record which expired; it goes when the member resumes or deregisters.
+  async setDraining (draining: boolean): Promise<void> {
+    this.#draining = draining
+
+    if (draining) {
+      await this.#redis.hset(this.#memberKey(), 'draining', '1')
+    } else {
+      await this.#redis.hdel(this.#memberKey(), 'draining')
+    }
+  }
+
+  get draining (): boolean {
+    return this.#draining
+  }
+
   async addToDestination (destinationId: string): Promise<void> {
-    await this.#redis.sadd(this.#destinationKey(destinationId), this.#memberId)
+    if (this.#memberIndex) {
+      await this.#redis.multi().sadd(this.#destinationKey(destinationId), this.#memberId).sadd(this.#ownedKey(), destinationId).exec()
+    } else {
+      await this.#redis.sadd(this.#destinationKey(destinationId), this.#memberId)
+    }
   }
 
   async removeFromDestination (destinationId: string): Promise<void> {
-    await this.#redis.srem(this.#destinationKey(destinationId), this.#memberId)
+    if (this.#memberIndex) {
+      await this.#redis.multi().srem(this.#destinationKey(destinationId), this.#memberId).srem(this.#ownedKey(), destinationId).exec()
+    } else {
+      await this.#redis.srem(this.#destinationKey(destinationId), this.#memberId)
+    }
+  }
+
+  // The destinations of this member, from the member index
+  async listOwned (): Promise<string[]> {
+    return this.#redis.smembers(this.#ownedKey())
   }
 
   async registerLock (
@@ -116,7 +158,7 @@ export class Member {
 
     const pipeline = this.#redis.pipeline()
     for (const id of memberIds) {
-      pipeline.hmget(this.#memberKey(id), 'address', 'load')
+      pipeline.hmget(this.#memberKey(id), 'address', 'load', 'draining')
     }
     const results = await pipeline.exec()
     if (!results) return []
@@ -128,7 +170,7 @@ export class Member {
       const address = fields[0]
       if (!address) continue
       const load = parseInt(fields[1] ?? '0', 10) || 0
-      live.push({ memberId: memberIds[i], address, load })
+      live.push({ memberId: memberIds[i], address, load, draining: fields[2] === '1' })
     }
     return live
   }