kafka-ts 0.0.1-beta

package/src/cluster.ts

@@ -0,0 +1,87 @@
+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;
+    bootstrapServers: TcpSocketConnectOpts[];
+    sasl: SASLOptions | null;
+    ssl: TLSSocketOptions | null;
+};
+
+export class Cluster {
+    private seedBroker: Broker;
+    private brokerById: Record<number, Broker> = {};
+
+    constructor(private options: ClusterOptions) {
+        this.seedBroker = new Broker({
+            clientId: this.options.clientId,
+            sasl: this.options.sasl,
+            ssl: this.options.ssl,
+            options: this.options.bootstrapServers[0],
+        });
+    }
+
+    public async connect() {
+        await this.connectSeedBroker();
+        const metadata = await this.sendRequest(API.METADATA, {
+            allowTopicAutoCreation: false,
+            includeTopicAuthorizedOperations: false,
+            topics: [],
+        });
+
+        this.brokerById = Object.fromEntries(
+            metadata.brokers.map(({ nodeId, ...options }) => [
+                nodeId,
+                new Broker({
+                    clientId: this.options.clientId,
+                    sasl: this.options.sasl,
+                    ssl: this.options.ssl,
+                    options,
+                }),
+            ]),
+        );
+        return this;
+    }
+
+    public async disconnect() {
+        await Promise.all([
+            this.seedBroker.disconnect(),
+            ...Object.values(this.brokerById).map((broker) => broker.disconnect()),
+        ]);
+    }
+
+    public sendRequest: SendRequest = (...args) => this.seedBroker.sendRequest(...args);
+
+    public sendRequestToNode =
+        (nodeId: number): SendRequest =>
+        async (...args) => {
+            const broker = this.brokerById[nodeId];
+            if (!broker) {
+                throw new ConnectionError(`Broker ${nodeId} is not available`);
+            }
+            await broker.ensureConnected();
+            return broker.sendRequest(...args);
+        };
+
+    private async connectSeedBroker() {
+        const randomizedBrokers = this.options.bootstrapServers.toSorted(() => Math.random() - 0.5);
+        for (const options of randomizedBrokers) {
+            try {
+                this.seedBroker = await new Broker({
+                    clientId: this.options.clientId,
+                    sasl: this.options.sasl,
+                    ssl: this.options.ssl,
+                    options,
+                }).connect();
+                return;
+            } catch (error) {
+                console.warn(`Failed to connect to seed broker ${options.host}:${options.port}`, error);
+            }
+        }
+        throw new KafkaTSError("No seed brokers found");
+    }
+}

package/src/connection.ts

