kafka-ts 0.0.3-beta → 0.0.3

package/src/consumer/offset-manager.ts

@@ -3,8 +3,11 @@ import { IsolationLevel } from '../api/fetch';
 import { Assignment } from '../api/sync-group';
 import { Cluster } from '../cluster';
 import { distributeMessagesToTopicPartitionLeaders } from '../distributors/messages-to-topic-partition-leaders';
+import { createTracer } from '../utils/tracer';
 import { ConsumerMetadata } from './consumer-metadata';
 
+const trace = createTracer('OffsetManager');
+
 type OffsetManagerOptions = {
     cluster: Cluster;
     metadata: ConsumerMetadata;
@@ -24,13 +27,18 @@ export class OffsetManager {
     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 flush(topicPartitions: Record<string, Set<number>>) {
+        Object.entries(topicPartitions).forEach(([topic, partitions]) => {
+            this.currentOffsets[topic] ??= {};
+            partitions.forEach((partition) => {
+                if (this.pendingOffsets[topic]?.[partition]) {
+                    this.currentOffsets[topic][partition] = this.pendingOffsets[topic][partition];
+                    delete this.pendingOffsets[topic][partition];
+                }
+            });
+        });
     }
 
     public async fetchOffsets(options: { fromBeginning: boolean }) {
@@ -58,7 +66,6 @@ export class OffsetManager {
                 }),
             ),
         );
-        this.flush();
     }
 
     private async listOffsets({
@@ -83,11 +90,15 @@ export class OffsetManager {
                 })),
         });
 
+        const topicPartitions: Record<string, Set<number>> = {};
         offsets.topics.forEach(({ name, partitions }) => {
+            topicPartitions[name] ??= new Set();
             partitions.forEach(({ partitionIndex, offset }) => {
+                topicPartitions[name].add(partitionIndex);
                 this.resolve(name, partitionIndex, fromBeginning ? 0n : offset);
             });
         });
-        this.flush();
+
+        this.flush(topicPartitions);
     }
 }

package/src/consumer/processor.ts

@@ -9,7 +9,7 @@ type ProcessorOptions = {
     process: (batch: Batch) => Promise<void>;
 };
 
