kafka-ts 0.0.1-beta → 0.0.3-beta

package/src/cluster.ts

@@ -1,9 +1,9 @@
-import { TcpSocketConnectOpts } from "net";
-import { TLSSocketOptions } from "tls";
-import { API } from "./api";
-import { Broker, SASLOptions } from "./broker";
-import { SendRequest } from "./connection";
-import { ConnectionError, KafkaTSError } from "./utils/error";
+import { TcpSocketConnectOpts } from 'net';
+import { TLSSocketOptions } from 'tls';
+import { API } from './api';
+import { Broker, SASLOptions } from './broker';
+import { SendRequest } from './connection';
+import { ConnectionError, KafkaTSError } from './utils/error';
 
 type ClusterOptions = {
     clientId: string | null;
@@ -82,6 +82,6 @@ export class Cluster {
                 console.warn(`Failed to connect to seed broker ${options.host}:${options.port}`, error);
             }
         }
-        throw new KafkaTSError("No seed brokers found");
+        throw new KafkaTSError('No seed brokers found');
     }
 }

package/src/connection.ts

@@ -1,12 +1,14 @@
-import assert from "assert";
-import net, { isIP, Socket, TcpSocketConnectOpts } from "net";
-import tls, { TLSSocketOptions } from "tls";
-import { getApiName } from "./api";
-import { Api } from "./utils/api";
-import { Decoder } from "./utils/decoder";
-import { Encoder } from "./utils/encoder";
-import { ConnectionError } from "./utils/error";
-import { trace } from "./utils/tracer";
+import assert from 'assert';
+import net, { isIP, Socket, TcpSocketConnectOpts } from 'net';
+import tls, { TLSSocketOptions } from 'tls';
+import { getApiName } from './api';
+import { Api } from './utils/api';
+import { Decoder } from './utils/decoder';
+import { Encoder } from './utils/encoder';
+import { ConnectionError } from './utils/error';
+import { createTracer } from './utils/tracer';
+
+const trace = createTracer('Connection');
 
 export type ConnectionOptions = {
     clientId: string | null;
@@ -44,14 +46,14 @@ export class Connection {
                       resolve,
                   )
                 : net.connect(connection, resolve);
-            this.socket.once("error", reject);
+            this.socket.once('error', reject);
         });
-        this.socket.removeAllListeners("error");
+        this.socket.removeAllListeners('error');
 
-        this.socket.on("data", (data) => this.handleData(data));
-        this.socket.once("close", async () => {
+        this.socket.on('data', (data) => this.handleData(data));
+        this.socket.once('close', async () => {
             Object.values(this.queue).forEach(({ reject }) => {
-                reject(new ConnectionError("Socket closed unexpectedly"));
+                reject(new ConnectionError('Socket closed unexpectedly'));
             });
             this.queue = {};
         });
@@ -100,7 +102,7 @@ export class Connection {
 
     private write(buffer: Buffer) {
         return new Promise<void>((resolve, reject) => {
-            const { stack } = new Error("Write error");
+            const { stack } = new Error('Write error');
             this.socket.write(buffer, (error) => {
                 if (error) {
                     const err = new ConnectionError(error.message);

package/src/consumer/consumer-group.ts

@@ -1,10 +1,10 @@
-import { API, API_ERROR } from "../api";
-import { KEY_TYPE } from "../api/find-coordinator";
-import { Assignment, MemberAssignment } from "../api/sync-group";
-import { Cluster } from "../cluster";
-import { KafkaTSApiError, KafkaTSError } from "../utils/error";
-import { ConsumerMetadata } from "./consumer-metadata";
-import { OffsetManager } from "./offset-manager";
+import { API, API_ERROR } from '../api';
+import { KEY_TYPE } from '../api/find-coordinator';
+import { Assignment, MemberAssignment } from '../api/sync-group';
+import { Cluster } from '../cluster';
+import { KafkaTSApiError, KafkaTSError } from '../utils/error';
+import { ConsumerMetadata } from './consumer-metadata';
+import { OffsetManager } from './offset-manager';
 
 type ConsumerGroupOptions = {
     cluster: Cluster;
@@ -19,9 +19,9 @@ type ConsumerGroupOptions = {
 
 export class ConsumerGroup {
     private coordinatorId = -1;
-    private memberId = "";
+    private memberId = '';
     private generationId = -1;
-    private leaderId = "";
+    private leaderId = '';
     private memberIds: string[] = [];
     private heartbeatInterval: NodeJS.Timeout | null = null;
     private heartbeatError: KafkaTSError | null = null;
@@ -76,8 +76,8 @@ export class ConsumerGroup {
                 memberId: this.memberId,
                 sessionTimeoutMs,
                 rebalanceTimeoutMs,
-                protocolType: "consumer",
-                protocols: [{ name: "RoundRobinAssigner", metadata: { version: 0, topics } }],
+                protocolType: 'consumer',
+                protocols: [{ name: 'RoundRobinAssigner', metadata: { version: 0, topics } }],
                 reason: null,
             });
             this.memberId = response.memberId;
@@ -118,11 +118,11 @@ export class ConsumerGroup {
             groupInstanceId,
             memberId: this.memberId,
             generationId: this.generationId,
-            protocolType: "consumer",
-            protocolName: "RoundRobinAssigner",
+            protocolType: 'consumer',
+            protocolName: 'RoundRobinAssigner',
             assignments,
         });
-        metadata.setAssignment(JSON.parse(response.assignments || "{}") as Assignment);
+        metadata.setAssignment(JSON.parse(response.assignments || '{}') as Assignment);
     }
 
     private async offsetFetch() {

package/src/consumer/consumer-metadata.ts

@@ -1,5 +1,5 @@
-import { Assignment } from "../api/sync-group";
-import { Metadata } from "../metadata";
+import { Assignment } from '../api/sync-group';
+import { Metadata } from '../metadata';
 
 export class ConsumerMetadata extends Metadata {
     private assignment: Assignment = {};

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;