kafka-ts 0.0.1-beta.3 → 0.0.1-beta.6

package/src/cluster.test.ts

@@ -8,20 +8,21 @@ import { createKafkaClient } from './client';
 import { Cluster } from './cluster';
 import { KafkaTSApiError } from './utils/error';
 
-export const kafka = createKafkaClient({
+const kafka = createKafkaClient({
     clientId: 'kafka-ts',
     bootstrapServers: [{ host: 'localhost', port: 9092 }],
     sasl: saslPlain({ username: 'admin', password: 'admin' }),
     ssl: { ca: readFileSync('./certs/ca.crt').toString() },
 });
 
-describe.sequential('Request handler', () => {
+describe.sequential('Low-level API', () => {
     const groupId = randomBytes(16).toString('hex');
 
     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

@@ -1,9 +1,11 @@
 import { TcpSocketConnectOpts } from 'net';
 import { TLSSocketOptions } from 'tls';
 import { API } from './api';
+import { Metadata } from './api/metadata';
 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 +15,69 @@ type ClusterOptions = {
 };
 
 export class Cluster {
-    private seedBroker: Broker;
+    private seedBroker = new Broker({ clientId: null, sasl: null, ssl: null, options: { port: 9092 } });
     private brokerById: Record<number, Broker> = {};
+    private brokerMetadata: Record<number, Metadata['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();
+        this.brokerById = {};
+
         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.seedBroker.disconnect(), ...Object.values(this.brokerById).map((x) => x.disconnect())]);
     }
 
+    public setSeedBroker = async (nodeId: number) => {
+        await this.seedBroker.disconnect();
+        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];
-            if (!broker) {
-                throw new ConnectionError(`Broker ${nodeId} is not available`);
+            if (!this.brokerById[nodeId]) {
+                this.brokerById[nodeId] = await this.acquireBroker(nodeId);
             }
-            await broker.ensureConnected();
-            return broker.sendRequest(...args);
+            return this.brokerById[nodeId].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],
+        });
+        await broker.connect();
+        return 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();
+                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

@@ -6,11 +6,12 @@ import { Api } from './utils/api';
 import { Decoder } from './utils/decoder';
 import { Encoder } from './utils/encoder';
 import { ConnectionError } from './utils/error';
+import { log } from './utils/logger';
 import { createTracer } from './utils/tracer';
 
 const trace = createTracer('Connection');
 
-export type ConnectionOptions = {
+type ConnectionOptions = {
     clientId: string | null;
     connection: TcpSocketConnectOpts;
     ssl: TLSSocketOptions | null;
@@ -24,14 +25,14 @@ export class Connection {
         [correlationId: number]: { resolve: (response: RawResonse) => void; reject: (error: Error) => void };
     } = {};
     private lastCorrelationId = 0;
-    private buffer: Buffer | null = null;
+    private chunks: Buffer[] = [];
 
     constructor(private options: ConnectionOptions) {}
 
     @trace()
     public async connect() {
         this.queue = {};
-        this.buffer = null;
+        this.chunks = [];
 
         await new Promise<void>((resolve, reject) => {
             const { ssl, connection } = this.options;
@@ -59,6 +60,7 @@ export class Connection {
         });
     }
 
+    @trace()
     public disconnect() {
         this.socket.removeAllListeners();
         return new Promise<void>((resolve) => {
@@ -69,9 +71,10 @@ 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();
+        const apiName = getApiName(api);
 
         const encoder = new Encoder()
             .writeInt16(api.apiKey)
@@ -79,18 +82,25 @@ export class Connection {
             .writeInt32(correlationId)
             .writeString(this.options.clientId);
 
-        const request = api.request(encoder, body).value();
-        const requestEncoder = new Encoder().writeInt32(request.length).write(request);
+        const request = api.request(encoder, body);
+        const requestEncoder = new Encoder().writeInt32(request.getByteLength()).writeEncoder(request);
 
+        let timeout: NodeJS.Timeout | undefined;
         const { responseDecoder, responseSize } = await new Promise<RawResonse>(async (resolve, reject) => {
+            timeout = setTimeout(() => {
+                delete this.queue[correlationId];
+                reject(new ConnectionError(`${apiName} timed out`));
+            }, 30_000);
+
             try {
-                await this.write(requestEncoder.value());
                 this.queue[correlationId] = { resolve, reject };
+                await this.write(requestEncoder.value());
             } catch (error) {
                 reject(error);
             }
         });
-        const response = api.response(responseDecoder);
+        clearTimeout(timeout);
+        const response = await api.response(responseDecoder);
 
         assert(
             responseDecoder.getOffset() - 4 === responseSize,
@@ -115,12 +125,13 @@ export class Connection {
     }
 
     private handleData(buffer: Buffer) {
-        this.buffer = this.buffer ? Buffer.concat([this.buffer, buffer]) : buffer;
-        if (this.buffer.length < 4) {
+        this.chunks.push(buffer);
+
+        const decoder = new Decoder(Buffer.concat(this.chunks));
+        if (decoder.getBufferLength() < 4) {
             return;
         }
 
-        const decoder = new Decoder(this.buffer);
         const size = decoder.readInt32();
         if (size !== decoder.getBufferLength() - 4) {
             return;
@@ -128,15 +139,18 @@ export class Connection {
 
         const correlationId = decoder.readInt32();
 
-        const { resolve } = this.queue[correlationId];
-        delete this.queue[correlationId];
-
-        resolve({ responseDecoder: decoder, responseSize: size });
-        this.buffer = null;
+        const context = this.queue[correlationId];
+        if (context) {
+            delete this.queue[correlationId];
+            context.resolve({ responseDecoder: decoder, responseSize: size });
+        } else {
+            log.debug('Could not find pending request for correlationId', { correlationId });
+        }
+        this.chunks = [];
     }
 
     private nextCorrelationId() {
-        return (this.lastCorrelationId = (this.lastCorrelationId + 1) % 2 ** 31);
+        return this.lastCorrelationId++;
     }
 }
 

package/src/consumer/consumer-group.ts

@@ -1,11 +1,15 @@
+import EventEmitter from 'events';
 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 { createTracer } from '../utils/tracer';
 import { ConsumerMetadata } from './consumer-metadata';
 import { OffsetManager } from './offset-manager';
 
+const trace = createTracer('ConsumerGroup');
+
 type ConsumerGroupOptions = {
     cluster: Cluster;
     topics: string[];
@@ -17,7 +21,7 @@ type ConsumerGroupOptions = {
     offsetManager: OffsetManager;
 };
 
-export class ConsumerGroup {
+export class ConsumerGroup extends EventEmitter<{ offsetCommit: [] }> {
     private coordinatorId = -1;
     private memberId = '';
     private generationId = -1;
@@ -26,10 +30,16 @@ export class ConsumerGroup {
     private heartbeatInterval: NodeJS.Timeout | null = null;
     private heartbeatError: KafkaTSError | null = null;
 
-    constructor(private options: ConsumerGroupOptions) {}
+    constructor(private options: ConsumerGroupOptions) {
+        super();
+    }
 
+    @trace()
     public async join() {
         await this.findCoordinator();
+        await this.options.cluster.setSeedBroker(this.coordinatorId);
+
+        this.memberId = '';
         await this.joinGroup();
         await this.syncGroup();
         await this.offsetFetch();
@@ -53,12 +63,16 @@ export class ConsumerGroup {
         }
     }
 
-    public async handleLastHeartbeat() {
+    public handleLastHeartbeat() {
         if (this.heartbeatError) {
             throw this.heartbeatError;
         }
     }
 
+    public resetHeartbeat() {
+        this.heartbeatError = null;
+    }
+
     private async findCoordinator() {
         const { coordinators } = await this.options.cluster.sendRequest(API.FIND_COORDINATOR, {
             keyType: KEY_TYPE.GROUP,
@@ -70,7 +84,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 +127,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,31 +158,35 @@ 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);
+
+        const topicPartitions: Record<string, Set<number>> = {};
         response.groups.forEach((group) => {
             group.topics.forEach((topic) => {
-                topic.partitions
-                    .filter(({ committedOffset }) => committedOffset >= 0)
-                    .forEach(({ partitionIndex, committedOffset }) =>
-                        offsetManager.resolve(topic.name, partitionIndex, committedOffset),
-                    );
+                topicPartitions[topic.name] ??= new Set();
+                topic.partitions.forEach(({ partitionIndex, committedOffset }) => {
+                    if (committedOffset >= 0) {
+                        topicPartitions[topic.name].add(partitionIndex);
+                        offsetManager.resolve(topic.name, partitionIndex, committedOffset);
+                    }
+                });
             });
         });
-        offsetManager.flush();
+        offsetManager.flush(topicPartitions);
     }
 
-    public async offsetCommit() {
+    public async offsetCommit(topicPartitions: Record<string, Set<number>>) {
         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]) => ({
+            topics: Object.entries(topicPartitions).map(([topic, partitions]) => ({
                 name: topic,
-                partitions: Object.entries(partitions).map(([partition, offset]) => ({
-                    partitionIndex: parseInt(partition),
-                    committedOffset: offset,
+                partitions: [...partitions].map((partitionIndex) => ({
+                    partitionIndex,
+                    committedOffset: offsetManager.pendingOffsets[topic][partitionIndex],
                     committedLeaderEpoch: -1,
                     committedMetadata: null,
                 })),
@@ -177,13 +195,13 @@ export class ConsumerGroup {
         if (!request.topics.length) {
             return;
         }
-        await cluster.sendRequestToNode(this.coordinatorId)(API.OFFSET_COMMIT, request);
-        offsetManager.flush();
+        await cluster.sendRequest(API.OFFSET_COMMIT, request);
+        this.emit('offsetCommit');
     }
 
     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,
@@ -192,10 +210,14 @@ export class ConsumerGroup {
     }
 
     public async leaveGroup() {
+        if (this.coordinatorId === -1) {
+            return;
+        }
+
         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 }],
             });