kafka-ts 0.0.2-beta → 0.0.2

package/src/consumer/fetch-manager.ts

@@ -0,0 +1,169 @@
+import { FetchResponse } from '../api/fetch';
+import { Assignment } from '../api/sync-group';
+import { Metadata } from '../metadata';
+import { Batch, Message } from '../types';
+import { KafkaTSError } from '../utils/error';
+import { createTracer } from '../utils/tracer';
+import { ConsumerGroup } from './consumer-group';
+import { Fetcher } from './fetcher';
+import { Processor } from './processor';
+
+const trace = createTracer('FetchManager');
+
+export type BatchGranularity = 'partition' | 'topic' | 'broker';
+
+type FetchManagerOptions = {
+    fetch: (nodeId: number, assignment: Assignment) => Promise<FetchResponse>;
+    process: (batch: Batch) => Promise<void>;
+    metadata: Metadata;
+    consumerGroup?: ConsumerGroup;
+    nodeAssignments: { nodeId: number; assignment: Assignment }[];
+    batchGranularity: BatchGranularity;
+    concurrency: number;
+};
+
+type Checkpoint = { kind: 'checkpoint'; fetcherId: number };
+type Entry = Batch | Checkpoint;
+
+export class FetchManager {
+    private queue: Entry[] = [];
+    private isRunning = false;
+    private fetchers: Fetcher[];
+    private processors: Processor[];
+    private pollQueue: (() => void)[] = [];
+    private fetcherCallbacks: Record<number, () => void> = {};
+
+    constructor(private options: FetchManagerOptions) {
+        const { fetch, process, consumerGroup, nodeAssignments, concurrency } = this.options;
+
+        this.fetchers = nodeAssignments.map(
+            ({ nodeId, assignment }, index) =>
+                new Fetcher(index, {
+                    nodeId,
+                    assignment,
+                    consumerGroup,
+                    fetch,
+                    onResponse: this.onResponse.bind(this),
+                }),
+        );
+        this.processors = Array.from({ length: concurrency }).map(
+            () => new Processor({ process, poll: this.poll.bind(this) }),
+        );
+    }
+
+    @trace(() => ({ root: true }))
+    public async start() {
+        this.queue = [];
+        this.isRunning = true;
+
+        try {
+            await Promise.all([
+                ...this.fetchers.map((fetcher) => fetcher.loop()),
+                ...this.processors.map((processor) => processor.loop()),
+            ]);
+        } finally {
+            await this.stop();
+        }
+    }
+
+    public async stop() {
+        this.isRunning = false;
+
+        const stopPromise = Promise.all([
+            ...this.fetchers.map((fetcher) => fetcher.stop()),
+            ...this.processors.map((processor) => processor.stop()),
+        ]);
+
+        this.pollQueue.forEach((resolve) => resolve());
+        this.pollQueue = [];
+
+        Object.values(this.fetcherCallbacks).forEach((callback) => callback());
+        this.fetcherCallbacks = {};
+
+        await stopPromise;
+    }
+
+    @trace()
+    public async poll(): Promise<Batch> {
+        if (!this.isRunning) {
+            return [];
+        }
+
+        const batch = this.queue.shift();
+        if (!batch) {
+            // wait until new data is available or fetch manager is requested to stop
+            await new Promise<void>((resolve) => {
+                this.pollQueue.push(resolve);
+            });
+            return this.poll();
+        }
+
+        if ('kind' in batch && batch.kind === 'checkpoint') {
+            this.fetcherCallbacks[batch.fetcherId]?.();
+            return this.poll();
+        }
+
+        this.pollQueue?.shift()?.();
+
+        return batch as Exclude<Entry, Checkpoint>;
+    }
+
+    @trace()
+    private async onResponse(fetcherId: number, response: FetchResponse) {
+        const { metadata, batchGranularity } = this.options;
+
+        const batches = fetchResponseToBatches(response, batchGranularity, metadata);
+        if (!batches.length) {
+            return;
+        }
+
+        // wait until all broker batches have been processed or fetch manager is requested to stop
+        await new Promise<void>((resolve) => {
+            this.fetcherCallbacks[fetcherId] = resolve;
+            this.queue.push(...batches, { kind: 'checkpoint', fetcherId });
+            this.pollQueue?.shift()?.();
+        });
+    }
+}
+
+const fetchResponseToBatches = (
+    batch: FetchResponse,
+    batchGranularity: BatchGranularity,
+    metadata: Metadata,
+): Batch[] => {
+    const brokerTopics = batch.responses.map(({ topicId, partitions }) =>
+        partitions.map(({ partitionIndex, records }) =>
+            records.flatMap(({ baseTimestamp, baseOffset, records }) =>
+                records.map(
+                    (message): Required<Message> => ({
+                        topic: 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),
+                    }),
+                ),
+            ),
+        ),
+    );
+
+    switch (batchGranularity) {
+        case 'broker':
+            const messages = brokerTopics.flatMap((topicPartition) =>
+                topicPartition.flatMap((partitionMessages) => partitionMessages),
+            );
+            return messages.length ? [messages] : [];
+        case 'topic':
+            return brokerTopics
+                .map((topicPartition) => topicPartition.flatMap((partitionMessages) => partitionMessages))
+                .filter((messages) => messages.length);
+        case 'partition':
+            return brokerTopics
+                .flatMap((topicPartition) => topicPartition.map((partitionMessages) => partitionMessages))
+                .filter((messages) => messages.length);
+        default:
+            throw new KafkaTSError(`Unhandled batch granularity: ${batchGranularity}`);
+    }
+};