@@ -0,0 +1,141 @@
+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";
+
+export type ConnectionOptions = {
+    clientId: string | null;
+    connection: TcpSocketConnectOpts;
+    ssl: TLSSocketOptions | null;
+};
+
+type RawResonse = { responseDecoder: Decoder; responseSize: number };
+
+export class Connection {
+    private socket = new Socket();
+    private queue: {
+        [correlationId: number]: { resolve: (response: RawResonse) => void; reject: (error: Error) => void };
+    } = {};
+    private lastCorrelationId = 0;
+    private buffer: Buffer | null = null;
+
+    constructor(private options: ConnectionOptions) {}
+
+    @trace()
+    public async connect() {
+        this.queue = {};
+        this.buffer = null;
+
+        await new Promise<void>((resolve, reject) => {
+            const { ssl, connection } = this.options;
+
+            this.socket = ssl
+                ? tls.connect(
+                      {
+                          ...connection,
+                          ...ssl,
+                          ...(connection.host && !isIP(connection.host) && { servername: connection.host }),
+                      },
+                      resolve,
+                  )
+                : net.connect(connection, resolve);
+            this.socket.once("error", reject);
+        });
+        this.socket.removeAllListeners("error");
+
+        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"));
+            });
+            this.queue = {};
+        });
+    }
+
+    public disconnect() {
+        this.socket.removeAllListeners();
+        return new Promise<void>((resolve) => {
+            if (this.socket.pending) {
+                return resolve();
+            }
+            this.socket.end(resolve);
+        });
+    }
+
+    @trace((api, body) => ({ apiName: getApiName(api), body }))
+    public async sendRequest<Request, Response>(api: Api<Request, Response>, body: Request): Promise<Response> {
+        const correlationId = this.nextCorrelationId();
+
+        const encoder = new Encoder()
+            .writeInt16(api.apiKey)
+            .writeInt16(api.apiVersion)
+            .writeInt32(correlationId)
+            .writeString(this.options.clientId);
+
+        const request = api.request(encoder, body).value();
+        const requestEncoder = new Encoder().writeInt32(request.length).write(request);
+
+        const { responseDecoder, responseSize } = await new Promise<RawResonse>(async (resolve, reject) => {
+            try {
+                await this.write(requestEncoder.value());
+                this.queue[correlationId] = { resolve, reject };
+            } catch (error) {
+                reject(error);
+            }
+        });
+        const response = api.response(responseDecoder);
+
+        assert(
+            responseDecoder.getOffset() - 4 === responseSize,
+            `Buffer not correctly consumed: ${responseDecoder.getOffset() - 4} !== ${responseSize}`,
+        );
+
+        return response;
+    }
+
+    private write(buffer: Buffer) {
+        return new Promise<void>((resolve, reject) => {
+            const { stack } = new Error("Write error");
+            this.socket.write(buffer, (error) => {
+                if (error) {
+                    const err = new ConnectionError(error.message);
+                    err.stack += `\n${stack}`;
+                    return reject(err);
+                }
+                resolve();
+            });
+        });
+    }
+
+    private handleData(buffer: Buffer) {
+        this.buffer = this.buffer ? Buffer.concat([this.buffer, buffer]) : buffer;
+        if (this.buffer.length < 4) {
+            return;
+        }
+
+        const decoder = new Decoder(this.buffer);
+        const size = decoder.readInt32();
+        if (size !== decoder.getBufferLength() - 4) {
+            return;
+        }
+
+        const correlationId = decoder.readInt32();
+
+        const { resolve } = this.queue[correlationId];
+        delete this.queue[correlationId];
+
+        resolve({ responseDecoder: decoder, responseSize: size });
+        this.buffer = null;
+    }
+
+    private nextCorrelationId() {
+        return (this.lastCorrelationId = (this.lastCorrelationId + 1) % 2 ** 31);
+    }
+}
+
+export type SendRequest = typeof Connection.prototype.sendRequest;

package/src/consumer/consumer-group.ts

