kafka-ts 0.0.1-beta.2 → 0.0.1-beta.4

package/src/api/fetch.ts

@@ -1,3 +1,4 @@
+import { findCodec } from '../codecs';
 import { createApi } from '../utils/api';
 import { Decoder } from '../utils/decoder';
 import { KafkaTSApiError } from '../utils/error';
@@ -68,7 +69,7 @@ export const FETCH = createApi({
             )
             .writeCompactString(data.rackId)
             .writeUVarInt(0),
-    response: (decoder) => {
+    response: async (decoder) => {
         const result = {
             _tag: decoder.readTagBuffer(),
             throttleTimeMs: decoder.readInt32(),
@@ -88,56 +89,105 @@ export const FETCH = createApi({
                         _tag: abortedTransaction.readTagBuffer(),
                     })),
                     preferredReadReplica: partition.readInt32(),
-                    records: decodeRecords(partition),
+                    records: decodeRecordBatch(partition),
                     _tag: partition.readTagBuffer(),
                 })),
                 _tag: response.readTagBuffer(),
             })),
             _tag2: decoder.readTagBuffer(),
         };
+
         if (result.errorCode) throw new KafkaTSApiError(result.errorCode, null, result);
         result.responses.forEach((response) => {
             response.partitions.forEach((partition) => {
                 if (partition.errorCode) throw new KafkaTSApiError(partition.errorCode, null, result);
             });
         });
-        return result;
+
+        const decompressedResponses = await Promise.all(
+            result.responses.map(async (response) => ({
+                ...response,
+                partitions: await Promise.all(
+                    response.partitions.map(async (partition) => ({
+                        ...partition,
+                        records: await Promise.all(
+                            partition.records.map(async ({ recordsLength, compressedRecords, ...record }) => {
+                                const { decompress } = findCodec(record.compression);
+                                const decompressedRecords = await decompress(compressedRecords);
+                                const decompressedDecoder = new Decoder(
+                                    Buffer.concat([recordsLength, decompressedRecords]),
+                                );
+                                return { ...record, records: decodeRecord(decompressedDecoder) };
+                            }),
+                        ),
+                    })),
+                ),
+            })),
+        );
+
+        return { ...result, responses: decompressedResponses };
     },
 });
 