package/src/consumer/fetcher.ts

@@ -0,0 +1,64 @@
+import { EventEmitter } from 'stream';
+import { FetchResponse } from '../api/fetch';
+import { Assignment } from '../api/sync-group';
+import { createTracer } from '../utils/tracer';
+import { ConsumerGroup } from './consumer-group';
+
+const trace = createTracer('Fetcher');
+
+type FetcherOptions = {
+    nodeId: number;
+    assignment: Assignment;
+    consumerGroup?: ConsumerGroup;
+    fetch: (nodeId: number, assignment: Assignment) => Promise<FetchResponse>;
+    onResponse: (fetcherId: number, response: FetchResponse) => Promise<void>;
+};
+
+export class Fetcher extends EventEmitter<{ stopped: [] }> {
+    private isRunning = false;
+
+    constructor(
+        private fetcherId: number,
+        private options: FetcherOptions,
+    ) {
+        super();
+    }
+
+    public async loop() {
+        this.isRunning = true;
+
+        try {
+            while (this.isRunning) {
+                await this.step();
+            }
+        } finally {
+            this.isRunning = false;
+            this.emit('stopped');
+        }
+    }
+
+    @trace()
+    private async step() {
+        const { nodeId, assignment, consumerGroup, fetch, onResponse } = this.options;
+
+        const response = await fetch(nodeId, assignment);
+        if (!this.isRunning) {
+            return;
+        }
+        consumerGroup?.handleLastHeartbeat();
+        await onResponse(this.fetcherId, response);
+        consumerGroup?.handleLastHeartbeat();
+    }
+
+    public async stop() {
+        if (!this.isRunning) {
+            return;
+        }
+
+        const stopPromise = new Promise<void>((resolve) => {
+            this.once('stopped', resolve);
+        });
+        this.isRunning = false;
+        return stopPromise;
+    }
+}

package/src/consumer/offset-manager.ts

