kafka-ts 0.0.3 → 0.0.5

package/src/consumer/fetch-manager.ts

@@ -1,169 +0,0 @@
-import { FetchResponse } from '../api/fetch';
-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 BatchGranularity = 'partition' | 'topic' | 'broker';
-
-type FetchManagerOptions = {
-    fetch: (nodeId: number, assignment: Assignment) => Promise<FetchResponse>;
-    process: (batch: Batch) => Promise<void>;
-    metadata: Metadata;
-    consumerGroup?: ConsumerGroup;
-    nodeAssignments: { nodeId: number; assignment: Assignment }[];
-    batchGranularity: BatchGranularity;
-    concurrency: number;
-};
-
-type Checkpoint = { kind: 'checkpoint'; fetcherId: number };
-type Entry = Batch | Checkpoint;
-
-export class FetchManager {
-    private queue: Entry[] = [];
-    private isRunning = false;
-    private fetchers: Fetcher[];
-    private processors: Processor[];
-    private pollQueue: (() => void)[] = [];
-    private fetcherCallbacks: Record<number, () => void> = {};
-
-    constructor(private options: FetchManagerOptions) {
-        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) }),
-        );
-    }
-
-    @trace(() => ({ root: true }))
-    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 {
-            await this.stop();
-        }
-    }
-
-    public async stop() {
-        this.isRunning = false;
-
-        const stopPromise = Promise.all([
-            ...this.fetchers.map((fetcher) => fetcher.stop()),
-            ...this.processors.map((processor) => processor.stop()),
-        ]);
-
-        this.pollQueue.forEach((resolve) => resolve());
-        this.pollQueue = [];
-
-        Object.values(this.fetcherCallbacks).forEach((callback) => callback());
-        this.fetcherCallbacks = {};
-
-        await stopPromise;
-    }
-
-    @trace()
-    public async poll(): Promise<Batch> {
-        if (!this.isRunning) {
-            return [];
-        }
-
-        const batch = this.queue.shift();
-        if (!batch) {
-            // wait until new data is available or fetch manager is requested to stop
-            await new Promise<void>((resolve) => {
-                this.pollQueue.push(resolve);
-            });
-            return this.poll();
-        }
-
-        if ('kind' in batch && batch.kind === 'checkpoint') {
-            this.fetcherCallbacks[batch.fetcherId]?.();
-            return this.poll();
-        }
-
-        this.pollQueue?.shift()?.();
-
-        return batch as Exclude<Entry, Checkpoint>;
-    }
-
-    @trace()
-    private async onResponse(fetcherId: number, response: FetchResponse) {
-        const { metadata, batchGranularity } = this.options;
-
-        const batches = fetchResponseToBatches(response, batchGranularity, metadata);
-        if (!batches.length) {
-            return;
-        }
-
-        // wait until all broker batches have been processed or fetch manager is requested to stop
-        await new Promise<void>((resolve) => {
-            this.fetcherCallbacks[fetcherId] = resolve;
-            this.queue.push(...batches, { kind: 'checkpoint', fetcherId });
-            this.pollQueue?.shift()?.();
-        });
-    }
-}
-
-const fetchResponseToBatches = (
-    batch: FetchResponse,
-    batchGranularity: BatchGranularity,
-    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 (batchGranularity) {
-        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))
-                .filter((messages) => messages.length);
-        default:
-            throw new KafkaTSError(`Unhandled batch granularity: ${batchGranularity}`);
-    }
-};

package/src/consumer/fetcher.ts

@@ -1,64 +0,0 @@
-import { EventEmitter } from 'stream';
-import { FetchResponse } from '../api/fetch';
-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<FetchResponse>;
-    onResponse: (fetcherId: number, response: FetchResponse) => Promise<void>;
-};
-
-export class Fetcher extends EventEmitter<{ stopped: [] }> {
-    private isRunning = false;
-
-    constructor(
-        private fetcherId: number,
-        private options: FetcherOptions,
-    ) {
-        super();
-    }
-
-    public async loop() {
-        this.isRunning = true;
-
-        try {
-            while (this.isRunning) {
-                await this.step();
-            }
-        } finally {
-            this.isRunning = false;
-            this.emit('stopped');
-        }
-    }
-
-    @trace()
-    private async step() {
-        const { nodeId, assignment, consumerGroup, fetch, onResponse } = this.options;
-
-        const response = await fetch(nodeId, assignment);
-        if (!this.isRunning) {
-            return;
-        }
-        consumerGroup?.handleLastHeartbeat();
-        await onResponse(this.fetcherId, response);
-        consumerGroup?.handleLastHeartbeat();
-    }
-
-    public async stop() {
-        if (!this.isRunning) {
-            return;
-        }
-
-        const stopPromise = new Promise<void>((resolve) => {
-            this.once('stopped', resolve);
-        });
-        this.isRunning = false;
-        return stopPromise;
-    }
-}