@@ -0,0 +1,209 @@
+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;
+    topics: string[];
+    groupId: string;
+    groupInstanceId: string | null;
+    sessionTimeoutMs: number;
+    rebalanceTimeoutMs: number;
+    metadata: ConsumerMetadata;
+    offsetManager: OffsetManager;
+};
+
+export class ConsumerGroup {
+    private coordinatorId = -1;
+    private memberId = "";
+    private generationId = -1;
+    private leaderId = "";
+    private memberIds: string[] = [];
+    private heartbeatInterval: NodeJS.Timeout | null = null;
+    private heartbeatError: KafkaTSError | null = null;
+
+    constructor(private options: ConsumerGroupOptions) {}
+
+    public async join() {
+        await this.findCoordinator();
+        await this.joinGroup();
+        await this.syncGroup();
+        await this.offsetFetch();
+        this.startHeartbeater();
+    }
+
+    private async startHeartbeater() {
+        this.heartbeatInterval = setInterval(async () => {
+            try {
+                await this.heartbeat();
+            } catch (error) {
+                this.heartbeatError = error as KafkaTSError;
+            }
+        }, 5000);
+    }
+
+    private async stopHeartbeater() {
+        if (this.heartbeatInterval) {
+            clearInterval(this.heartbeatInterval);
+            this.heartbeatInterval = null;
+        }
+    }
+
+    public async handleLastHeartbeat() {
+        if (this.heartbeatError) {
+            throw this.heartbeatError;
+        }
+    }
+
+    private async findCoordinator() {
+        const { coordinators } = await this.options.cluster.sendRequest(API.FIND_COORDINATOR, {
+            keyType: KEY_TYPE.GROUP,
+            keys: [this.options.groupId],
+        });
+        this.coordinatorId = coordinators[0].nodeId;
+    }
+
+    private async joinGroup(): Promise<void> {
+        const { cluster, groupId, groupInstanceId, sessionTimeoutMs, rebalanceTimeoutMs, topics } = this.options;
+        try {
+            const response = await cluster.sendRequestToNode(this.coordinatorId)(API.JOIN_GROUP, {
+                groupId,
+                groupInstanceId,
+                memberId: this.memberId,
+                sessionTimeoutMs,
+                rebalanceTimeoutMs,
+                protocolType: "consumer",
+                protocols: [{ name: "RoundRobinAssigner", metadata: { version: 0, topics } }],
+                reason: null,
+            });
+            this.memberId = response.memberId;
+            this.generationId = response.generationId;
+            this.leaderId = response.leader;
+            this.memberIds = response.members.map((member) => member.memberId);
+        } catch (error) {
+            if ((error as KafkaTSApiError).errorCode === API_ERROR.MEMBER_ID_REQUIRED) {
+                this.memberId = (error as KafkaTSApiError).response.memberId;
+                return this.joinGroup();
+            }
+            throw error;
+        }
+    }
+
+    private async syncGroup() {
+        const { cluster, metadata, groupId, groupInstanceId } = this.options;
+
+        let assignments: MemberAssignment[] = [];
+        if (this.memberId === this.leaderId) {
+            const memberAssignments = Object.entries(metadata.getTopicPartitions())
+                .flatMap(([topic, partitions]) => partitions.map((partition) => ({ topic, partition })))
+                .reduce(
+                    (acc, { topic, partition }, index) => {
+                        const memberId = this.memberIds[index % this.memberIds.length];
+                        acc[memberId] ??= {};
+                        acc[memberId][topic] ??= [];
+                        acc[memberId][topic].push(partition);
+                        return acc;
+                    },
+                    {} as Record<string, Record<string, number[]>>,
+                );
+            assignments = Object.entries(memberAssignments).map(([memberId, assignment]) => ({ memberId, assignment }));
+        }
+
+        const response = await cluster.sendRequestToNode(this.coordinatorId)(API.SYNC_GROUP, {
+            groupId,
+            groupInstanceId,
+            memberId: this.memberId,
+            generationId: this.generationId,
+            protocolType: "consumer",
+            protocolName: "RoundRobinAssigner",
+            assignments,
+        });
+        metadata.setAssignment(JSON.parse(response.assignments || "{}") as Assignment);
+    }
+
+    private async offsetFetch() {
+        const { cluster, groupId, topics, metadata, offsetManager } = this.options;
+
+        const assignment = metadata.getAssignment();
+        const request = {
+            groups: [
+                {
+                    groupId,
+                    memberId: this.memberId,
+                    memberEpoch: -1,
+                    topics: topics
+                        .map((topic) => ({ name: topic, partitionIndexes: assignment[topic] ?? [] }))
+                        .filter(({ partitionIndexes }) => partitionIndexes.length),
+                },
+            ].filter(({ topics }) => topics.length),
+            requireStable: true,
+        };
+        if (!request.groups.length) return;
+
+        const response = await cluster.sendRequestToNode(this.coordinatorId)(API.OFFSET_FETCH, request);
+        response.groups.forEach((group) => {
+            group.topics.forEach((topic) => {
+                topic.partitions
+                    .filter(({ committedOffset }) => committedOffset >= 0)
+                    .forEach(({ partitionIndex, committedOffset }) =>
+                        offsetManager.resolve(topic.name, partitionIndex, committedOffset),
+                    );
+            });
+        });
+        offsetManager.flush();
+    }
+
+    public async offsetCommit() {
+        const { cluster, groupId, groupInstanceId, offsetManager } = this.options;
+        const request = {
+            groupId,
+            groupInstanceId,
+            memberId: this.memberId,
+            generationIdOrMemberEpoch: this.generationId,
+            topics: Object.entries(offsetManager.pendingOffsets).map(([topic, partitions]) => ({
+                name: topic,
+                partitions: Object.entries(partitions).map(([partition, offset]) => ({
+                    partitionIndex: parseInt(partition),
+                    committedOffset: offset,
+                    committedLeaderEpoch: -1,
+                    committedMetadata: null,
+                })),
+            })),
+        };
+        if (!request.topics.length) {
+            return;
+        }
+        await cluster.sendRequestToNode(this.coordinatorId)(API.OFFSET_COMMIT, request);
+        offsetManager.flush();
+    }
+
+    public async heartbeat() {
+        const { cluster, groupId, groupInstanceId } = this.options;
+        await cluster.sendRequestToNode(this.coordinatorId)(API.HEARTBEAT, {
+            groupId,
+            groupInstanceId,
+            memberId: this.memberId,
+            generationId: this.generationId,
+        });
+    }
+
+    public async leaveGroup() {
+        const { cluster, groupId, groupInstanceId } = this.options;
+        this.stopHeartbeater();
+        try {
+            await cluster.sendRequestToNode(this.coordinatorId)(API.LEAVE_GROUP, {
+                groupId,
+                members: [{ memberId: this.memberId, groupInstanceId, reason: null }],
+            });
+        } catch (error) {
+            if ((error as KafkaTSApiError).errorCode === API_ERROR.FENCED_INSTANCE_ID) {
+                return;
+            }
+            throw error;
+        }
+    }
+}

