kafka-ts 0.0.2-beta → 0.0.3-beta

package/src/consumer/consumer.ts

@@ -1,15 +1,16 @@
-import { API, API_ERROR } from "../api";
-import { IsolationLevel } from "../api/fetch";
-import { Assignment } from "../api/sync-group";
-import { Cluster } from "../cluster";
-import { distributeAssignmentsToNodes } from "../distributors/assignments-to-replicas";
-import { Message } from "../types";
-import { delay } from "../utils/delay";
-import { ConnectionError, KafkaTSApiError } from "../utils/error";
-import { defaultRetrier, Retrier } from "../utils/retrier";
-import { ConsumerGroup } from "./consumer-group";
-import { ConsumerMetadata } from "./consumer-metadata";
-import { OffsetManager } from "./offset-manager";
+import { API, API_ERROR } from '../api';
+import { IsolationLevel } from '../api/fetch';
+import { Assignment } from '../api/sync-group';
+import { Cluster } from '../cluster';
+import { distributeAssignmentsToNodes } from '../distributors/assignments-to-replicas';
+import { Message } from '../types';
+import { delay } from '../utils/delay';
+import { ConnectionError, KafkaTSApiError } from '../utils/error';
+import { defaultRetrier, Retrier } from '../utils/retrier';
+import { ConsumerGroup } from './consumer-group';
+import { ConsumerMetadata } from './consumer-metadata';
+import { FetchManager, Granularity } from './fetch-manager';
+import { OffsetManager } from './offset-manager';
 
 export type ConsumerOptions = {
     topics: string[];
@@ -26,13 +27,16 @@ export type ConsumerOptions = {
     allowTopicAutoCreation?: boolean;
     fromBeginning?: boolean;
     retrier?: Retrier;
-} & ({ onMessage: (message: Message) => unknown } | { onBatch: (messages: Message[]) => unknown });
+    granularity?: Granularity;
+    concurrency?: number;
+} & ({ onBatch: (messages: Required<Message>[]) => unknown } | { onMessage: (message: Required<Message>) => unknown });
 
 export class Consumer {
     private options: Required<ConsumerOptions>;
     private metadata: ConsumerMetadata;
     private consumerGroup: ConsumerGroup | undefined;
     private offsetManager: OffsetManager;
+    private fetchManager?: FetchManager;
     private stopHook: (() => void) | undefined;
 
     constructor(
@@ -43,7 +47,7 @@ export class Consumer {
             ...options,
             groupId: options.groupId ?? null,
             groupInstanceId: options.groupInstanceId ?? null,
-            rackId: options.rackId ?? "",
+            rackId: options.rackId ?? '',
             sessionTimeoutMs: options.sessionTimeoutMs ?? 30_000,
             rebalanceTimeoutMs: options.rebalanceTimeoutMs ?? 60_000,
             maxWaitMs: options.maxWaitMs ?? 5000,
@@ -54,6 +58,8 @@ export class Consumer {
             allowTopicAutoCreation: options.allowTopicAutoCreation ?? false,
             fromBeginning: options.fromBeginning ?? false,
             retrier: options.retrier ?? defaultRetrier,
+            granularity: options.granularity ?? 'broker',
+            concurrency: options.concurrency ?? 1,
         };
 
         this.metadata = new ConsumerMetadata({ cluster: this.cluster });
@@ -95,81 +101,63 @@ export class Consumer {
             if (this.stopHook) return (this.stopHook as () => void)();
             return this.close(true).then(() => this.start());
         }
-        this.fetchLoop();
+        setImmediate(() => this.startFetchManager());
     }
 
-    private fetchLoop = async () => {
-        const { options } = this;
-        const { retrier } = options;
+    public async close(force = false): Promise<void> {
+        if (!force) {
+            await new Promise<void>(async (resolve) => {
+                this.stopHook = resolve;
+                await this.fetchManager?.stop();
+            });
+        }
+        await this.consumerGroup
+            ?.leaveGroup()
+            .catch((error) => console.warn(`Failed to leave group: ${error.message}`));
+        await this.cluster.disconnect().catch((error) => console.warn(`Failed to disconnect: ${error.message}`));
+    }
 
-        let nodeAssignments: { nodeId: number; assignment: Assignment }[] = [];
-        let shouldReassign = true;
+    private startFetchManager = async () => {
+        const { granularity, concurrency } = this.options;
 
         while (!this.stopHook) {
-            if (shouldReassign || !nodeAssignments) {
-                nodeAssignments = Object.entries(
-                    distributeAssignmentsToNodes(
-                        this.metadata.getAssignment(),
-                        this.metadata.getTopicPartitionReplicaIds(),
-                    ),
-                ).map(([nodeId, assignment]) => ({ nodeId: parseInt(nodeId), assignment }));
-                shouldReassign = false;
-            }
+            const nodeAssignments = Object.entries(
+                distributeAssignmentsToNodes(
+                    this.metadata.getAssignment(),
+                    this.metadata.getTopicPartitionReplicaIds(),
+                ),
+            ).map(([nodeId, assignment]) => ({ nodeId: parseInt(nodeId), assignment }));
+
+            const numPartitions = Object.values(this.metadata.getAssignment()).flat().length;
+            const numProcessors = Math.min(concurrency, numPartitions);
+
+            this.fetchManager = new FetchManager({
+                fetch: this.fetch.bind(this),
+                process: this.process.bind(this),
+                metadata: this.metadata,
+                consumerGroup: this.consumerGroup,
+                nodeAssignments,
+                granularity,
+                concurrency: numProcessors,
+            });
 
             try {
-                for (const { nodeId, assignment } of nodeAssignments) {
-                    const batch = await this.fetch(nodeId, assignment);
-                    const messages = batch.responses.flatMap(({ topicId, partitions }) =>
-                        partitions.flatMap(({ partitionIndex, records }) =>
-                            records.flatMap(({ baseTimestamp, baseOffset, records }) =>
-                                records.map(
-                                    (message): Required<Message> => ({
-                                        topic: this.metadata.getTopicNameById(topicId),
-                                        partition: partitionIndex,
-                                        key: message.key ?? null,
-                                        value: message.value ?? null,
-                                        headers: Object.fromEntries(
-                                            message.headers.map(({ key, value }) => [key, value]),
-                                        ),
-                                        timestamp: baseTimestamp + BigInt(message.timestampDelta),
-                                        offset: baseOffset + BigInt(message.offsetDelta),
-                                    }),
-                                ),
-                            ),
-                        ),
-                    );
-
-                    if ("onBatch" in options) {
-                        await retrier(() => options.onBatch(messages));
-
-                        messages.forEach(({ topic, partition, offset }) =>
-                            this.offsetManager.resolve(topic, partition, offset + 1n),
-                        );
-                    } else if ("onMessage" in options) {
-                        for (const message of messages) {
-                            await retrier(() => options.onMessage(message));
-
-                            const { topic, partition, offset } = message;
-                            this.offsetManager.resolve(topic, partition, offset + 1n);
-                        }
-                    }
-                    await this.consumerGroup?.offsetCommit();
-                    await this.consumerGroup?.handleLastHeartbeat();
-                }
+                await this.fetchManager.start();
 
                 if (!nodeAssignments.length) {
-                    console.debug("No partitions assigned. Waiting for reassignment...");
+                    console.debug('No partitions assigned. Waiting for reassignment...');
                     await delay(this.options.maxWaitMs);
                     await this.consumerGroup?.handleLastHeartbeat();
                 }
             } catch (error) {
+                await this.fetchManager.stop();
+
                 if ((error as KafkaTSApiError).errorCode === API_ERROR.REBALANCE_IN_PROGRESS) {
-                    console.debug("Rebalance in progress...");
-                    shouldReassign = true;
+                    console.debug('Rebalance in progress...');
                     continue;
                 }
                 if ((error as KafkaTSApiError).errorCode === API_ERROR.FENCED_INSTANCE_ID) {
-                    console.debug("New consumer with the same groupInstanceId joined. Exiting the consumer...");
+                    console.debug('New consumer with the same groupInstanceId joined. Exiting the consumer...');
                     this.close();
                     break;
                 }
@@ -182,23 +170,37 @@ export class Consumer {
                     break;
                 }
                 console.error(error);
-                await this.consumerGroup?.offsetCommit();
+                this.close();
                 break;
             }
         }
         this.stopHook?.();
     };
 
-    public async close(force = false): Promise<void> {
-        if (!force) {
-            await new Promise<void>((resolve) => {
-                this.stopHook = resolve;
-            });
+    private async process(messages: Required<Message>[]) {
+        const { options } = this;
+        const { retrier } = options;
+
+        if ('onBatch' in options) {
+            await retrier(() => options.onBatch(messages));
+
+            messages.forEach(({ topic, partition, offset }) =>
+                this.offsetManager.resolve(topic, partition, offset + 1n),
+            );
+        } else if ('onMessage' in options) {
+            try {
+                for (const message of messages) {
+                    await retrier(() => options.onMessage(message));
+
+                    const { topic, partition, offset } = message;
+                    this.offsetManager.resolve(topic, partition, offset + 1n);
+                }
+            } catch (error) {
+                await this.consumerGroup?.offsetCommit().catch(() => {});
+                throw error;
+            }
         }
-        await this.consumerGroup
-            ?.leaveGroup()
-            .catch((error) => console.warn(`Failed to leave group: ${error.message}`));
-        await this.cluster.disconnect().catch((error) => console.warn(`Failed to disconnect: ${error.message}`));
+        await this.consumerGroup?.offsetCommit();
     }
 
     private fetch(nodeId: number, assignment: Assignment) {

package/src/consumer/fetch-manager.ts

@@ -0,0 +1,179 @@
+import EventEmitter from 'events';
+import { API } from '../api';
+import { Assignment } from '../api/sync-group';
+import { Metadata } from '../metadata';
+import { Batch, Message } from '../types';
+import { KafkaTSError } from '../utils/error';
+import { createTracer } from '../utils/tracer';
+import { ConsumerGroup } from './consumer-group';
+import { Fetcher } from './fetcher';
+import { Processor } from './processor';
+
+const trace = createTracer('FetchManager');
+
+export type Granularity = 'partition' | 'topic' | 'broker';
+
+type FetchManagerOptions = {
+    fetch: (nodeId: number, assignment: Assignment) => Promise<ReturnType<(typeof API.FETCH)['response']>>;
+    process: (batch: Batch) => Promise<void>;
+    metadata: Metadata;
+    consumerGroup?: ConsumerGroup;
+    nodeAssignments: { nodeId: number; assignment: Assignment }[];
+    granularity: Granularity;
+    concurrency: number;
+};
+
+type Checkpoint = { kind: 'checkpoint'; fetcherId: number };
+type Entry = Batch | Checkpoint;
+
+export class FetchManager extends EventEmitter<{ data: []; checkpoint: [number]; stop: [] }> {
+    private queue: Entry[] = [];
+    private isRunning = false;
+    private fetchers: Fetcher[];
+    private processors: Processor[];
+
+    constructor(private options: FetchManagerOptions) {
+        super();
+
+        const { fetch, process, consumerGroup, nodeAssignments, concurrency } = this.options;
+
+        this.fetchers = nodeAssignments.map(
+            ({ nodeId, assignment }, index) =>
+                new Fetcher(index, {
+                    nodeId,
+                    assignment,
+                    consumerGroup,
+                    fetch,
+                    onResponse: this.onResponse.bind(this),
+                }),
+        );
+        this.processors = Array.from({ length: concurrency }).map(
+            () => new Processor({ process, poll: this.poll.bind(this) }),
+        );
+    }
+
+    public async start() {
+        this.queue = [];
+        this.isRunning = true;
+
+        try {
+            await Promise.all([
+                ...this.fetchers.map((fetcher) => fetcher.loop()),
+                ...this.processors.map((processor) => processor.loop()),
+            ]);
+        } finally {
+            this.isRunning = false;
+            this.emit('stop');
+        }
+    }
+
+    @trace()
+    public async stop() {
+        this.isRunning = false;
+        this.emit('stop');
+
+        await Promise.all([
+            ...this.fetchers.map((fetcher) => fetcher.stop()),
+            ...this.processors.map((processor) => processor.stop()),
+        ]);
+    }
+
+    @trace()
+    public async poll(): Promise<Batch> {
+        if (!this.isRunning) {
+            return [];
+        }
+
+        const batch = this.queue.shift();
+        if (!batch) {
+            await new Promise<void>((resolve) => {
+                const onData = () => {
+                    this.removeListener('stop', onStop);
+                    resolve();
+                };
+                const onStop = () => {
+                    this.removeListener('data', onData);
+                    resolve();
+                };
+                this.once('data', onData);
+                this.once('stop', onStop);
+            });
+            return this.poll();
+        }
+
+        if ('kind' in batch && batch.kind === 'checkpoint') {
+            this.emit('checkpoint', batch.fetcherId);
+            return this.poll();
+        }
+
+        return batch as Exclude<Entry, Checkpoint>;
+    }
+
+    private async onResponse(fetcherId: number, response: ReturnType<(typeof API.FETCH)['response']>) {
+        const { metadata, granularity } = this.options;
+
+        const batches = fetchResponseToBatches(response, granularity, metadata);
+        if (batches.length) {
+            this.queue.push(...batches);
+            this.queue.push({ kind: 'checkpoint', fetcherId });
+
+            this.emit('data');
+            await new Promise<void>((resolve) => {
+                const onCheckpoint = (id: number) => {
+                    if (id === fetcherId) {
+                        this.removeListener('stop', onStop);
+                        resolve();
+                    }
+                };
+                const onStop = () => {
+                    this.removeListener('checkpoint', onCheckpoint);
+                    resolve();
+                };
+                this.once('checkpoint', onCheckpoint);
+                this.once('stop', onStop);
+            });
+        }
+    }
+}
+
+const fetchResponseToBatches = (
+    batch: ReturnType<typeof API.FETCH.response>,
+    granularity: Granularity,
+    metadata: Metadata,
+): Batch[] => {
+    const brokerTopics = batch.responses.map(({ topicId, partitions }) =>
+        partitions.map(({ partitionIndex, records }) =>
+            records.flatMap(({ baseTimestamp, baseOffset, records }) =>
+                records.map(
+                    (message): Required<Message> => ({
+                        topic: metadata.getTopicNameById(topicId),
+                        partition: partitionIndex,
+                        key: message.key ?? null,
+                        value: message.value ?? null,
+                        headers: Object.fromEntries(message.headers.map(({ key, value }) => [key, value])),
+                        timestamp: baseTimestamp + BigInt(message.timestampDelta),
+                        offset: baseOffset + BigInt(message.offsetDelta),
+                    }),
+                ),
+            ),
+        ),
+    );
+
+    switch (granularity) {
+        case 'broker':
+            const messages = brokerTopics.flatMap((topicPartition) =>
+                topicPartition.flatMap((partitionMessages) => partitionMessages),
+            );
+            return messages.length ? [messages] : [];
+        case 'topic':
+            return brokerTopics
+                .map((topicPartition) => topicPartition.flatMap((partitionMessages) => partitionMessages))
+                .filter((messages) => messages.length);
+        case 'partition':
+            return brokerTopics.flatMap((topicPartition) =>
+                topicPartition.map((partitionMessages) => partitionMessages),
+            );
+        default:
+            throw new KafkaTSError(`Unhandled batch granularity: ${granularity}`);
+    }
+};

package/src/consumer/fetcher.ts

@@ -0,0 +1,57 @@
+import { EventEmitter } from 'stream';
+import { API } from '../api';
+import { Assignment } from '../api/sync-group';
+import { createTracer } from '../utils/tracer';
+import { ConsumerGroup } from './consumer-group';
+
+const trace = createTracer('Fetcher');
+
+type FetcherOptions = {
+    nodeId: number;
+    assignment: Assignment;
+    consumerGroup?: ConsumerGroup;
+    fetch: (nodeId: number, assignment: Assignment) => Promise<ReturnType<(typeof API.FETCH)['response']>>;
+    onResponse: (fetcherId: number, response: ReturnType<(typeof API.FETCH)['response']>) => Promise<void>;
+};
+
+export class Fetcher extends EventEmitter<{ stop: []; stopped: []; data: []; drain: [] }> {
+    private isRunning = false;
+
+    constructor(
+        private fetcherId: number,
+        private options: FetcherOptions,
+    ) {
+        super();
+    }
+
+    public async loop() {
+        const { nodeId, assignment, consumerGroup, fetch, onResponse } = this.options;
+
+        this.isRunning = true;
+        this.once('stop', () => (this.isRunning = false));
+
+        try {
+            while (this.isRunning) {
+                const response = await fetch(nodeId, assignment);
+                await consumerGroup?.handleLastHeartbeat();
+                await onResponse(this.fetcherId, response);
+                await consumerGroup?.handleLastHeartbeat();
+            }
+        } finally {
+            this.isRunning = false;
+            this.emit('stopped');
+        }
+    }
+
+    @trace()
+    public async stop() {
+        if (!this.isRunning) {
+            return;
+        }
+
+        this.emit('stop');
+        return new Promise<void>((resolve) => {
+            this.once('stopped', resolve);
+        });
+    }
+}

package/src/consumer/offset-manager.ts

@@ -1,9 +1,9 @@
-import { API } from "../api";
-import { IsolationLevel } from "../api/fetch";
-import { Assignment } from "../api/sync-group";
-import { Cluster } from "../cluster";
-import { distributeMessagesToTopicPartitionLeaders } from "../distributors/messages-to-topic-partition-leaders";
-import { ConsumerMetadata } from "./consumer-metadata";
+import { API } from '../api';
+import { IsolationLevel } from '../api/fetch';
+import { Assignment } from '../api/sync-group';
+import { Cluster } from '../cluster';
+import { distributeMessagesToTopicPartitionLeaders } from '../distributors/messages-to-topic-partition-leaders';
+import { ConsumerMetadata } from './consumer-metadata';
 
 type OffsetManagerOptions = {
     cluster: Cluster;

package/src/consumer/processor.ts

@@ -0,0 +1,47 @@
+import { EventEmitter } from 'stream';
+import { Batch } from '../types';
+import { createTracer } from '../utils/tracer';
+
+const trace = createTracer('Processor');
+
+type ProcessorOptions = {
+    poll: () => Promise<Batch>;
+    process: (batch: Batch) => Promise<void>;
+};
+
+export class Processor extends EventEmitter<{ stop: []; stopped: [] }> {
+    private isRunning = false;
+
+    constructor(private options: ProcessorOptions) {
+        super();
+    }
+
+    public async loop() {
+        const { poll, process } = this.options;
+
+        this.isRunning = true;
+        this.once('stop', () => (this.isRunning = false));
+
+        try {
+            while (this.isRunning) {
+                const batch = await poll();
+                await process(batch);
+            }
+        } finally {
+            this.isRunning = false;
+            this.emit('stopped');
+        }
+    }
+
+    @trace()
+    public async stop() {
+        if (!this.isRunning) {
+            return;
+        }
+
+        return new Promise<void>((resolve) => {
+            this.once('stopped', resolve);
+            this.emit('stop');
+        });
+    }
+}

package/src/distributors/assignments-to-replicas.test.ts

@@ -1,9 +1,9 @@
-import { describe, expect, it } from "vitest";
-import { distributeAssignmentsToNodesBalanced, distributeAssignmentsToNodesOptimized } from "./assignments-to-replicas";
+import { describe, expect, it } from 'vitest';
+import { distributeAssignmentsToNodesBalanced, distributeAssignmentsToNodesOptimized } from './assignments-to-replicas';
 
-describe("Distribute assignments to replica ids", () => {
-    describe("distributeAssignmentsToNodesBalanced", () => {
-        it("smoke", () => {
+describe('Distribute assignments to replica ids', () => {
+    describe('distributeAssignmentsToNodesBalanced', () => {
+        it('smoke', () => {
             const result = distributeAssignmentsToNodesBalanced({ topic: [0, 1] }, { topic: { 0: [0, 1], 1: [1, 2] } });
             expect(result).toMatchInlineSnapshot(`
               {
@@ -22,8 +22,8 @@ describe("Distribute assignments to replica ids", () => {
         });
     });
 
-    describe("distributeAssignmentsToNodesOptimized", () => {
-        it("smoke", () => {
+    describe('distributeAssignmentsToNodesOptimized', () => {
+        it('smoke', () => {
             const result = distributeAssignmentsToNodesOptimized(
                 { topic: [0, 1] },
                 { topic: { 0: [0, 1], 1: [1, 2] } },

package/src/distributors/assignments-to-replicas.ts

@@ -50,7 +50,7 @@ export const distributeAssignmentsToNodesOptimized = (
         }
 
         result[parseInt(replicaId)] = partitions.reduce((acc, partition) => {
-            const [topicName, partitionId] = partition.split(":");
+            const [topicName, partitionId] = partition.split(':');
             acc[topicName] ??= [];
             acc[topicName].push(parseInt(partitionId));
             return acc;

package/src/distributors/messages-to-topic-partition-leaders.test.ts

@@ -1,11 +1,11 @@
-import { describe, expect, it } from "vitest";
-import { distributeMessagesToTopicPartitionLeaders } from "./messages-to-topic-partition-leaders";
+import { describe, expect, it } from 'vitest';
+import { distributeMessagesToTopicPartitionLeaders } from './messages-to-topic-partition-leaders';
 
-describe("Distribute messages to partition leader ids", () => {
-    describe("distributeMessagesToTopicPartitionLeaders", () => {
-        it("snoke", () => {
+describe('Distribute messages to partition leader ids', () => {
+    describe('distributeMessagesToTopicPartitionLeaders', () => {
+        it('snoke', () => {
             const result = distributeMessagesToTopicPartitionLeaders(
-                [{ topic: "topic", partition: 0, key: null, value: null, offset: 0n, timestamp: 0n, headers: {} }],
+                [{ topic: 'topic', partition: 0, key: null, value: null, offset: 0n, timestamp: 0n, headers: {} }],
                 { topic: { 0: 1 } },
             );
             expect(result).toMatchInlineSnapshot(`

package/src/index.ts

@@ -1,3 +1,4 @@
-export * from "./client";
-export * from "./api";
-export * from "./types";
+export * from './utils/error';
+export * from './client';
+export * from './api';
+export * from './types';

package/src/metadata.ts

@@ -1,7 +1,7 @@
-import { API, API_ERROR } from "./api";
-import { Cluster } from "./cluster";
-import { delay } from "./utils/delay";
-import { KafkaTSApiError } from "./utils/error";
+import { API, API_ERROR } from './api';
+import { Cluster } from './cluster';
+import { delay } from './utils/delay';
+import { KafkaTSApiError } from './utils/error';
 
 type MetadataOptions = {
     cluster: Cluster;

package/src/producer/producer.ts

@@ -1,11 +1,11 @@
-import { API, API_ERROR } from "../api";
-import { Cluster } from "../cluster";
-import { distributeMessagesToTopicPartitionLeaders } from "../distributors/messages-to-topic-partition-leaders";
-import { Metadata } from "../metadata";
-import { Message } from "../types";
-import { delay } from "../utils/delay";
-import { KafkaTSApiError } from "../utils/error";
-import { memo } from "../utils/memo";
+import { API, API_ERROR } from '../api';
+import { Cluster } from '../cluster';
+import { distributeMessagesToTopicPartitionLeaders } from '../distributors/messages-to-topic-partition-leaders';
+import { Metadata } from '../metadata';
+import { Message } from '../types';
+import { delay } from '../utils/delay';
+import { KafkaTSApiError } from '../utils/error';
+import { memo } from '../utils/memo';
 
 export type ProducerOptions = {
     allowTopicAutoCreation?: boolean;

package/src/types.ts

@@ -7,3 +7,5 @@ export type Message = {
     value: string | null;
     headers?: Record<string, string>;
 };
+
+export type Batch = Required<Message>[];

package/src/utils/api.ts

@@ -1,11 +1,11 @@
-import { Decoder } from "./decoder";
-import { Encoder } from "./encoder";
+import { Decoder } from './decoder';
+import { Encoder } from './encoder';
 
 export type Api<Request, Response> = {
     apiKey: number;
     apiVersion: number;
     request: (encoder: Encoder, body: Request) => Encoder;
     response: (buffer: Decoder) => Response;
-}
+};
 
-export const createApi = <Request, Response>(api: Api<Request, Response>) => api;
+export const createApi = <Request, Response>(api: Api<Request, Response>) => api;