package/src/consumer/offset-manager.ts

@@ -1,104 +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 { createTracer } from '../utils/tracer';
-import { ConsumerMetadata } from './consumer-metadata';
-
-const trace = createTracer('OffsetManager');
-
-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;
-    }
-
-    public flush(topicPartitions: Record<string, Set<number>>) {
-        Object.entries(topicPartitions).forEach(([topic, partitions]) => {
-            this.currentOffsets[topic] ??= {};
-            partitions.forEach((partition) => {
-                if (this.pendingOffsets[topic]?.[partition]) {
-                    this.currentOffsets[topic][partition] = this.pendingOffsets[topic][partition];
-                    delete this.pendingOffsets[topic][partition];
-                }
-            });
-        });
-    }
-
-    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[]],
-                        ),
-                    ),
-                }),
-            ),
-        );
-    }
-
-    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 }],
-                })),
-        });
-
-        const topicPartitions: Record<string, Set<number>> = {};
-        offsets.topics.forEach(({ name, partitions }) => {
-            topicPartitions[name] ??= new Set();
-            partitions.forEach(({ partitionIndex, offset }) => {
-                topicPartitions[name].add(partitionIndex);
-                this.resolve(name, partitionIndex, fromBeginning ? 0n : offset);
-            });
-        });
-
-        this.flush(topicPartitions);
-    }
-}

package/src/consumer/processor.ts

@@ -1,53 +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<{ stopped: [] }> {
-    private isRunning = false;
-
-    constructor(private options: ProcessorOptions) {
-        super();
-    }
-
-    public async loop() {
-        this.isRunning = true;
-
-        try {
-            while (this.isRunning) {
-                await this.step();
-            }
-        } finally {
-            this.isRunning = false;
-            this.emit('stopped');
-        }
-    }
-
-    @trace()
-    private async step() {
-        const { poll, process } = this.options;
-
-        const batch = await poll();
-        if (batch.length) {
-            await process(batch);
-        }
-    }
-
-    public async stop() {
-        if (!this.isRunning) {
-            return;
-        }
-
-        const stopPromise = new Promise<void>((resolve) => {
-            this.once('stopped', resolve);
-        });
-        this.isRunning = false;
-        return stopPromise;
-    }
-}

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/assignments-to-replicas.ts