-const decodeRecords = (decoder: Decoder) => {
+const decodeRecordBatch = (decoder: Decoder) => {
     const size = decoder.readUVarInt() - 1;
     if (size <= 0) {
         return [];
     }
 
+    const recordBatchDecoder = new Decoder(decoder.read(size));
+
     const results = [];
-    while (decoder.getBufferLength() > decoder.getOffset() + 49) {
+    while (recordBatchDecoder.getBufferLength() > recordBatchDecoder.getOffset()) {
+        const baseOffset = recordBatchDecoder.readInt64();
+        const batchLength = recordBatchDecoder.readInt32();
+
+        const batchDecoder = new Decoder(recordBatchDecoder.read(batchLength));
+
+        const result = {
+            baseOffset,
+            batchLength,
+            partitionLeaderEpoch: batchDecoder.readInt32(),
+            magic: batchDecoder.readInt8(),
+            crc: batchDecoder.readUInt32(),
+            attributes: batchDecoder.readInt16(),
+            lastOffsetDelta: batchDecoder.readInt32(),
+            baseTimestamp: batchDecoder.readInt64(),
+            maxTimestamp: batchDecoder.readInt64(),
+            producerId: batchDecoder.readInt64(),
+            producerEpoch: batchDecoder.readInt16(),
+            baseSequence: batchDecoder.readInt32(),
+            recordsLength: batchDecoder.read(4),
+            compressedRecords: batchDecoder.read(),
+        };
+
+        const compression = result.attributes & 0x07;
+        const timestampType = (result.attributes & 0x08) >> 3 ? 'LogAppendTime' : 'CreateTime';
+        const isTransactional = !!((result.attributes & 0x10) >> 4);
+        const isControlBatch = !!((result.attributes & 0x20) >> 5);
+        const hasDeleteHorizonMs = !!((result.attributes & 0x40) >> 6);
+
         results.push({
-            baseOffset: decoder.readInt64(),
-            batchLength: decoder.readInt32(),
-            partitionLeaderEpoch: decoder.readInt32(),
-            magic: decoder.readInt8(),
-            crc: decoder.readUInt32(),
-            attributes: decoder.readInt16(),
-            lastOffsetDelta: decoder.readInt32(),
-            baseTimestamp: decoder.readInt64(),
-            maxTimestamp: decoder.readInt64(),
-            producerId: decoder.readInt64(),
-            producerEpoch: decoder.readInt16(),
-            baseSequence: decoder.readInt32(),
-            records: decoder.readRecords((record) => ({
-                attributes: record.readInt8(),
-                timestampDelta: record.readVarLong(),
-                offsetDelta: record.readVarInt(),
-                key: record.readVarIntBuffer(),
-                value: record.readVarIntBuffer(),
-                headers: record.readCompactArray((header) => ({
-                    key: header.readVarIntBuffer(),
-                    value: header.readVarIntBuffer(),
-                })),
-            })),
+            ...result,
+            compression,
+            timestampType,
+            isTransactional,
+            isControlBatch,
+            hasDeleteHorizonMs,
         });
     }
     return results;
 };
+
+const decodeRecord = (decoder: Decoder) =>
+    decoder.readRecords((record) => ({
+        attributes: record.readInt8(),
+        timestampDelta: record.readVarLong(),
+        offsetDelta: record.readVarInt(),
+        key: record.readVarIntBuffer(),
+        value: record.readVarIntBuffer(),
+        headers: record.readVarIntArray((header) => ({
+            key: header.readVarIntBuffer(),
+            value: header.readVarIntBuffer(),
+        })),
+    }));

package/src/cluster.test.ts

@@ -21,7 +21,8 @@ describe.sequential('Request handler', () => {
     let cluster: Cluster;
 
     beforeAll(async () => {
-        cluster = await kafka.createCluster().connect();
+        cluster = await kafka.createCluster();
+        await cluster.connect();
 
         const metadataResult = await cluster.sendRequest(API.METADATA, {
             topics: null,
@@ -52,8 +53,8 @@ describe.sequential('Request handler', () => {
             topics: [
                 {
                     name: 'kafka-ts-test-topic',
-                    numPartitions: 1,
-                    replicationFactor: 1,
+                    numPartitions: 10,
+                    replicationFactor: 3,
                     assignments: [],
                     configs: [],
                 },
@@ -89,6 +90,7 @@ describe.sequential('Request handler', () => {
         expect(result).toMatchSnapshot();
     });
 
+    let partitionIndex = 0;
     let leaderId = 0;
 
     it('should request metadata for a topic', async () => {
@@ -97,6 +99,7 @@ describe.sequential('Request handler', () => {
             allowTopicAutoCreation: false,
             includeTopicAuthorizedOperations: false,
         });
+        partitionIndex = result.topics[0].partitions[0].partitionIndex;
         leaderId = result.topics[0].partitions[0].leaderId;
         result.controllerId = 0;
         result.topics.forEach((topic) => {
@@ -134,7 +137,7 @@ describe.sequential('Request handler', () => {
                     name: 'kafka-ts-test-topic',
                     partitionData: [
                         {
-                            index: 0,
+                            index: partitionIndex,
                             baseOffset: 0n,
                             partitionLeaderEpoch: 0,
                             attributes: 0,
@@ -180,7 +183,7 @@ describe.sequential('Request handler', () => {
                     topicId,
                     partitions: [
                         {
-                            partition: 0,
+                            partition: partitionIndex,
                             currentLeaderEpoch: -1,
                             fetchOffset: 0n,
                             lastFetchedEpoch: 0,

package/src/cluster.ts

@@ -3,7 +3,8 @@ import { TLSSocketOptions } from 'tls';
 import { API } from './api';
 import { Broker, SASLProvider } from './broker';
 import { SendRequest } from './connection';
-import { ConnectionError, KafkaTSError } from './utils/error';
+import { KafkaTSError } from './utils/error';
+import { log } from './utils/logger';
 
 type ClusterOptions = {
     clientId: string | null;
@@ -13,73 +14,77 @@ type ClusterOptions = {
 };
 
 export class Cluster {
-    private seedBroker: Broker;
-    private brokerById: Record<number, Broker> = {};
+    private seedBroker = new Broker({ clientId: null, sasl: null, ssl: null, options: { port: 9092 } });
+    private brokers: { nodeId: number; broker: Broker }[] = [];
+    private brokerMetadata: Record<number, Awaited<ReturnType<(typeof API.METADATA)['response']>>['brokers'][number]> =
+        {};
 
-    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],
-        });
-    }
+    constructor(private options: ClusterOptions) {}
 
     public async connect() {
-        await this.connectSeedBroker();
+        this.seedBroker = await this.findSeedBroker();
+
         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;
+        this.brokerMetadata = Object.fromEntries(metadata.brokers.map((options) => [options.nodeId, options]));
     }
 
     public async disconnect() {
-        await Promise.all([
-            this.seedBroker.disconnect(),
-            ...Object.values(this.brokerById).map((broker) => broker.disconnect()),
-        ]);
+        await Promise.all(this.brokers.map((x) => x.broker.disconnect()));
     }
 
+    public setSeedBroker = async (nodeId: number) => {
+        await this.releaseBroker(this.seedBroker);
+        this.seedBroker = await this.acquireBroker(nodeId);
+    };
+
     public sendRequest: SendRequest = (...args) => this.seedBroker.sendRequest(...args);
 
     public sendRequestToNode =
         (nodeId: number): SendRequest =>
         async (...args) => {
-            const broker = this.brokerById[nodeId];
+            let broker = this.brokers.find((x) => x.nodeId === nodeId)?.broker;
             if (!broker) {
-                throw new ConnectionError(`Broker ${nodeId} is not available`);
+                broker = await this.acquireBroker(nodeId);
             }
-            await broker.ensureConnected();
             return broker.sendRequest(...args);
         };
 
-    private async connectSeedBroker() {
+    public async acquireBroker(nodeId: number) {
+        const broker = new Broker({
+            clientId: this.options.clientId,
+            sasl: this.options.sasl,
+            ssl: this.options.ssl,
+            options: this.brokerMetadata[nodeId],
+        });
+        this.brokers.push({ nodeId, broker });
+        await broker.connect();
+        return broker;
+    }
+
+    public async releaseBroker(broker: Broker) {
+        await broker.disconnect();
+        this.brokers = this.brokers.filter((x) => x.broker !== broker);
+    };
+
+    private async findSeedBroker() {
         const randomizedBrokers = this.options.bootstrapServers.toSorted(() => Math.random() - 0.5);
         for (const options of randomizedBrokers) {
             try {
-                this.seedBroker = await new Broker({
+                const broker = await new Broker({
                     clientId: this.options.clientId,
                     sasl: this.options.sasl,
                     ssl: this.options.ssl,
                     options,
-                }).connect();
-                return;
+                });
+                await broker.connect();
+                this.brokers.push({ nodeId: -1, broker });
+                return broker;
             } catch (error) {
-                console.warn(`Failed to connect to seed broker ${options.host}:${options.port}`, error);
+                log.warn(`Failed to connect to seed broker ${options.host}:${options.port}`, error);
             }
         }
         throw new KafkaTSError('No seed brokers found');

package/src/codecs/gzip.ts

@@ -0,0 +1,9 @@
+import { gzip, unzip } from 'zlib';
+import { Codec } from './types';
+
+export const GZIP: Codec = {
+    compress: async (data) =>
+        new Promise<Buffer>((resolve, reject) => gzip(data, (err, result) => (err ? reject(err) : resolve(result)))),
+    decompress: async (data) =>
+        new Promise<Buffer>((resolve, reject) => unzip(data, (err, result) => (err ? reject(err) : resolve(result)))),
+};

package/src/codecs/index.ts

@@ -0,0 +1,16 @@
+import { GZIP } from './gzip';
+import { NONE } from './none';
+import { Codec } from './types';
+
+const codecs: Record<number, Codec> = {
+    0: NONE,
+    1: GZIP,
+};
+
+export const findCodec = (type: number) => {
+    const codec = codecs[type];
+    if (!codec) {
+        throw new Error(`Unsupported codec: ${type}`);
+    }
+    return codec;
+};

package/src/codecs/none.ts

@@ -0,0 +1,6 @@
+import { Codec } from './types';
+
+export const NONE: Codec = {
+    compress: (data: Buffer) => Promise.resolve(data),
+    decompress: (data: Buffer) => Promise.resolve(data),
+};

package/src/codecs/types.ts

@@ -0,0 +1,4 @@
+export type Codec = {
+    compress: (data: Buffer) => Promise<Buffer>;
+    decompress: (data: Buffer) => Promise<Buffer>;
+};

package/src/connection.ts

@@ -59,6 +59,7 @@ export class Connection {
         });
     }
 
+    @trace()
     public disconnect() {
         this.socket.removeAllListeners();
         return new Promise<void>((resolve) => {
@@ -69,7 +70,7 @@ export class Connection {
         });
     }
 
-    @trace((api, body) => ({ apiName: getApiName(api), body }))
+    @trace((api, body) => ({ message: getApiName(api), body }))
     public async sendRequest<Request, Response>(api: Api<Request, Response>, body: Request): Promise<Response> {
         const correlationId = this.nextCorrelationId();
 
@@ -90,7 +91,7 @@ export class Connection {
                 reject(error);
             }
         });
-        const response = api.response(responseDecoder);
+        const response = await api.response(responseDecoder);
 
         assert(
             responseDecoder.getOffset() - 4 === responseSize,

package/src/consumer/consumer-group.ts

@@ -30,6 +30,8 @@ export class ConsumerGroup {
 
     public async join() {
         await this.findCoordinator();
+        await this.options.cluster.setSeedBroker(this.coordinatorId);
+
         await this.joinGroup();
         await this.syncGroup();
         await this.offsetFetch();
@@ -70,7 +72,7 @@ export class ConsumerGroup {
     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, {
+            const response = await cluster.sendRequest(API.JOIN_GROUP, {
                 groupId,
                 groupInstanceId,
                 memberId: this.memberId,
@@ -113,7 +115,7 @@ export class ConsumerGroup {
             assignments = Object.entries(memberAssignments).map(([memberId, assignment]) => ({ memberId, assignment }));
         }
 
-        const response = await cluster.sendRequestToNode(this.coordinatorId)(API.SYNC_GROUP, {
+        const response = await cluster.sendRequest(API.SYNC_GROUP, {
             groupId,
             groupInstanceId,
             memberId: this.memberId,
@@ -144,7 +146,7 @@ export class ConsumerGroup {
         };
         if (!request.groups.length) return;
 
-        const response = await cluster.sendRequestToNode(this.coordinatorId)(API.OFFSET_FETCH, request);
+        const response = await cluster.sendRequest(API.OFFSET_FETCH, request);
         response.groups.forEach((group) => {
             group.topics.forEach((topic) => {
                 topic.partitions
@@ -177,13 +179,13 @@ export class ConsumerGroup {
         if (!request.topics.length) {
             return;
         }
-        await cluster.sendRequestToNode(this.coordinatorId)(API.OFFSET_COMMIT, request);
+        await cluster.sendRequest(API.OFFSET_COMMIT, request);
         offsetManager.flush();
     }
 
     public async heartbeat() {
         const { cluster, groupId, groupInstanceId } = this.options;
-        await cluster.sendRequestToNode(this.coordinatorId)(API.HEARTBEAT, {
+        await cluster.sendRequest(API.HEARTBEAT, {
             groupId,
             groupInstanceId,
             memberId: this.memberId,
@@ -195,7 +197,7 @@ export class ConsumerGroup {
         const { cluster, groupId, groupInstanceId } = this.options;
         this.stopHeartbeater();
         try {
-            await cluster.sendRequestToNode(this.coordinatorId)(API.LEAVE_GROUP, {
+            await cluster.sendRequest(API.LEAVE_GROUP, {
                 groupId,
                 members: [{ memberId: this.memberId, groupInstanceId, reason: null }],
             });

package/src/consumer/consumer.ts

@@ -2,16 +2,19 @@ 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 { distributeMessagesToTopicPartitionLeaders } from '../distributors/messages-to-topic-partition-leaders';
 import { Message } from '../types';
 import { delay } from '../utils/delay';
 import { ConnectionError, KafkaTSApiError } from '../utils/error';
-import { defaultRetrier, Retrier } from '../utils/retrier';
+import { log } from '../utils/logger';
+import { createTracer } from '../utils/tracer';
 import { ConsumerGroup } from './consumer-group';
 import { ConsumerMetadata } from './consumer-metadata';
-import { FetchManager, BatchGranularity } from './fetch-manager';
+import { BatchGranularity, FetchManager } from './fetch-manager';
 import { OffsetManager } from './offset-manager';
 
+const trace = createTracer('Consumer');
+
 export type ConsumerOptions = {
     topics: string[];
     groupId?: string | null;
@@ -26,7 +29,6 @@ export type ConsumerOptions = {
     partitionMaxBytes?: number;
     allowTopicAutoCreation?: boolean;
     fromBeginning?: boolean;
-    retrier?: Retrier;
     batchGranularity?: BatchGranularity;
     concurrency?: number;
 } & ({ onBatch: (messages: Required<Message>[]) => unknown } | { onMessage: (message: Required<Message>) => unknown });
@@ -57,7 +59,6 @@ export class Consumer {
             isolationLevel: options.isolationLevel ?? IsolationLevel.READ_UNCOMMITTED,
             allowTopicAutoCreation: options.allowTopicAutoCreation ?? false,
             fromBeginning: options.fromBeginning ?? false,
-            retrier: options.retrier ?? defaultRetrier,
             batchGranularity: options.batchGranularity ?? 'partition',
             concurrency: options.concurrency ?? 1,
         };
@@ -94,16 +95,17 @@ export class Consumer {
             await this.offsetManager.fetchOffsets({ fromBeginning });
             await this.consumerGroup?.join();
         } catch (error) {
-            console.error(error);
-            console.debug(`Restarting consumer in 1 second...`);
+            log.error('Failed to start consumer', error);
+            log.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());
+        this.startFetchManager();
     }
 
+    @trace()
     public async close(force = false): Promise<void> {
         if (!force) {
             await new Promise<void>(async (resolve) => {
@@ -111,22 +113,31 @@ export class Consumer {
                 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}`));
+        await this.consumerGroup?.leaveGroup().catch((error) => log.warn(`Failed to leave group: ${error.message}`));
+        await this.cluster.disconnect().catch((error) => log.warn(`Failed to disconnect: ${error.message}`));
     }
 
     private startFetchManager = async () => {
         const { batchGranularity, concurrency } = this.options;
 
         while (!this.stopHook) {
+            // TODO: If leader is not available, find another read replica
             const nodeAssignments = Object.entries(
-                distributeAssignmentsToNodes(
-                    this.metadata.getAssignment(),
-                    this.metadata.getTopicPartitionReplicaIds(),
+                distributeMessagesToTopicPartitionLeaders(
+                    Object.entries(this.metadata.getAssignment()).flatMap(([topic, partitions]) =>
+                        partitions.map((partition) => ({ topic, partition })),
+                    ),
+                    this.metadata.getTopicPartitionLeaderIds(),
+                ),
+            ).map(([nodeId, assignment]) => ({
+                nodeId: parseInt(nodeId),
+                assignment: Object.fromEntries(
+                    Object.entries(assignment).map(([topic, partitions]) => [
+                        topic,
+                        Object.keys(partitions).map(Number),
+                    ]),
                 ),
-            ).map(([nodeId, assignment]) => ({ nodeId: parseInt(nodeId), assignment }));
+            }));
 
             const numPartitions = Object.values(this.metadata.getAssignment()).flat().length;
             const numProcessors = Math.min(concurrency, numPartitions);
@@ -145,7 +156,7 @@ export class Consumer {
                 await this.fetchManager.start();
 
                 if (!nodeAssignments.length) {
-                    console.debug('No partitions assigned. Waiting for reassignment...');
+                    log.debug('No partitions assigned. Waiting for reassignment...');
                     await delay(this.options.maxWaitMs);
                     await this.consumerGroup?.handleLastHeartbeat();
                 }
@@ -153,11 +164,11 @@ export class Consumer {
                 await this.fetchManager.stop();
 
                 if ((error as KafkaTSApiError).errorCode === API_ERROR.REBALANCE_IN_PROGRESS) {
-                    console.debug('Rebalance in progress...');
+                    log.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...');
+                    log.debug('New consumer with the same groupInstanceId joined. Exiting the consumer...');
                     this.close();
                     break;
                 }
@@ -165,11 +176,11 @@ export class Consumer {
                     error instanceof ConnectionError ||
                     (error instanceof KafkaTSApiError && error.errorCode === API_ERROR.NOT_COORDINATOR)
                 ) {
-                    console.debug(`${error.message}. Restarting consumer...`);
+                    log.debug(`${error.message}. Restarting consumer...`);
                     this.close().then(() => this.start());
                     break;
                 }
-                console.error(error);
+                log.error((error as Error).message, error);
                 this.close();
                 break;
             }
@@ -177,12 +188,12 @@ export class Consumer {
         this.stopHook?.();
     };
 
+    @trace()
     private async process(messages: Required<Message>[]) {
         const { options } = this;
-        const { retrier } = options;
 
         if ('onBatch' in options) {
-            await retrier(() => options.onBatch(messages));
+            await options.onBatch(messages);
 
             messages.forEach(({ topic, partition, offset }) =>
                 this.offsetManager.resolve(topic, partition, offset + 1n),
@@ -190,7 +201,7 @@ export class Consumer {
         } else if ('onMessage' in options) {
             try {
                 for (const message of messages) {
-                    await retrier(() => options.onMessage(message));
+                    await options.onMessage(message);
 
                     const { topic, partition, offset } = message;
                     this.offsetManager.resolve(topic, partition, offset + 1n);