@@ -1,9 +1,12 @@
-import { API } from "../api";
-import { IsolationLevel } from "../api/fetch";
-import { Assignment } from "../api/sync-group";
-import { Cluster } from "../cluster";
-import { distributeMessagesToTopicPartitionLeaders } from "../distributors/messages-to-topic-partition-leaders";
-import { ConsumerMetadata } from "./consumer-metadata";
+import { API } from '../api';
+import { IsolationLevel } from '../api/fetch';
+import { Assignment } from '../api/sync-group';
+import { Cluster } from '../cluster';
+import { distributeMessagesToTopicPartitionLeaders } from '../distributors/messages-to-topic-partition-leaders';
+import { createTracer } from '../utils/tracer';
+import { ConsumerMetadata } from './consumer-metadata';
+
+const trace = createTracer('OffsetManager');
 
 type OffsetManagerOptions = {
     cluster: Cluster;
@@ -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

@@ -0,0 +1,53 @@
+import { EventEmitter } from 'stream';
+import { Batch } from '../types';
+import { createTracer } from '../utils/tracer';
+
+const trace = createTracer('Processor');
+
+type ProcessorOptions = {
+    poll: () => Promise<Batch>;
+    process: (batch: Batch) => Promise<void>;
+};
+
+export class Processor extends EventEmitter<{ stopped: [] }> {
+    private isRunning = false;
+
+    constructor(private options: ProcessorOptions) {
+        super();
+    }
+
+    public async loop() {
+        this.isRunning = true;
+
+        try {
+            while (this.isRunning) {
+                await this.step();
+            }
+        } finally {
+            this.isRunning = false;
+            this.emit('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;
+        }
+
+        const stopPromise = new Promise<void>((resolve) => {
+            this.once('stopped', resolve);
+        });
+        this.isRunning = false;
+        return stopPromise;
+    }
+}

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

@@ -1,9 +1,9 @@
-import { describe, expect, it } from "vitest";
-import { distributeAssignmentsToNodesBalanced, distributeAssignmentsToNodesOptimized } from "./assignments-to-replicas";
+import { describe, expect, it } from 'vitest';
+import { distributeAssignmentsToNodesBalanced, distributeAssignmentsToNodesOptimized } from './assignments-to-replicas';
 
-describe("Distribute assignments to replica ids", () => {
-    describe("distributeAssignmentsToNodesBalanced", () => {
-        it("smoke", () => {
+describe('Distribute assignments to replica ids', () => {
+    describe('distributeAssignmentsToNodesBalanced', () => {
+        it('smoke', () => {
             const result = distributeAssignmentsToNodesBalanced({ topic: [0, 1] }, { topic: { 0: [0, 1], 1: [1, 2] } });
             expect(result).toMatchInlineSnapshot(`
               {
@@ -22,8 +22,8 @@ describe("Distribute assignments to replica ids", () => {
         });
     });
 
-    describe("distributeAssignmentsToNodesOptimized", () => {
-        it("smoke", () => {
+    describe('distributeAssignmentsToNodesOptimized', () => {
+        it('smoke', () => {
             const result = distributeAssignmentsToNodesOptimized(
                 { topic: [0, 1] },
                 { topic: { 0: [0, 1], 1: [1, 2] } },

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 = (
@@ -50,7 +50,7 @@ export const distributeAssignmentsToNodesOptimized = (
         }
 
         result[parseInt(replicaId)] = partitions.reduce((acc, partition) => {
-            const [topicName, partitionId] = partition.split(":");
+            const [topicName, partitionId] = partition.split(':');
             acc[topicName] ??= [];
             acc[topicName].push(parseInt(partitionId));
             return acc;
@@ -81,5 +81,3 @@ const getPartitionsByReplica = (assignment: Assignment, topicPartitionReplicaIds
     }
     return Object.entries(partitionsByReplicaId);
 };
-
-export const distributeAssignmentsToNodes = distributeAssignmentsToNodesBalanced;

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

@@ -1,11 +1,11 @@
-import { describe, expect, it } from "vitest";
-import { distributeMessagesToTopicPartitionLeaders } from "./messages-to-topic-partition-leaders";
+import { describe, expect, it } from 'vitest';
+import { distributeMessagesToTopicPartitionLeaders } from './messages-to-topic-partition-leaders';
 
-describe("Distribute messages to partition leader ids", () => {
-    describe("distributeMessagesToTopicPartitionLeaders", () => {
-        it("snoke", () => {
+describe('Distribute messages to partition leader ids', () => {
+    describe('distributeMessagesToTopicPartitionLeaders', () => {
+        it('snoke', () => {
             const result = distributeMessagesToTopicPartitionLeaders(
-                [{ topic: "topic", partition: 0, key: null, value: null, offset: 0n, timestamp: 0n, headers: {} }],
+                [{ topic: 'topic', partition: 0, key: null, value: null, offset: 0n, timestamp: 0n, headers: {} }],
                 { topic: { 0: 1 } },
             );
             expect(result).toMatchInlineSnapshot(`

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,3 +1,9 @@
-export * from "./client";
-export * from "./api";
-export * from "./types";
+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

@@ -1,7 +1,10 @@
-import { API, API_ERROR } from "./api";
-import { Cluster } from "./cluster";
-import { delay } from "./utils/delay";
-import { KafkaTSApiError } from "./utils/error";
+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 { Metadata } from "../metadata";
-import { Message } from "../types";
-import { delay } from "../utils/delay";
-import { KafkaTSApiError } from "../utils/error";
-import { memo } from "../utils/memo";
+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,9 +1,11 @@
 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>;
 };
+
+export type Batch = Required<Message>[];

package/src/utils/api.ts

@@ -1,11 +1,11 @@
-import { Decoder } from "./decoder";
-import { Encoder } from "./encoder";
+import { Decoder } from './decoder';
+import { Encoder } from './encoder';
 
 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;
+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]));