-export class Processor extends EventEmitter<{ stop: []; stopped: [] }> {
+export class Processor extends EventEmitter<{ stopped: [] }> {
     private isRunning = false;
 
     constructor(private options: ProcessorOptions) {
@@ -17,15 +17,11 @@ export class Processor extends EventEmitter<{ stop: []; stopped: [] }> {
     }
 
     public async loop() {
-        const { poll, process } = this.options;
-
         this.isRunning = true;
-        this.once('stop', () => (this.isRunning = false));
 
         try {
             while (this.isRunning) {
-                const batch = await poll();
-                await process(batch);
+                await this.step();
             }
         } finally {
             this.isRunning = false;
@@ -34,14 +30,24 @@ export class Processor extends EventEmitter<{ stop: []; stopped: [] }> {
     }
 
     @trace()
+    private async step() {
+        const { poll, process } = this.options;
+
+        const batch = await poll();
+        if (batch.length) {
+            await process(batch);
+        }
+    }
+
     public async stop() {
         if (!this.isRunning) {
             return;
         }
 
-        return new Promise<void>((resolve) => {
+        const stopPromise = new Promise<void>((resolve) => {
             this.once('stopped', resolve);
-            this.emit('stop');
         });
+        this.isRunning = false;
+        return stopPromise;
     }
 }

package/src/distributors/assignments-to-replicas.ts

@@ -1,6 +1,6 @@
 type Assignment = { [topicName: string]: number[] };
 type TopicPartitionReplicaIds = { [topicName: string]: { [partition: number]: number[] } };
-export type NodeAssignment = { [replicaId: number]: Assignment };
+type NodeAssignment = { [replicaId: number]: Assignment };
 
 /** From replica ids pick the one with fewest assignments to balance the load across brokers */
 export const distributeAssignmentsToNodesBalanced = (
@@ -81,5 +81,3 @@ const getPartitionsByReplica = (assignment: Assignment, topicPartitionReplicaIds
     }
     return Object.entries(partitionsByReplicaId);
 };
-
-export const distributeAssignmentsToNodes = distributeAssignmentsToNodesBalanced;

package/src/distributors/partitioner.ts

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

package/src/index.ts

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

package/src/metadata.ts

@@ -2,6 +2,9 @@ import { API, API_ERROR } from './api';
 import { Cluster } from './cluster';
 import { delay } from './utils/delay';
 import { KafkaTSApiError } from './utils/error';
+import { createTracer } from './utils/tracer';
+
+const trace = createTracer('Metadata');
 
 type MetadataOptions = {
     cluster: Cluster;
@@ -36,6 +39,7 @@ export class Metadata {
         return this.topicNameById[id];
     }
 
+    @trace()
     public async fetchMetadataIfNecessary({
         topics,
         allowTopicAutoCreation,

package/src/producer/producer.ts

@@ -1,14 +1,19 @@
 import { API, API_ERROR } from '../api';
 import { Cluster } from '../cluster';
 import { distributeMessagesToTopicPartitionLeaders } from '../distributors/messages-to-topic-partition-leaders';
+import { defaultPartitioner, Partition, Partitioner } from '../distributors/partitioner';
 import { Metadata } from '../metadata';
 import { Message } from '../types';
 import { delay } from '../utils/delay';
 import { KafkaTSApiError } from '../utils/error';
 import { memo } from '../utils/memo';
+import { createTracer } from '../utils/tracer';
+
+const trace = createTracer('Producer');
 
 export type ProducerOptions = {
     allowTopicAutoCreation?: boolean;
+    partitioner?: Partitioner;
 };
 
 export class Producer {
@@ -17,6 +22,7 @@ export class Producer {
     private producerId = 0n;
     private producerEpoch = 0;
     private sequences: Record<string, Record<number, number>> = {};
+    private partition: Partition;
 
     constructor(
         private cluster: Cluster,
@@ -25,11 +31,14 @@ export class Producer {
         this.options = {
             ...options,
             allowTopicAutoCreation: options.allowTopicAutoCreation ?? false,
+            partitioner: options.partitioner ?? defaultPartitioner,
         };
         this.metadata = new Metadata({ cluster });
+        this.partition = this.options.partitioner({ metadata: this.metadata });
     }
 
-    public async send(messages: Message[]) {
+    @trace(() => ({ root: true }))
+    public async send(messages: Message[], { acks = -1 }: { acks?: -1 | 1 } = {}) {
         await this.ensureConnected();
 
         const { allowTopicAutoCreation } = this.options;
@@ -39,19 +48,20 @@ export class Producer {
         await this.metadata.fetchMetadataIfNecessary({ topics, allowTopicAutoCreation });
 
         const nodeTopicPartitionMessages = distributeMessagesToTopicPartitionLeaders(
-            messages,
+            messages.map((message) => ({ ...message, partition: this.partition(message) })),
             this.metadata.getTopicPartitionLeaderIds(),
         );
 
         await Promise.all(
-            Object.entries(nodeTopicPartitionMessages).map(async ([nodeId, topicPartitionMessages]) => {
-                await this.cluster.sendRequestToNode(parseInt(nodeId))(API.PRODUCE, {
+            Object.entries(nodeTopicPartitionMessages).map(([nodeId, topicPartitionMessages]) =>
+                this.cluster.sendRequestToNode(parseInt(nodeId))(API.PRODUCE, {
                     transactionalId: null,
-                    acks: 1,
+                    acks,
                     timeoutMs: 5000,
                     topicData: Object.entries(topicPartitionMessages).map(([topic, partitionMessages]) => ({
                         name: topic,
                         partitionData: Object.entries(partitionMessages).map(([partition, messages]) => {
+                            const partitionIndex = parseInt(partition);
                             let baseTimestamp: bigint | undefined;
                             let maxTimestamp: bigint | undefined;
 
@@ -64,9 +74,9 @@ export class Producer {
                                 }
                             });
 
-                            const baseSequence = this.nextSequence(topic, parseInt(partition), messages.length);
+                            const baseSequence = this.nextSequence(topic, partitionIndex, messages.length);
                             return {
-                                index: parseInt(partition),
+                                index: partitionIndex,
                                 baseOffset: 0n,
                                 partitionLeaderEpoch: -1,
                                 attributes: 0,
@@ -80,18 +90,18 @@ export class Producer {
                                     attributes: 0,
                                     timestampDelta: (message.timestamp ?? defaultTimestamp) - (baseTimestamp ?? 0n),
                                     offsetDelta: index,
-                                    key: message.key,
+                                    key: message.key ?? null,
                                     value: message.value,
                                     headers: Object.entries(message.headers ?? {}).map(([key, value]) => ({
-                                        key,
-                                        value,
+                                        key: Buffer.from(key),
+                                        value: Buffer.from(value),
                                     })),
                                 })),
                             };
                         }),
                     })),
-                });
-            }),
+                }),
+            ),
         );
     }
 

package/src/types.ts

@@ -1,10 +1,10 @@
 export type Message = {
     topic: string;
-    partition: number;
+    partition?: number;
     offset?: bigint;
     timestamp?: bigint;
-    key: string | null;
-    value: string | null;
+    key?: Buffer | null;
+    value: Buffer | null;
     headers?: Record<string, string>;
 };
 

package/src/utils/api.ts

@@ -5,7 +5,7 @@ export type Api<Request, Response> = {
     apiKey: number;
     apiVersion: number;
     request: (encoder: Encoder, body: Request) => Encoder;
-    response: (buffer: Decoder) => Response;
+    response: (buffer: Decoder) => Promise<Response> | Response;
 };
 
 export const createApi = <Request, Response>(api: Api<Request, Response>) => api;

package/src/utils/crypto.ts

@@ -0,0 +1,15 @@
+import { createHash, createHmac, pbkdf2, randomBytes } from 'crypto';
+
+export const generateNonce = () => randomBytes(16).toString('base64').replace(/[\/=]/g, '');
+
+export const saltPassword = (password: string, salt: string, iterations: number, keyLength: number, digest: string) =>
+    new Promise<Buffer>((resolve, reject) =>
+        pbkdf2(password, salt, iterations, keyLength, digest, (err, key) => (err ? reject(err) : resolve(key))),
+    );
+
+export const base64Encode = (input: Buffer | string) => Buffer.from(input).toString('base64');
+export const base64Decode = (input: string) => Buffer.from(input, 'base64').toString();
+export const hash = (data: Buffer, digest: string) => createHash(digest).update(data).digest();
+export const hmac = (key: Buffer, data: Buffer | string, digest: string) =>
+    createHmac(digest, key).update(data).digest();
+export const xor = (a: Buffer, b: Buffer) => Buffer.from(a.map((byte, i) => byte ^ b[i]));

package/src/utils/decoder.ts

@@ -97,13 +97,13 @@ export class Decoder {
         return value;
     }
 
-    public readVarIntString() {
+    public readVarIntBuffer() {
         const length = this.readVarInt();
         if (length < 0) {
             return null;
         }
 
-        const value = this.buffer.toString('utf-8', this.offset, this.offset + length);
+        const value = this.buffer.subarray(this.offset, this.offset + length);
         this.offset += length;
         return value;
     }
@@ -132,6 +132,12 @@ export class Decoder {
         return results;
     }
 
+    public readVarIntArray<T>(callback: (opts: Decoder) => T): T[] {
+        const length = this.readVarInt();
+        const results = Array.from({ length }).map(() => callback(this));
+        return results;
+    }
+
     public readRecords<T>(callback: (opts: Decoder) => T): T[] {
         const length = this.readInt32();
 
@@ -143,9 +149,9 @@ export class Decoder {
         });
     }
 
-    public read(length: number) {
-        const value = this.buffer.subarray(this.offset, this.offset + length);
-        this.offset += length;
+    public read(length?: number) {
+        const value = this.buffer.subarray(this.offset, length !== undefined ? this.offset + length : undefined);
+        this.offset += Buffer.byteLength(value);
         return value;
     }
 

package/src/utils/encoder.ts

@@ -1,41 +1,49 @@
 export class Encoder {
-    private buffer: Buffer;
+    private chunks: Buffer[] = [];
 
-    constructor({ buffer = Buffer.alloc(0) }: { buffer?: Buffer } = {}) {
-        this.buffer = buffer;
+    public getChunks() {
+        return this.chunks;
     }
 
-    public write(rightBuffer: Buffer) {
-        this.buffer = Buffer.concat([this.buffer, rightBuffer]);
+    public getByteLength() {
+        return this.chunks.reduce((acc, chunk) => acc + chunk.byteLength, 0);
+    }
+
+    public write(...buffers: Buffer[]) {
+        this.chunks.push(...buffers);
         return this;
     }
 
+    public writeEncoder(encoder: Encoder) {
+        return this.write(...encoder.getChunks());
+    }
+
     public writeInt8(value: number) {
-        const buffer = Buffer.alloc(1);
+        const buffer = Buffer.allocUnsafe(1);
         buffer.writeInt8(value);
         return this.write(buffer);
     }
 
     public writeInt16(value: number) {
-        const buffer = Buffer.alloc(2);
+        const buffer = Buffer.allocUnsafe(2);
         buffer.writeInt16BE(value);
         return this.write(buffer);
     }
 
     public writeInt32(value: number) {
-        const buffer = Buffer.alloc(4);
+        const buffer = Buffer.allocUnsafe(4);
         buffer.writeInt32BE(value);
         return this.write(buffer);
     }
 
     public writeUInt32(value: number) {
-        const buffer = Buffer.alloc(4);
+        const buffer = Buffer.allocUnsafe(4);
         buffer.writeUInt32BE(value);
         return this.write(buffer);
     }
 
     public writeInt64(value: bigint) {
-        const buffer = Buffer.alloc(8);
+        const buffer = Buffer.allocUnsafe(8);
         buffer.writeBigInt64BE(value);
         return this.write(buffer);
     }
@@ -75,7 +83,7 @@ export class Encoder {
             return this.writeInt16(-1);
         }
         const byteLength = Buffer.byteLength(value, 'utf-8');
-        const buffer = Buffer.alloc(byteLength);
+        const buffer = Buffer.allocUnsafe(byteLength);
         buffer.write(value, 0, byteLength, 'utf-8');
         return this.writeInt16(byteLength).write(buffer);
     }
@@ -86,16 +94,16 @@ export class Encoder {
         }
 
         const byteLength = Buffer.byteLength(value, 'utf-8');
-        const buffer = Buffer.alloc(byteLength);
+        const buffer = Buffer.allocUnsafe(byteLength);
         buffer.write(value, 0, byteLength, 'utf-8');
         return this.writeUVarInt(byteLength + 1).write(buffer);
     }
 
-    public writeVarIntString(value: string | null) {
-        if (value === null) {
+    public writeVarIntBuffer(buffer: Buffer | null) {
+        if (buffer === null) {
             return this.writeVarInt(-1);
         }
-        return this.writeVarInt(Buffer.byteLength(value, 'utf-8')).write(Buffer.from(value, 'utf-8'));
+        return this.writeVarInt(buffer.byteLength).write(buffer);
     }
 
     public writeUUID(value: string | null) {
@@ -110,21 +118,20 @@ export class Encoder {
     }
 
     public writeArray<T>(arr: T[], callback: (encoder: Encoder, item: T) => Encoder) {
-        const buffers = arr.map((item) => callback(new Encoder(), item).value());
-        return this.writeInt32(arr.length).write(Buffer.concat(buffers));
+        return this.writeInt32(arr.length).write(...arr.flatMap((item) => callback(new Encoder(), item).getChunks()));
     }
 
     public writeCompactArray<T>(arr: T[] | null, callback: (encoder: Encoder, item: T) => Encoder) {
         if (arr === null) {
             return this.writeUVarInt(0);
         }
-        const buffers = arr.map((item) => callback(new Encoder(), item).value());
-        return this.writeUVarInt(buffers.length + 1).write(Buffer.concat(buffers));
+        return this.writeUVarInt(arr.length + 1).write(
+            ...arr.flatMap((item) => callback(new Encoder(), item).getChunks()),
+        );
     }
 
     public writeVarIntArray<T>(arr: T[], callback: (encoder: Encoder, item: T) => Encoder) {
-        const buffers = arr.map((item) => callback(new Encoder(), item).value());
-        return this.writeVarInt(buffers.length).write(Buffer.concat(buffers));
+        return this.writeVarInt(arr.length).write(...arr.flatMap((item) => callback(new Encoder(), item).getChunks()));
     }
 
     public writeBytes(value: Buffer) {
@@ -136,6 +143,6 @@ export class Encoder {
     }
 
     public value() {
-        return this.buffer;
+        return Buffer.concat(this.chunks);
     }
 }

package/src/utils/logger.ts

@@ -0,0 +1,37 @@
+export interface Logger {
+    debug: (message: string, metadata?: unknown) => void;
+    info: (message: string, metadata?: unknown) => void;
+    warn: (message: string, metadata?: unknown) => void;
+    error: (message: string, metadata?: unknown) => void;
+}
+
+export const jsonSerializer = (_: unknown, v: unknown) => {
+    if (v instanceof Error) {
+        return { name: v.name, message: v.message, stack: v.stack, cause: v.cause };
+    }
+    if (typeof v === 'bigint') {
+        return v.toString();
+    }
+    return v;
+};
+
+class JsonLogger implements Logger {
+    debug(message: string, metadata?: unknown) {
+        console.debug(JSON.stringify({ message, metadata, level: 'debug' }, jsonSerializer));
+    }
+    info(message: string, metadata?: unknown) {
+        console.info(JSON.stringify({ message, metadata, level: 'info' }, jsonSerializer));
+    }
+    warn(message: string, metadata?: unknown) {
+        console.warn(JSON.stringify({ message, metadata, level: 'warning' }, jsonSerializer));
+    }
+    error(message: string, metadata?: unknown) {
+        console.error(JSON.stringify({ message, metadata, level: 'error' }, jsonSerializer));
+    }
+}
+
+export let log: Logger = new JsonLogger();
+
+export const setLogger = (newLogger: Logger) => {
+    log = newLogger;
+};

package/src/utils/murmur2.ts

@@ -0,0 +1,44 @@
+/* https://github.com/apache/kafka/blob/0.10.2/clients/src/main/java/org/apache/kafka/common/utils/Utils.java#L364 */
+
+export const murmur2 = (data: Buffer): number => {
+    const length = data.length;
+    const seed = 0x9747b28c;
+
+    const m = 0x5bd1e995;
+    const r = 24;
+
+    let h = seed ^ length;
+    let length4 = Math.floor(length / 4);
+
+    for (let i = 0; i < length4; i++) {
+        const i4 = i * 4;
+        let k =
+            (data[i4 + 0] & 0xff) +
+            ((data[i4 + 1] & 0xff) << 8) +
+            ((data[i4 + 2] & 0xff) << 16) +
+            ((data[i4 + 3] & 0xff) << 24);
+        k *= m;
+        k ^= k >> r;
+        k *= m;
+        h *= m;
+        h ^= k;
+    }
+
+    switch (length % 4) {
+        case 3:
+            h = h ^ ((data[(length & ~3) + 2] & 0xff) << 16);
+        case 2:
+            h = h ^ ((data[(length & ~3) + 1] & 0xff) << 8);
+        case 1:
+            h = h ^ (data[length & ~3] & 0xff);
+            h *= m;
+    }
+
+    h ^= h >> 13;
+    h *= m;
+    h ^= h >> 15;
+
+    return h;
+};
+
+export const toPositive = (input: number) => input & 0x7fffffff;

package/src/utils/tracer.ts

@@ -1,31 +1,49 @@
-import { serializer } from './debug';
+import { log } from './logger';
+
+export interface Tracer {
+    startActiveSpan<T>(module: string, method: string, metadata: Record<string, unknown>, callback: () => T): T;
+}
+
+class DebugTracer implements Tracer {
+    private isEnabled = process.env.DEBUG?.includes('kafka-ts');
+
+    startActiveSpan<T>(module: string, method: string, metadata: Record<string, unknown>, callback: () => T): T {
+        if (!this.isEnabled) {
+            return callback();
+        }
+
+        const startTime = Date.now();
+
+        const onEnd = <T>(result: T): T => {
+            log.debug(`[${module}.${method}] ${metadata?.message ?? ''} +${Date.now() - startTime}ms`, {
+                ...metadata,
+                ...(!!result && { result }),
+            });
+            return result;
+        };
+
+        const result = callback();
+        if (result instanceof Promise) {
+            return result.then(onEnd) as T;
+        }
+        onEnd(result);
+        return result;
+    }
+}
+
+let tracer: Tracer = new DebugTracer();
+
+export const setTracer = <T>(newTracer: Tracer) => {
+    tracer = newTracer;
+};
 
 export const createTracer =
-    (module: string, attributes?: Record<string, unknown>) =>
+    (module: string) =>
     (fn?: (...args: any[]) => Record<string, unknown> | undefined) =>
     (target: any, propertyKey: string, descriptor: PropertyDescriptor) => {
-        if (!process.env.DEBUG?.includes('kafka-ts')) return;
-
         const original = descriptor.value;
         descriptor.value = function (...args: any[]) {
-            const startTime = Date.now();
             const metadata = fn?.(...args);
-
-            const onEnd = <T>(result: T): T => {
-                console.log(
-                    `[${module}.${propertyKey}] +${Date.now() - startTime}ms ${JSON.stringify({ ...attributes, ...metadata, result }, serializer)}`,
-                );
-                return result;
-            };
-
-            const result = original.apply(this, args);
-            if (result instanceof Promise) {
-                return result.then(onEnd);
-            } else {
-                onEnd(result);
-                return result;
-            }
+            return tracer.startActiveSpan(module, propertyKey, { ...metadata }, () => original.apply(this, args));
         };
     };
-
-export const trace = createTracer('GLOBAL');

package/.github/workflows/release.yml

@@ -1,17 +0,0 @@
-name: Publish package
-on:
-    release:
-        types: [published]
-jobs:
-    build:
-        runs-on: ubuntu-latest
-        steps:
-            - uses: actions/checkout@v4
-            - uses: actions/setup-node@v4
-              with:
-                  node-version: '20.x'
-                  registry-url: 'https://registry.npmjs.org'
-            - run: npm ci
-            - run: npm publish
-              env:
-                  NODE_AUTH_TOKEN: ${{secrets.NPM_TOKEN}}