kafka-ts 0.0.1-beta

package/src/consumer/offset-manager.ts

@@ -0,0 +1,93 @@
+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/distributors/assignments-to-replicas.test.ts

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

@@ -0,0 +1,85 @@
+type Assignment = { [topicName: string]: number[] };
+type TopicPartitionReplicaIds = { [topicName: string]: { [partition: number]: number[] } };
+export 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);
+};
+
+export const distributeAssignmentsToNodes = distributeAssignmentsToNodesBalanced;

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

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

@@ -0,0 +1,19 @@
+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/index.ts

@@ -0,0 +1,3 @@
+export * from "./client";
+export * from "./api";
+export * from "./types";

package/src/metadata.ts

@@ -0,0 +1,122 @@
+import { API, API_ERROR } from "./api";
+import { Cluster } from "./cluster";
+import { delay } from "./utils/delay";
+import { KafkaTSApiError } from "./utils/error";
+
+type MetadataOptions = {
+    cluster: Cluster;
+};
+
+export class Metadata {
+    private topicPartitions: Record<string, number[]> = {};
+    private topicNameById: Record<string, string> = {};
+    private topicIdByName: Record<string, string> = {};
+    private leaderIdByTopicPartition: Record<string, Record<number, number>> = {};
+    private isrNodesByTopicPartition: Record<string, Record<number, number[]>> = {};
+
+    constructor(private options: MetadataOptions) {}
+
+    public getTopicPartitionLeaderIds() {
+        return this.leaderIdByTopicPartition;
+    }
+
+    public getTopicPartitionReplicaIds() {
+        return this.isrNodesByTopicPartition;
+    }
+
+    public getTopicPartitions() {
+        return this.topicPartitions;
+    }
+
+    public getTopicIdByName(name: string) {
+        return this.topicIdByName[name];
+    }
+
+    public getTopicNameById(id: string) {
+        return this.topicNameById[id];
+    }
+
+    public async fetchMetadataIfNecessary({
+        topics,
+        allowTopicAutoCreation,
+    }: {
+        topics: string[];
+        allowTopicAutoCreation: boolean;
+    }) {
+        const missingTopics = topics.filter((topic) => !this.topicPartitions[topic]);
+        if (!missingTopics.length) {
+            return;
+        }
+
+        try {
+            return await this.fetchMetadata({ topics: missingTopics, allowTopicAutoCreation });
+        } catch (error) {
+            if (
+                error instanceof KafkaTSApiError &&
+                error.errorCode === API_ERROR.UNKNOWN_TOPIC_OR_PARTITION &&
+                allowTopicAutoCreation
+            ) {
+                // TODO: investigate if we can avoid the delay
+                await delay(1000);
+                return await this.fetchMetadata({ topics: missingTopics, allowTopicAutoCreation });
+            }
+            throw error;
+        }
+    }
+
+    private async fetchMetadata({
+        topics,
+        allowTopicAutoCreation,
+    }: {
+        topics: string[] | null;
+        allowTopicAutoCreation: boolean;
+    }) {
+        const { cluster } = this.options;
+
+        const response = await cluster.sendRequest(API.METADATA, {
+            allowTopicAutoCreation,
+            includeTopicAuthorizedOperations: false,
+            topics: topics?.map((name) => ({ id: null, name })) ?? null,
+        });
+
+        this.topicPartitions = {
+            ...this.topicPartitions,
+            ...Object.fromEntries(
+                response.topics.map((topic) => [
+                    topic.name,
+                    topic.partitions.map((partition) => partition.partitionIndex),
+                ]),
+            ),
+        };
+        this.topicNameById = {
+            ...this.topicNameById,
+            ...Object.fromEntries(response.topics.map((topic) => [topic.topicId, topic.name])),
+        };
+        this.topicIdByName = {
+            ...this.topicIdByName,
+            ...Object.fromEntries(response.topics.map((topic) => [topic.name, topic.topicId])),
+        };
+        this.leaderIdByTopicPartition = {
+            ...this.leaderIdByTopicPartition,
+            ...Object.fromEntries(
+                response.topics.map((topic) => [
+                    topic.name,
+                    Object.fromEntries(
+                        topic.partitions.map((partition) => [partition.partitionIndex, partition.leaderId]),
+                    ),
+                ]),
+            ),
+        };
+        this.isrNodesByTopicPartition = {
+            ...this.isrNodesByTopicPartition,
+            ...Object.fromEntries(
+                response.topics.map((topic) => [
+                    topic.name,
+                    Object.fromEntries(
+                        topic.partitions.map((partition) => [partition.partitionIndex, partition.isrNodes]),
+                    ),
+                ]),
+            ),
+        };
+    }
+}