package/src/consumer/consumer-metadata.ts

@@ -0,0 +1,14 @@
+import { Assignment } from "../api/sync-group";
+import { Metadata } from "../metadata";
+
+export class ConsumerMetadata extends Metadata {
+    private assignment: Assignment = {};
+
+    public getAssignment() {
+        return this.assignment;
+    }
+
+    public setAssignment(newAssignment: Assignment) {
+        this.assignment = newAssignment;
+    }
+}

package/src/consumer/consumer.ts

@@ -0,0 +1,229 @@
+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";
+
+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;
+} & ({ onMessage: (message: Message) => unknown } | { onBatch: (messages: Message[]) => unknown });
+
+export class Consumer {
+    private options: Required<ConsumerOptions>;
+    private metadata: ConsumerMetadata;
+    private consumerGroup: ConsumerGroup | undefined;
+    private offsetManager: OffsetManager;
+    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,
+        };
+
+        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());
+        }
+        this.fetchLoop();
+    }
+
+    private fetchLoop = async () => {
+        const { options } = this;
+        const { retrier } = options;
+
+        let nodeAssignments: { nodeId: number; assignment: Assignment }[] = [];
+        let shouldReassign = true;
+
+        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;
+            }
+
+            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();
+                }
+
+                if (!nodeAssignments.length) {
+                    console.debug("No partitions assigned. Waiting for reassignment...");
+                    await delay(this.options.maxWaitMs);
+                    await this.consumerGroup?.handleLastHeartbeat();
+                }
+            } catch (error) {
+                if ((error as KafkaTSApiError).errorCode === API_ERROR.REBALANCE_IN_PROGRESS) {
+                    console.debug("Rebalance in progress...");
+                    shouldReassign = true;
+                    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);
+                await this.consumerGroup?.offsetCommit();
+                break;
+            }
+        }
+        this.stopHook?.();
+    };
+
+    public async close(force = false): Promise<void> {
+        if (!force) {
+            await new Promise<void>((resolve) => {
+                this.stopHook = resolve;
+            });
+        }
+        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 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,
+        });
+    }
+}