kafka-ts 0.0.3-beta → 0.0.4

package/src/consumer/consumer.ts

@@ -1,231 +0,0 @@
-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[];
-    groupId?: string | null;
-    groupInstanceId?: string | null;
-    rackId?: string;
-    isolationLevel?: IsolationLevel;
-    sessionTimeoutMs?: number;
-    rebalanceTimeoutMs?: number;
-    maxWaitMs?: number;
-    minBytes?: number;
-    maxBytes?: number;
-    partitionMaxBytes?: number;
-    allowTopicAutoCreation?: boolean;
-    fromBeginning?: boolean;
-    retrier?: Retrier;
-    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(
-        private cluster: Cluster,
-        options: ConsumerOptions,
-    ) {
-        this.options = {
-            ...options,
-            groupId: options.groupId ?? null,
-            groupInstanceId: options.groupInstanceId ?? null,
-            rackId: options.rackId ?? '',
-            sessionTimeoutMs: options.sessionTimeoutMs ?? 30_000,
-            rebalanceTimeoutMs: options.rebalanceTimeoutMs ?? 60_000,
-            maxWaitMs: options.maxWaitMs ?? 5000,
-            minBytes: options.minBytes ?? 1,
-            maxBytes: options.maxBytes ?? 1_000_000,
-            partitionMaxBytes: options.partitionMaxBytes ?? 1_000_000,
-            isolationLevel: options.isolationLevel ?? IsolationLevel.READ_UNCOMMITTED,
-            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 });
-        this.offsetManager = new OffsetManager({
-            cluster: this.cluster,
-            metadata: this.metadata,
-            isolationLevel: this.options.isolationLevel,
-        });
-        this.consumerGroup = this.options.groupId
-            ? new ConsumerGroup({
-                  cluster: this.cluster,
-                  topics: this.options.topics,
-                  groupId: this.options.groupId,
-                  groupInstanceId: this.options.groupInstanceId,
-                  sessionTimeoutMs: this.options.sessionTimeoutMs,
-                  rebalanceTimeoutMs: this.options.rebalanceTimeoutMs,
-                  metadata: this.metadata,
-                  offsetManager: this.offsetManager,
-              })
-            : undefined;
-    }
-
-    public async start(): Promise<void> {
-        const { topics, allowTopicAutoCreation, fromBeginning } = this.options;
-
-        this.stopHook = undefined;
-
-        try {
-            await this.cluster.connect();
-            await this.metadata.fetchMetadataIfNecessary({ topics, allowTopicAutoCreation });
-            this.metadata.setAssignment(this.metadata.getTopicPartitions());
-            await this.offsetManager.fetchOffsets({ fromBeginning });
-            await this.consumerGroup?.join();
-        } catch (error) {
-            console.error(error);
-            console.debug(`Restarting consumer in 1 second...`);
-            await delay(1000);
-
-            if (this.stopHook) return (this.stopHook as () => void)();
-            return this.close(true).then(() => this.start());
-        }
-        setImmediate(() => this.startFetchManager());
-    }
-
-    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}`));
-    }
-
-    private startFetchManager = async () => {
-        const { granularity, concurrency } = this.options;
-
-        while (!this.stopHook) {
-            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 {
-                await this.fetchManager.start();
-
-                if (!nodeAssignments.length) {
-                    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...');
-                    continue;
-                }
-                if ((error as KafkaTSApiError).errorCode === API_ERROR.FENCED_INSTANCE_ID) {
-                    console.debug('New consumer with the same groupInstanceId joined. Exiting the consumer...');
-                    this.close();
-                    break;
-                }
-                if (
-                    error instanceof ConnectionError ||
-                    (error instanceof KafkaTSApiError && error.errorCode === API_ERROR.NOT_COORDINATOR)
-                ) {
-                    console.debug(`${error.message}. Restarting consumer...`);
-                    this.close().then(() => this.start());
-                    break;
-                }
-                console.error(error);
-                this.close();
-                break;
-            }
-        }
-        this.stopHook?.();
-    };
-
-    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?.offsetCommit();
-    }
-
-    private fetch(nodeId: number, assignment: Assignment) {
-        const { rackId, maxWaitMs, minBytes, maxBytes, partitionMaxBytes, isolationLevel } = this.options;
-
-        return this.cluster.sendRequestToNode(nodeId)(API.FETCH, {
-            maxWaitMs,
-            minBytes,
-            maxBytes,
-            isolationLevel,
-            sessionId: 0,
-            sessionEpoch: -1,
-            topics: Object.entries(assignment).map(([topic, partitions]) => ({
-                topicId: this.metadata.getTopicIdByName(topic),
-                partitions: partitions.map((partition) => ({
-                    partition,
-                    currentLeaderEpoch: -1,
-                    fetchOffset: this.offsetManager.getCurrentOffset(topic, partition),
-                    lastFetchedEpoch: -1,
-                    logStartOffset: 0n,
-                    partitionMaxBytes,
-                })),
-            })),
-            forgottenTopicsData: [],
-            rackId,
-        });
-    }
-}

package/src/consumer/fetch-manager.ts

@@ -1,179 +0,0 @@
-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

@@ -1,57 +0,0 @@
-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,93 +0,0 @@
-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;
-    metadata: ConsumerMetadata;
-    isolationLevel: IsolationLevel;
-};
-
-export class OffsetManager {
-    private currentOffsets: Record<string, Record<number, bigint>> = {};
-    public pendingOffsets: Record<string, Record<number, bigint>> = {};
-
-    constructor(private options: OffsetManagerOptions) {}
-
-    public getCurrentOffset(topic: string, partition: number) {
-        return this.currentOffsets[topic]?.[partition] ?? 0n;
-    }
-
-    public resolve(topic: string, partition: number, offset: bigint) {
-        this.pendingOffsets[topic] ??= {};
-        this.pendingOffsets[topic][partition] = offset;
-
-        this.currentOffsets[topic] ??= {};
-        this.currentOffsets[topic][partition] = offset;
-    }
-
-    public flush() {
-        this.pendingOffsets = {};
-    }
-
-    public async fetchOffsets(options: { fromBeginning: boolean }) {
-        const { metadata } = this.options;
-
-        const topicPartitions = Object.entries(metadata.getAssignment()).flatMap(([topic, partitions]) =>
-            partitions.map((partition) => ({ topic, partition })),
-        );
-        const nodeTopicPartitions = distributeMessagesToTopicPartitionLeaders(
-            topicPartitions,
-            metadata.getTopicPartitionLeaderIds(),
-        );
-
-        await Promise.all(
-            Object.entries(nodeTopicPartitions).map(([nodeId, topicPartitions]) =>
-                this.listOffsets({
-                    ...options,
-                    nodeId: parseInt(nodeId),
-                    nodeAssignment: Object.fromEntries(
-                        Object.entries(topicPartitions).map(
-                            ([topicName, partitions]) =>
-                                [topicName, Object.keys(partitions).map(Number)] as [string, number[]],
-                        ),
-                    ),
-                }),
-            ),
-        );
-        this.flush();
-    }
-
-    private async listOffsets({
-        nodeId,
-        nodeAssignment,
-        fromBeginning,
-    }: {
-        nodeId: number;
-        nodeAssignment: Assignment;
-        fromBeginning: boolean;
-    }) {
-        const { cluster, isolationLevel } = this.options;
-
-        const offsets = await cluster.sendRequestToNode(nodeId)(API.LIST_OFFSETS, {
-            replicaId: -1,
-            isolationLevel,
-            topics: Object.entries(nodeAssignment)
-                .flatMap(([topic, partitions]) => partitions.map((partition) => ({ topic, partition })))
-                .map(({ topic, partition }) => ({
-                    name: topic,
-                    partitions: [{ partitionIndex: partition, currentLeaderEpoch: -1, timestamp: -1n }],
-                })),
-        });
-
-        offsets.topics.forEach(({ name, partitions }) => {
-            partitions.forEach(({ partitionIndex, offset }) => {
-                this.resolve(name, partitionIndex, fromBeginning ? 0n : offset);
-            });
-        });
-        this.flush();
-    }
-}

package/src/consumer/processor.ts

@@ -1,47 +0,0 @@
-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,43 +0,0 @@
-import { describe, expect, it } from 'vitest';
-import { distributeAssignmentsToNodesBalanced, distributeAssignmentsToNodesOptimized } from './assignments-to-replicas';
-
-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(`
-              {
-                "1": {
-                  "topic": [
-                    0,
-                  ],
-                },
-                "2": {
-                  "topic": [
-                    1,
-                  ],
-                },
-              }
-            `);
-        });
-    });
-
-    describe('distributeAssignmentsToNodesOptimized', () => {
-        it('smoke', () => {
-            const result = distributeAssignmentsToNodesOptimized(
-                { topic: [0, 1] },
-                { topic: { 0: [0, 1], 1: [1, 2] } },
-            );
-            expect(result).toMatchInlineSnapshot(`
-              {
-                "1": {
-                  "topic": [
-                    0,
-                    1,
-                  ],
-                },
-              }
-            `);
-        });
-    });
-});

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

@@ -1,32 +0,0 @@
-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', () => {
-            const result = distributeMessagesToTopicPartitionLeaders(
-                [{ topic: 'topic', partition: 0, key: null, value: null, offset: 0n, timestamp: 0n, headers: {} }],
-                { topic: { 0: 1 } },
-            );
-            expect(result).toMatchInlineSnapshot(`
-              {
-                "1": {
-                  "topic": {
-                    "0": [
-                      {
-                        "headers": {},
-                        "key": null,
-                        "offset": 0n,
-                        "partition": 0,
-                        "timestamp": 0n,
-                        "topic": "topic",
-                        "value": null,
-                      },
-                    ],
-                  },
-                },
-              }
-            `);
-        });
-    });
-});