package/src/producer/producer.ts

@@ -0,0 +1,132 @@
+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;
+};
+
+export class Producer {
+    private options: Required<ProducerOptions>;
+    private metadata: Metadata;
+    private producerId = 0n;
+    private producerEpoch = 0;
+    private sequences: Record<string, Record<number, number>> = {};
+
+    constructor(
+        private cluster: Cluster,
+        options: ProducerOptions,
+    ) {
+        this.options = {
+            ...options,
+            allowTopicAutoCreation: options.allowTopicAutoCreation ?? false,
+        };
+        this.metadata = new Metadata({ cluster });
+    }
+
+    public async send(messages: Message[]) {
+        await this.ensureConnected();
+
+        const { allowTopicAutoCreation } = this.options;
+        const defaultTimestamp = BigInt(Date.now());
+
+        const topics = Array.from(new Set(messages.map((message) => message.topic)));
+        await this.metadata.fetchMetadataIfNecessary({ topics, allowTopicAutoCreation });
+
+        const nodeTopicPartitionMessages = distributeMessagesToTopicPartitionLeaders(
+            messages,
+            this.metadata.getTopicPartitionLeaderIds(),
+        );
+
+        await Promise.all(
+            Object.entries(nodeTopicPartitionMessages).map(async ([nodeId, topicPartitionMessages]) => {
+                await this.cluster.sendRequestToNode(parseInt(nodeId))(API.PRODUCE, {
+                    transactionalId: null,
+                    acks: 1,
+                    timeoutMs: 5000,
+                    topicData: Object.entries(topicPartitionMessages).map(([topic, partitionMessages]) => ({
+                        name: topic,
+                        partitionData: Object.entries(partitionMessages).map(([partition, messages]) => {
+                            let baseTimestamp: bigint | undefined;
+                            let maxTimestamp: bigint | undefined;
+
+                            messages.forEach(({ timestamp = defaultTimestamp }) => {
+                                if (!baseTimestamp || timestamp < baseTimestamp) {
+                                    baseTimestamp = timestamp;
+                                }
+                                if (!maxTimestamp || timestamp > maxTimestamp) {
+                                    maxTimestamp = timestamp;
+                                }
+                            });
+
+                            const baseSequence = this.nextSequence(topic, parseInt(partition), messages.length);
+                            return {
+                                index: parseInt(partition),
+                                baseOffset: 0n,
+                                partitionLeaderEpoch: -1,
+                                attributes: 0,
+                                lastOffsetDelta: messages.length - 1,
+                                baseTimestamp: baseTimestamp ?? 0n,
+                                maxTimestamp: maxTimestamp ?? 0n,
+                                producerId: this.producerId,
+                                producerEpoch: 0,
+                                baseSequence,
+                                records: messages.map((message, index) => ({
+                                    attributes: 0,
+                                    timestampDelta: (message.timestamp ?? defaultTimestamp) - (baseTimestamp ?? 0n),
+                                    offsetDelta: index,
+                                    key: message.key,
+                                    value: message.value,
+                                    headers: Object.entries(message.headers ?? {}).map(([key, value]) => ({
+                                        key,
+                                        value,
+                                    })),
+                                })),
+                            };
+                        }),
+                    })),
+                });
+            }),
+        );
+    }
+
+    public async close() {
+        await this.cluster.disconnect();
+    }
+
+    private ensureConnected = memo(async () => {
+        await this.cluster.connect();
+        await this.initProducerId();
+    });
+
+    private async initProducerId(): Promise<void> {
+        try {
+            const result = await this.cluster.sendRequest(API.INIT_PRODUCER_ID, {
+                transactionalId: null,
+                transactionTimeoutMs: 0,
+                producerId: this.producerId,
+                producerEpoch: this.producerEpoch,
+            });
+            this.producerId = result.producerId;
+            this.producerEpoch = result.producerEpoch;
+            this.sequences = {};
+        } catch (error) {
+            if ((error as KafkaTSApiError).errorCode === API_ERROR.COORDINATOR_LOAD_IN_PROGRESS) {
+                await delay(100);
+                return this.initProducerId();
+            }
+            throw error;
+        }
+    }
+
+    private nextSequence(topic: string, partition: number, messagesCount: number) {
+        this.sequences[topic] ??= {};
+        this.sequences[topic][partition] ??= 0;
+        return (this.sequences[topic][partition] += messagesCount || 1);
+    }
+}