@@ -1,83 +0,0 @@
-type Assignment = { [topicName: string]: number[] };
-type TopicPartitionReplicaIds = { [topicName: string]: { [partition: number]: number[] } };
-type NodeAssignment = { [replicaId: number]: Assignment };
-
-/** From replica ids pick the one with fewest assignments to balance the load across brokers */
-export const distributeAssignmentsToNodesBalanced = (
-    assignment: Assignment,
-    topicPartitionReplicaIds: TopicPartitionReplicaIds,
-) => {
-    const replicaPartitions = getPartitionsByReplica(assignment, topicPartitionReplicaIds);
-
-    const result: NodeAssignment = {};
-
-    for (const [topicName, partitions] of Object.entries(assignment)) {
-        for (const partition of partitions) {
-            const replicaIds = topicPartitionReplicaIds[topicName][partition];
-            const replicaId = replicaIds.reduce((prev, curr) => {
-                if (!prev) {
-                    return curr;
-                }
-                return (replicaPartitions[prev]?.length ?? 0) < (replicaPartitions[curr]?.length ?? 0) ? prev : curr;
-            });
-            result[replicaId] ??= {};
-            result[replicaId][topicName] ??= [];
-            result[replicaId][topicName].push(partition);
-        }
-    }
-
-    return result;
-};
-
-/** Minimize the total number of replicas in the result to reduce the number of requests to different brokers */
-export const distributeAssignmentsToNodesOptimized = (
-    assignment: Assignment,
-    topicPartitionReplicaIds: TopicPartitionReplicaIds,
-) => {
-    const result: NodeAssignment = {};
-
-    const sortFn = ([, partitionsA]: [string, string[]], [, partitionsB]: [string, string[]]) =>
-        partitionsB.length - partitionsA.length;
-
-    let replicaPartitions = getPartitionsByReplica(assignment, topicPartitionReplicaIds);
-
-    while (replicaPartitions.length) {
-        replicaPartitions.sort(sortFn);
-
-        const [replicaId, partitions] = replicaPartitions.shift()!;
-        if (!partitions.length) {
-            continue;
-        }
-
-        result[parseInt(replicaId)] = partitions.reduce((acc, partition) => {
-            const [topicName, partitionId] = partition.split(':');
-            acc[topicName] ??= [];
-            acc[topicName].push(parseInt(partitionId));
-            return acc;
-        }, {} as Assignment);
-
-        replicaPartitions = replicaPartitions.map(
-            ([replicaId, replicaPartitions]) =>
-                [replicaId, replicaPartitions.filter((partition) => !partitions.includes(partition))] as [
-                    string,
-                    string[],
-                ],
-        );
-    }
-
-    return result;
-};
-
-const getPartitionsByReplica = (assignment: Assignment, topicPartitionReplicaIds: TopicPartitionReplicaIds) => {
-    const partitionsByReplicaId: { [replicaId: number]: string[] } = {};
-    for (const [topicName, partitions] of Object.entries(assignment)) {
-        for (const partition of partitions) {
-            const replicaIds = topicPartitionReplicaIds[topicName][partition];
-            for (const replicaId of replicaIds) {
-                partitionsByReplicaId[replicaId] ??= [];
-                partitionsByReplicaId[replicaId].push(`${topicName}:${partition}`);
-            }
-        }
-    }
-    return Object.entries(partitionsByReplicaId);
-};

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,
-                      },
-                    ],
-                  },
-                },
-              }
-            `);
-        });
-    });
-});

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

@@ -1,19 +0,0 @@
-type TopicPartitionLeader = { [topicName: string]: { [partitionId: number]: number } };
-type MessagesByNodeTopicPartition<T> = {
-    [nodeId: number]: { [topicName: string]: { [partitionId: number]: T[] } };
-};
-
-export const distributeMessagesToTopicPartitionLeaders = <T extends { topic: string; partition: number }>(
-    messages: T[],
-    topicPartitionLeader: TopicPartitionLeader,
-) => {
-    const result: MessagesByNodeTopicPartition<T> = {};
-    messages.forEach((message) => {
-        const leaderId = topicPartitionLeader[message.topic][message.partition];
-        result[leaderId] ??= {};
-        result[leaderId][message.topic] ??= {};
-        result[leaderId][message.topic][message.partition] ??= [];
-        result[leaderId][message.topic][message.partition].push(message);
-    });
-    return result;
-};

package/src/distributors/partitioner.ts

@@ -1,27 +0,0 @@
-import { Metadata } from '../metadata';
-import { Message } from '../types';
-import { murmur2, toPositive } from '../utils/murmur2';
-
-export type Partition = (message: Message) => number;
-export type Partitioner = (context: { metadata: Metadata }) => Partition;
-
-export const defaultPartitioner: Partitioner = ({ metadata }) => {
-    const topicCounterMap: Record<string, number> = {};
-
-    const getNextValue = (topic: string) => {
-        topicCounterMap[topic] ??= 0;
-        return topicCounterMap[topic]++;
-    };
-
-    return ({ topic, partition, key }: Message) => {
-        if (partition !== null && partition !== undefined) {
-            return partition;
-        }
-        const partitions = metadata.getTopicPartitions()[topic];
-        const numPartitions = partitions.length;
-        if (key) {
-            return toPositive(murmur2(key)) % numPartitions;
-        }
-        return toPositive(getNextValue(topic)) % numPartitions;
-    };
-};

package/src/index.ts

@@ -1,9 +0,0 @@
-export * from './api';
-export * from './auth';
-export { SASLProvider } from './broker';
-export * from './client';
-export * from './distributors/partitioner';
-export * from './types';
-export * from './utils/error';
-export * from './utils/logger';
-export { Tracer, setTracer } from './utils/tracer';