From f396aba73cedbfdc5c34dfedd0ceab847b98b1df Mon Sep 17 00:00:00 2001 From: Thomas Hart Date: Thu, 20 Aug 2026 22:32:05 +0000 Subject: [PATCH 1/2] feat: add consumer groups with partition assignment Split the log into partitions, hash keys with FNV-1a, and assign each partition to one member of a consumer group using range or round-robin. Offsets live on the group so a rebalance continues from the last commit. --- README.md | 10 ++ src/consumer-group.ts | 227 +++++++++++++++++++++++++++++++++++ src/index.ts | 15 +++ test/consumer-group.test.ts | 232 ++++++++++++++++++++++++++++++++++++ 4 files changed, 484 insertions(+) create mode 100644 src/consumer-group.ts create mode 100644 test/consumer-group.test.ts diff --git a/README.md b/README.md index 01804fa..d6a5199 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,9 @@ Message brokers hide a lot of machinery behind `publish` and `subscribe`. This r - **Time lag** as the age of the oldest unconsumed produce - **Sliding-window throughput** (circular time buckets) for produce and consume rates, so a quiet period after a burst reports ~0 +- **Partitioned log** with **key-based partitioning** (FNV-1a) so a key is sticky to one partition +- **Consumer groups** with **range** and **round-robin partition assignment**, **eager rebalance** on join/leave, and **group-level committed offsets** +- **Per-key ordering**: records for one key append in order on one partition, and only the assigned member reads them ## What's implemented - **Topic-based publish/subscribe with fan-out.** A `Broker` where subscribers register handlers against a topic and every publish fans out to all matching subscribers, with monotonic message ids, idempotent unsubscribe, and snapshot-consistent delivery (subscribing or unsubscribing during dispatch never changes who receives the in-flight message). @@ -55,6 +58,13 @@ Message brokers hide a lot of machinery behind `publish` and `subscribe`. This r - **Time lag** as the age of the oldest unconsumed produce - **Sliding-window throughput** (circular time buckets) for produce and consume rates, so a quiet period after a burst reports ~0 - **Throughput and lag metrics per topic.** `TopicMetrics` records `produced` / `consumed` per topic. Lag is `logEndOffset - committedOffset` (Kafka's formula), time lag is the age of the oldest unconsumed record, and produce/consume rates use a sliding window of circular time buckets. Two consumers on one topic keep independent commits and rates. +- **Partitioned log** with **key-based partitioning** (FNV-1a) so a key is sticky to one partition +- **Consumer groups** with **range** and **round-robin partition assignment**, **eager rebalance** on join/leave, and **group-level committed offsets** +- **Per-key ordering**: records for one key append in order on one partition, and only the assigned member reads them +- **Partitioned log** with **key-based partitioning** (FNV-1a) so a key is sticky to one partition +- **Consumer groups** with **range** and **round-robin partition assignment**, **eager rebalance** on join/leave, and **group-level committed offsets** +- **Per-key ordering**: records for one key append in order on one partition, and only the assigned member reads them +- **Consumer groups with partition assignment.** A `PartitionedTopic` is an append-only log split into N partitions. `produce(key, payload)` hashes the key with FNV-1a so that key always lands on the same partition. A `ConsumerGroup` assigns each partition to at most one member using Kafka's range assignor (consecutive slices, remainder on the first members) or round-robin (interleaved). Independent groups each see the full log. Offsets are stored on the group, so a rebalance hands a partition to a peer at the last committed offset instead of replaying from zero. A thrown handler stalls that partition until the next pump. ## Usage ```ts diff --git a/src/consumer-group.ts b/src/consumer-group.ts new file mode 100644 index 0000000..0d4dd83 --- /dev/null +++ b/src/consumer-group.ts @@ -0,0 +1,227 @@ +export type AssignmentStrategy = 'range' | 'round-robin' + +export interface LogRecord { + readonly partition: number + readonly offset: number + readonly key: string + readonly payload: T +} + +export type RecordHandler = (record: LogRecord) => void + +export interface GroupMember { + readonly id: number + assignment(): readonly number[] + leave(): void +} + +export function partitionForKey(key: string, partitionCount: number): number { + if (!Number.isInteger(partitionCount) || partitionCount < 1) { + throw new Error(`partitionCount must be a positive integer, got ${partitionCount}`) + } + let hash = 2166136261 + for (let i = 0; i < key.length; i++) { + hash ^= key.charCodeAt(i) + hash = Math.imul(hash, 16777619) + } + return (hash >>> 0) % partitionCount +} + +export function rangeAssign( + memberIds: readonly number[], + partitionCount: number, +): Map { + const ids = [...memberIds].sort((a, b) => a - b) + const assignment = new Map() + for (const id of ids) assignment.set(id, []) + if (ids.length === 0 || partitionCount < 1) return assignment + const n = ids.length + const base = Math.floor(partitionCount / n) + const extra = partitionCount % n + let partition = 0 + for (let i = 0; i < n; i++) { + const count = base + (i < extra ? 1 : 0) + const owned = assignment.get(ids[i]!)! + for (let j = 0; j < count; j++) owned.push(partition++) + } + return assignment +} + +export function roundRobinAssign( + memberIds: readonly number[], + partitionCount: number, +): Map { + const ids = [...memberIds].sort((a, b) => a - b) + const assignment = new Map() + for (const id of ids) assignment.set(id, []) + if (ids.length === 0) return assignment + for (let p = 0; p < partitionCount; p++) { + assignment.get(ids[p % ids.length]!)!.push(p) + } + return assignment +} + +export class PartitionedTopic { + readonly partitionCount: number + private readonly partitions: LogRecord[][] + private readonly listeners = new Set<() => void>() + + constructor(partitionCount: number) { + if (!Number.isInteger(partitionCount) || partitionCount < 1) { + throw new Error(`partitionCount must be a positive integer, got ${partitionCount}`) + } + this.partitionCount = partitionCount + this.partitions = Array.from({ length: partitionCount }, () => []) + } + + partitionFor(key: string): number { + return partitionForKey(key, this.partitionCount) + } + + produce(key: string, payload: T): LogRecord { + const partition = this.partitionFor(key) + const bucket = this.partitions[partition]! + const record: LogRecord = { + partition, + offset: bucket.length, + key, + payload, + } + bucket.push(record) + for (const listener of [...this.listeners]) listener() + return record + } + + log(partition: number): readonly LogRecord[] { + const bucket = this.partitions[partition] + if (!bucket) { + throw new Error(`partition ${partition} out of range (0..${this.partitionCount - 1})`) + } + return bucket + } + + endOffset(partition: number): number { + return this.log(partition).length + } + + watch(listener: () => void): () => void { + this.listeners.add(listener) + return () => { + this.listeners.delete(listener) + } + } +} + +interface MemberState { + readonly id: number + readonly handler: RecordHandler + active: boolean + partitions: number[] +} + +export class ConsumerGroup { + private readonly topic: PartitionedTopic + private readonly strategy: AssignmentStrategy + private readonly members: MemberState[] = [] + private readonly committed: number[] + private nextMemberId = 1 + private pumping = false + private pendingPump = false + + constructor(topic: PartitionedTopic, options?: { strategy?: AssignmentStrategy }) { + this.topic = topic + this.strategy = options?.strategy ?? 'range' + this.committed = Array.from({ length: topic.partitionCount }, () => 0) + topic.watch(() => this.pump()) + } + + join(handler: RecordHandler): GroupMember { + const member: MemberState = { + id: this.nextMemberId++, + handler, + active: true, + partitions: [], + } + this.members.push(member) + this.rebalance() + this.pump() + + return { + id: member.id, + assignment: () => member.partitions.slice(), + leave: () => this.leave(member), + } + } + + memberCount(): number { + return this.members.length + } + + committedOffset(partition: number): number { + const offset = this.committed[partition] + if (offset === undefined) { + throw new Error(`partition ${partition} out of range (0..${this.topic.partitionCount - 1})`) + } + return offset + } + + private leave(member: MemberState): void { + if (!member.active) return + member.active = false + member.partitions = [] + const i = this.members.indexOf(member) + if (i !== -1) this.members.splice(i, 1) + this.rebalance() + this.pump() + } + + private rebalance(): void { + const ids = this.members.map((m) => m.id) + const assigned = + this.strategy === 'round-robin' + ? roundRobinAssign(ids, this.topic.partitionCount) + : rangeAssign(ids, this.topic.partitionCount) + for (const member of this.members) { + member.partitions = assigned.get(member.id) ?? [] + } + } + + private pump(): void { + if (this.pumping) { + this.pendingPump = true + return + } + this.pumping = true + try { + do { + this.pendingPump = false + for (const member of [...this.members]) { + if (!member.active) continue + for (const partition of [...member.partitions]) { + this.drain(member, partition) + } + } + } while (this.pendingPump) + } finally { + this.pumping = false + } + } + + private drain(member: MemberState, partition: number): void { + const log = this.topic.log(partition) + while (member.active && member.partitions.includes(partition)) { + const offset = this.committed[partition] + if (offset === undefined || offset >= log.length) return + const record = log[offset] + if (!record) return + try { + member.handler(record) + } catch { + return + } + if (this.committed[partition] === offset) { + this.committed[partition] = offset + 1 + } + } + } +} diff --git a/src/index.ts b/src/index.ts index c08d152..8de6e48 100644 --- a/src/index.ts +++ b/src/index.ts @@ -76,3 +76,18 @@ export type { TopicSnapshot, TopicMetricsOptions, } from './metrics.js' + +export { + PartitionedTopic, + ConsumerGroup, + partitionForKey, + rangeAssign, + roundRobinAssign, +} from './consumer-group.js' + +export type { + AssignmentStrategy, + LogRecord, + RecordHandler, + GroupMember, +} from './consumer-group.js' diff --git a/test/consumer-group.test.ts b/test/consumer-group.test.ts new file mode 100644 index 0000000..01f95cc --- /dev/null +++ b/test/consumer-group.test.ts @@ -0,0 +1,232 @@ +import { describe, it, expect } from 'vitest' +import { + ConsumerGroup, + PartitionedTopic, + partitionForKey, + rangeAssign, + roundRobinAssign, + type LogRecord, +} from '../src/index.js' + +describe('partitionForKey', () => { + it('throws on a non-positive partition count', () => { + expect(() => partitionForKey('k', 0)).toThrow(/partitionCount/) + expect(() => partitionForKey('k', 1.5)).toThrow(/partitionCount/) + }) + + it('is stable and maps the same key to the same partition', () => { + const n = 8 + const a = partitionForKey('user-42', n) + expect(a).toBe(partitionForKey('user-42', n)) + expect(a).toBeGreaterThanOrEqual(0) + expect(a).toBeLessThan(n) + expect(partitionForKey('', n)).toBe(partitionForKey('', n)) + }) + + it('spreads distinct keys across partitions', () => { + const seen = new Set() + for (let i = 0; i < 40; i++) seen.add(partitionForKey(`k-${i}`, 4)) + expect(seen.size).toBe(4) + }) +}) + +describe('rangeAssign and roundRobinAssign', () => { + it('gives consecutive ranges and dumps remainder on the first members', () => { + expect(rangeAssign([2, 1], 6)).toEqual( + new Map([ + [1, [0, 1, 2]], + [2, [3, 4, 5]], + ]), + ) + expect(rangeAssign([1, 2, 3], 5)).toEqual( + new Map([ + [1, [0, 1]], + [2, [2, 3]], + [3, [4]], + ]), + ) + }) + + it('leaves extra members idle when there are more members than partitions', () => { + expect(rangeAssign([1, 2, 3, 4, 5], 3)).toEqual( + new Map([ + [1, [0]], + [2, [1]], + [3, [2]], + [4, []], + [5, []], + ]), + ) + }) + + it('interleaves partitions for round-robin and sorts members by id', () => { + expect(roundRobinAssign([3, 1], 6)).toEqual( + new Map([ + [1, [0, 2, 4]], + [3, [1, 3, 5]], + ]), + ) + expect(rangeAssign([], 4).size).toBe(0) + expect(roundRobinAssign([], 4).size).toBe(0) + }) +}) + +describe('PartitionedTopic', () => { + it('rejects a bad partition count', () => { + expect(() => new PartitionedTopic(0)).toThrow(/partitionCount/) + expect(() => new PartitionedTopic(-1)).toThrow(/partitionCount/) + }) + + it('appends monotonic offsets per partition and preserves per-key order', () => { + const topic = new PartitionedTopic(4) + const key = 'acct-7' + const p = topic.partitionFor(key) + const first = topic.produce(key, 'a') + const second = topic.produce(key, 'b') + expect(first.partition).toBe(p) + expect(second.partition).toBe(p) + expect(first.offset).toBe(0) + expect(second.offset).toBe(1) + expect(topic.log(p).map((r) => r.payload)).toEqual(['a', 'b']) + expect(topic.endOffset(p)).toBe(2) + expect(() => topic.log(4)).toThrow(/out of range/) + }) +}) + +describe('ConsumerGroup', () => { + it('delivers each record to exactly one member of the group', () => { + const topic = new PartitionedTopic(4) + const group = new ConsumerGroup(topic) + const a: number[] = [] + const b: number[] = [] + group.join((r) => a.push(r.payload)) + group.join((r) => b.push(r.payload)) + + for (let i = 0; i < 12; i++) topic.produce(`k-${i}`, i) + + const all = [...a, ...b].sort((x, y) => x - y) + expect(all).toEqual([...Array(12).keys()]) + expect(a.length).toBeGreaterThan(0) + expect(b.length).toBeGreaterThan(0) + expect(new Set(all).size).toBe(12) + }) + + it('gives every record to each independent group', () => { + const topic = new PartitionedTopic(3) + const g1 = new ConsumerGroup(topic) + const g2 = new ConsumerGroup(topic) + const one: string[] = [] + const two: string[] = [] + g1.join((r) => one.push(r.payload)) + g2.join((r) => two.push(r.payload)) + topic.produce('a', 'x') + topic.produce('b', 'y') + expect(one.sort()).toEqual(['x', 'y']) + expect(two.sort()).toEqual(['x', 'y']) + }) + + it('applies range and round-robin strategies on join', () => { + const rangeTopic = new PartitionedTopic(6) + const range = new ConsumerGroup(rangeTopic) + const r1 = range.join(() => {}) + const r2 = range.join(() => {}) + expect(r1.assignment()).toEqual([0, 1, 2]) + expect(r2.assignment()).toEqual([3, 4, 5]) + + const rr = new ConsumerGroup(new PartitionedTopic(6), { + strategy: 'round-robin', + }) + const a = rr.join(() => {}) + const b = rr.join(() => {}) + expect(a.assignment()).toEqual([0, 2, 4]) + expect(b.assignment()).toEqual([1, 3, 5]) + }) + + it('rebalances on join and leave, continuing from the group commit', () => { + const topic = new PartitionedTopic(2) + const group = new ConsumerGroup(topic) + const seen: string[] = [] + const hold: LogRecord[] = [] + const first = group.join((r) => { + seen.push(`1:${r.payload}`) + hold.push(r) + }) + expect(first.assignment()).toEqual([0, 1]) + + topic.produce('p0-only', 'a') + expect(seen).toEqual(['1:a']) + expect(group.committedOffset(hold[0]!.partition)).toBe(1) + + const secondSeen: string[] = [] + const second = group.join((r) => secondSeen.push(r.payload)) + expect(first.assignment().length + second.assignment().length).toBe(2) + expect(new Set([...first.assignment(), ...second.assignment()]).size).toBe(2) + + first.leave() + expect(group.memberCount()).toBe(1) + expect(second.assignment()).toEqual([0, 1]) + + topic.produce('p0-only', 'b') + topic.produce('other', 'c') + expect(secondSeen).toContain('b') + expect(secondSeen).toContain('c') + expect(seen).toEqual(['1:a']) + }) + + it('catches up a late joiner from offset zero', () => { + const topic = new PartitionedTopic(2) + topic.produce('k', 'old') + const group = new ConsumerGroup(topic) + const seen: string[] = [] + group.join((r) => seen.push(r.payload)) + expect(seen).toEqual(['old']) + }) + + it('stalls a partition on handler throw and retries on the next pump', () => { + const topic = new PartitionedTopic(1) + const group = new ConsumerGroup(topic) + let fail = true + const seen: string[] = [] + group.join((r) => { + if (fail) throw new Error('boom') + seen.push(r.payload) + }) + topic.produce('k', 'one') + expect(group.committedOffset(0)).toBe(0) + expect(seen).toEqual([]) + + fail = false + topic.produce('k', 'two') + expect(seen).toEqual(['one', 'two']) + expect(group.committedOffset(0)).toBe(2) + }) + + it('keeps per-key order on the member that owns the partition', () => { + const topic = new PartitionedTopic(8) + const group = new ConsumerGroup(topic) + const byKey = new Map() + const take = (r: LogRecord) => { + const list = byKey.get(r.key) ?? [] + list.push(r.payload) + byKey.set(r.key, list) + } + group.join(take) + group.join(take) + for (let i = 0; i < 5; i++) topic.produce('alpha', i) + expect(byKey.get('alpha')).toEqual([0, 1, 2, 3, 4]) + }) + + it('leave is idempotent and extra members stay idle', () => { + const topic = new PartitionedTopic(1) + const group = new ConsumerGroup(topic) + const owner = group.join(() => {}) + const idle = group.join(() => {}) + expect(owner.assignment()).toEqual([0]) + expect(idle.assignment()).toEqual([]) + owner.leave() + owner.leave() + expect(group.memberCount()).toBe(1) + expect(idle.assignment()).toEqual([0]) + expect(() => group.committedOffset(1)).toThrow(/out of range/) + }) +}) From e7e57f42ab07cddc78ff8c3a50ad6ca7233471ed Mon Sep 17 00:00:00 2001 From: Thomas Hart Date: Thu, 20 Aug 2026 22:42:25 +0000 Subject: [PATCH 2/2] test: assert group commits survive rebalance and empty groups Pin rebalance keys to partitions 0/1 and require the successor not to see the already-committed payload. Cover last-member leave and a throw that stalls only one partition. --- test/consumer-group.test.ts | 63 ++++++++++++++++++++++++++++++++++--- 1 file changed, 58 insertions(+), 5 deletions(-) diff --git a/test/consumer-group.test.ts b/test/consumer-group.test.ts index 01f95cc..2a4991b 100644 --- a/test/consumer-group.test.ts +++ b/test/consumer-group.test.ts @@ -14,6 +14,12 @@ describe('partitionForKey', () => { expect(() => partitionForKey('k', 1.5)).toThrow(/partitionCount/) }) + it('matches the FNV-1a 32-bit golden vectors', () => { + const fnvMod = 0x1_0000_0000 + expect(partitionForKey('', fnvMod)).toBe(0x811c9dc5) + expect(partitionForKey('hello', fnvMod)).toBe(0x4f9f2cab) + }) + it('is stable and maps the same key to the same partition', () => { const n = 8 const a = partitionForKey('user-42', n) @@ -144,6 +150,9 @@ describe('ConsumerGroup', () => { it('rebalances on join and leave, continuing from the group commit', () => { const topic = new PartitionedTopic(2) + expect(partitionForKey('p0-only', 2)).toBe(0) + expect(partitionForKey('other', 2)).toBe(1) + const group = new ConsumerGroup(topic) const seen: string[] = [] const hold: LogRecord[] = [] @@ -155,24 +164,50 @@ describe('ConsumerGroup', () => { topic.produce('p0-only', 'a') expect(seen).toEqual(['1:a']) - expect(group.committedOffset(hold[0]!.partition)).toBe(1) + expect(hold[0]!.partition).toBe(0) + expect(group.committedOffset(0)).toBe(1) + expect(group.committedOffset(1)).toBe(0) const secondSeen: string[] = [] const second = group.join((r) => secondSeen.push(r.payload)) - expect(first.assignment().length + second.assignment().length).toBe(2) - expect(new Set([...first.assignment(), ...second.assignment()]).size).toBe(2) + expect(first.assignment()).toEqual([0]) + expect(second.assignment()).toEqual([1]) first.leave() expect(group.memberCount()).toBe(1) expect(second.assignment()).toEqual([0, 1]) + expect(group.committedOffset(0)).toBe(1) topic.produce('p0-only', 'b') topic.produce('other', 'c') - expect(secondSeen).toContain('b') - expect(secondSeen).toContain('c') + expect(secondSeen.sort()).toEqual(['b', 'c']) + expect(secondSeen).not.toContain('a') expect(seen).toEqual(['1:a']) }) + it('keeps committed offsets after the last member leaves', () => { + const topic = new PartitionedTopic(2) + expect(partitionForKey('p0-only', 2)).toBe(0) + + const group = new ConsumerGroup(topic) + const firstSeen: string[] = [] + const first = group.join((r) => firstSeen.push(r.payload)) + topic.produce('p0-only', 'first') + expect(firstSeen).toEqual(['first']) + expect(group.committedOffset(0)).toBe(1) + + first.leave() + expect(group.memberCount()).toBe(0) + expect(group.committedOffset(0)).toBe(1) + + topic.produce('p0-only', 'second') + const later: string[] = [] + group.join((r) => later.push(r.payload)) + expect(later).toEqual(['second']) + expect(later).not.toContain('first') + expect(group.committedOffset(0)).toBe(2) + }) + it('catches up a late joiner from offset zero', () => { const topic = new PartitionedTopic(2) topic.produce('k', 'old') @@ -201,6 +236,24 @@ describe('ConsumerGroup', () => { expect(group.committedOffset(0)).toBe(2) }) + it('stalls only the throwing partition and still drains the rest', () => { + const topic = new PartitionedTopic(2) + expect(partitionForKey('p0-only', 2)).toBe(0) + expect(partitionForKey('other', 2)).toBe(1) + + const group = new ConsumerGroup(topic) + const seen: string[] = [] + group.join((r) => { + if (r.partition === 0) throw new Error('boom') + seen.push(r.payload) + }) + topic.produce('p0-only', 'bad') + topic.produce('other', 'ok') + expect(seen).toEqual(['ok']) + expect(group.committedOffset(0)).toBe(0) + expect(group.committedOffset(1)).toBe(1) + }) + it('keeps per-key order on the member that owns the partition', () => { const topic = new PartitionedTopic(8) const group = new ConsumerGroup(topic)