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

package/README.md

@@ -92,10 +92,56 @@ The existing high-level libraries (e.g. kafkajs) are missing a few crucial featu
 -   **Consuming messages without consumer groups** - When you don't need the consumer to track the partition offsets, you can simply create a consumer without groupId and always either start consuming messages from the beginning or from the latest partition offset.
 -   **Low-level API requests** - It's possible to communicate directly with the Kafka cluster using the kafka api protocol.
 
+## Configuration
+
+### `createKafkaClient()`
+
+| Name             | Type                   | Required | Default | Description                                           |
+| ---------------- | ---------------------- | -------- | ------- | ----------------------------------------------------- |
+| clientId         | string                 | false    | _null_  | The client id used for all requests.                  |
+| bootstrapServers | TcpSocketConnectOpts[] | true     |         | List of kafka brokers for initial cluster discovery.  |
+| sasl             | SASLProvider           | false    |         | SASL provider (see "Supported SASL mechanisms" below) |
+| ssl              | TLSSocketOptions       | false    |         | SSL configuration.                                    |
+
+### `kafka.startConsumer()`
+
+| Name                   | Type                                   | Required | Default   | Description                                                                                                                                                                                                                                                                      |
+| ---------------------- | -------------------------------------- | -------- | --------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| topics                 | string[]                               | true     |           | List of topics to subscribe to                                                                                                                                                                                                                                                   |
+| groupId                | string                                 | false    | _null_    | Consumer group id                                                                                                                                                                                                                                                                |
+| groupInstanceId        | string                                 | false    | _null_    | Consumer group instance id                                                                                                                                                                                                                                                       |
+| rackId                 | string                                 | false    | _null_    | Rack id                                                                                                                                                                                                                                                                          |
+| isolationLevel         | IsolationLevel                         | false    |           | Isolation level                                                                                                                                                                                                                                                                  |
+| sessionTimeoutMs       | number                                 | false    | 30000     | Session timeout in milliseconds                                                                                                                                                                                                                                                  |
+| rebalanceTimeoutMs     | number                                 | false    | 60000     | Rebalance timeout in milliseconds                                                                                                                                                                                                                                                |
+| maxWaitMs              | number                                 | false    | 5000      | Fetch long poll timeout in milliseconds                                                                                                                                                                                                                                          |
+| minBytes               | number                                 | false    | 1         | Minimum number of bytes to wait for before returning a fetch response                                                                                                                                                                                                            |
+| maxBytes               | number                                 | false    | 1_048_576 | Maximum number of bytes to return in the fetch response                                                                                                                                                                                                                          |
+| partitionMaxBytes      | number                                 | false    | 1_048_576 | Maximum number of bytes to return per partition in the fetch response                                                                                                                                                                                                            |
+| allowTopicAutoCreation | boolean                                | false    | false     | Allow kafka to auto-create topic when it doesn't exist                                                                                                                                                                                                                           |
+| fromBeginning          | boolean                                | false    | false     | Start consuming from the beginning of the topic                                                                                                                                                                                                                                  |
+| batchGranularity       | BatchGranularity                       | false    | partition | Controls messages split from fetch response. Also controls how often offsets are committed. **onBatch** will include messages:<br/>- **partition** - from a single batch<br/>- **topic** - from all topic partitions<br/>- **broker** - from all assignned topics and partitions |
+| concurrency            | number                                 | false    | 1         | How many batches to process concurrently                                                                                                                                                                                                                                         |
+| onMessage              | (message: Message) => Promise<unknown> | true     |           | Callback executed on every message                                                                                                                                                                                                                                               |
+| onBatch                | (batch: Message[]) => Promise<unknown> | true     |           | Callback executed on every batch of messages (based on **batchGranuality**)                                                                                                                                                                                                      |
+
+### `kafka.createProducer()`
+
+| Name                   | Type        | Required | Default | Description                                                                             |
+| ---------------------- | ----------- | -------- | ------- | --------------------------------------------------------------------------------------- |
+| allowTopicAutoCreation | boolean     | false    | false   | Allow kafka to auto-create topic when it doesn't exist                                  |
+| partitioner            | Partitioner | false    |         | Custom partitioner function. By default, it uses a default java-compatible partitioner. |
+
+### Supported SASL mechanisms
+
+-   PLAIN: `saslPlain({ username, password })`
+-   SCRAM-SHA-256: `saslScramSha256({ username, password })`
+-   SCRAM-SHA-512: `saslScramSha512({ username, password })`
+
+Custom SASL mechanisms can be implemented following the `SASLProvider` interface. See [src/auth](./src/auth) for examples.
+
 ## Backlog
 
-Minimal set of features required before a stable release:
+Minimal set of features left to implement before a stable release:
 
--   Partitioner (Currently have to specify the partition on producer.send())
 -   API versioning (Currently only tested against Kafka 3.7+)
--   SASL SCRAM-SHA-512 support

package/examples/package.json

@@ -4,7 +4,8 @@
     "description": "",
     "main": "dist/replicator.js",
     "scripts": {
-        "test": "echo \"Error: no test specified\" && exit 1"
+        "test": "echo \"Error: no test specified\" && exit 1",
+        "build": "tsc"
     },
     "dependencies": {
         "kafka-ts": "file:../"

package/examples/src/client.ts

@@ -1,9 +1,9 @@
 import { readFileSync } from 'fs';
-import { createKafkaClient } from 'kafka-ts';
+import { createKafkaClient, saslScramSha512 } from 'kafka-ts';
 
 export const kafka = createKafkaClient({
     clientId: 'examples',
     bootstrapServers: [{ host: 'localhost', port: 9092 }],
-    sasl: { mechanism: 'SCRAM-SHA-256', username: 'admin', password: 'admin' },
+    sasl: saslScramSha512({ username: 'admin', password: 'admin' }),
     ssl: { ca: readFileSync('../certs/ca.crt').toString() },
 });

package/examples/src/consumer.ts

@@ -8,7 +8,7 @@ import { kafka } from './client';
         onBatch: (batch) => {
             console.log(batch);
         },
-        granularity: 'broker',
+        batchGranularity: 'broker',
         concurrency: 10,
     });
 

package/examples/src/producer.ts

@@ -11,7 +11,7 @@ rl.on('line', async (line) => {
         {
             topic: 'example-topic-f',
             key: null,
-            value: line,
+            value: Buffer.from(line),
             partition: 0,
         },
     ]);

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "kafka-ts",
-    "version": "0.0.1-beta.0",
+    "version": "0.0.1-beta.2",
     "main": "dist/index.js",
     "author": "Priit Käärd",
     "license": "MIT",

package/src/__snapshots__/{request-handler.test.ts.snap → cluster.test.ts.snap}

@@ -360,14 +360,60 @@ exports[`Request handler > should fetch messages 1`] = `
                   "attributes": 0,
                   "headers": [
                     {
-                      "key": "header-key",
-                      "value": "header-value",
+                      "key": {
+                        "data": [
+                          104,
+                          101,
+                          97,
+                          100,
+                          101,
+                          114,
+                          45,
+                          107,
+                          101,
+                          121,
+                        ],
+                        "type": "Buffer",
+                      },
+                      "value": {
+                        "data": [
+                          104,
+                          101,
+                          97,
+                          100,
+                          101,
+                          114,
+                          45,
+                          118,
+                          97,
+                          108,
+                          117,
+                          101,
+                        ],
+                        "type": "Buffer",
+                      },
                     },
                   ],
-                  "key": "key",
+                  "key": {
+                    "data": [
+                      107,
+                      101,
+                      121,
+                    ],
+                    "type": "Buffer",
+                  },
                   "offsetDelta": 0,
                   "timestampDelta": 0n,
-                  "value": "value",
+                  "value": {
+                    "data": [
+                      118,
+                      97,
+                      108,
+                      117,
+                      101,
+                    ],
+                    "type": "Buffer",
+                  },
                 },
               ],
             },
@@ -937,6 +983,156 @@ exports[`Request handler > should request metadata for all topics 1`] = `
   "controllerId": 0,
   "throttleTimeMs": 0,
   "topics": [
+    {
+      "_tag": undefined,
+      "errorCode": 0,
+      "isInternal": false,
+      "name": "my-topic",
+      "partitions": [
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 5,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 6,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 0,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 2,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 7,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 1,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 8,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 9,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 3,
+          "replicaNodes": [
+            0,
+          ],
+        },
+        {
+          "_tag": undefined,
+          "errorCode": 0,
+          "isrNodes": [
+            0,
+          ],
+          "leaderEpoch": 0,
+          "leaderId": 0,
+          "offlineReplicas": [],
+          "partitionIndex": 4,
+          "replicaNodes": [
+            0,
+          ],
+        },
+      ],
+      "topicAuthorizedOperations": -2147483648,
+      "topicId": "Any<UUID>",
+    },
     {
       "_tag": undefined,
       "errorCode": 0,

package/src/api/fetch.ts

@@ -130,11 +130,11 @@ const decodeRecords = (decoder: Decoder) => {
                 attributes: record.readInt8(),
                 timestampDelta: record.readVarLong(),
                 offsetDelta: record.readVarInt(),
-                key: record.readVarIntString(),
-                value: record.readVarIntString(),
+                key: record.readVarIntBuffer(),
+                value: record.readVarIntBuffer(),
                 headers: record.readCompactArray((header) => ({
-                    key: header.readVarIntString(),
-                    value: header.readVarIntString(),
+                    key: header.readVarIntBuffer(),
+                    value: header.readVarIntBuffer(),
                 })),
             })),
         });

package/src/api/produce.ts

@@ -28,11 +28,11 @@ export const PRODUCE = createApi({
                         attributes: number;
                         timestampDelta: bigint;
                         offsetDelta: number;
-                        key: string | null;
-                        value: string | null;
+                        key: Buffer | null;
+                        value: Buffer | null;
                         headers: {
-                            key: string;
-                            value: string;
+                            key: Buffer;
+                            value: Buffer;
                         }[];
                     }[];
                 }[];
@@ -61,10 +61,10 @@ export const PRODUCE = createApi({
                                     .writeInt8(record.attributes)
                                     .writeVarLong(record.timestampDelta)
                                     .writeVarInt(record.offsetDelta)
-                                    .writeVarIntString(record.key)
-                                    .writeVarIntString(record.value)
+                                    .writeVarIntBuffer(record.key)
+                                    .writeVarIntBuffer(record.value)
                                     .writeVarIntArray(record.headers, (encoder, header) =>
-                                        encoder.writeVarIntString(header.key).writeVarIntString(header.value),
+                                        encoder.writeVarIntBuffer(header.key).writeVarIntBuffer(header.value),
                                     )
                                     .value();
 

package/src/auth/index.ts

@@ -0,0 +1,2 @@
+export { saslPlain } from './plain';
+export { saslScramSha256, saslScramSha512 } from './scram';

package/src/auth/plain.ts

@@ -0,0 +1,10 @@
+import { API } from "../api";
+import { SASLProvider } from "../broker";
+
+export const saslPlain = ({ username, password }: { username: string; password: string }): SASLProvider => ({
+    mechanism: 'PLAIN',
+    authenticate: async ({ sendRequest }) => {
+        const authBytes = [null, username, password].join('\u0000');
+        await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(authBytes) });
+    },
+});

package/src/auth/scram.ts

@@ -0,0 +1,52 @@
+import { API } from '../api';
+import { SASLProvider } from '../broker';
+import { base64Decode, base64Encode, generateNonce, hash, hmac, saltPassword, xor } from '../utils/crypto';
+import { KafkaTSError } from '../utils/error';
+
+const saslScram =
+    ({ mechanism, keyLength, digest }: { mechanism: string; keyLength: number; digest: string }) =>
+    ({ username, password }: { username: string; password: string }): SASLProvider => ({
+        mechanism,
+        authenticate: async ({ sendRequest }) => {
+            const nonce = generateNonce();
+            const firstMessage = `n=${username},r=${nonce}`;
+
+            const { authBytes } = await sendRequest(API.SASL_AUTHENTICATE, {
+                authBytes: Buffer.from(`n,,${firstMessage}`),
+            });
+            if (!authBytes) {
+                throw new KafkaTSError('No auth response');
+            }
+
+            const response = Object.fromEntries(
+                authBytes
+                    .toString()
+                    .split(',')
+                    .map((pair) => pair.split('=')),
+            ) as { r: string; s: string; i: string };
+
+            const rnonce = response.r;
+            if (!rnonce.startsWith(nonce)) {
+                throw new KafkaTSError('Invalid nonce');
+            }
+            const iterations = parseInt(response.i);
+            const salt = base64Decode(response.s);
+
+            const saltedPassword = await saltPassword(password, salt, iterations, keyLength, digest);
+            const clientKey = hmac(saltedPassword, 'Client Key', digest);
+            const clientKeyHash = hash(clientKey, digest);
+
+            let finalMessage = `c=${base64Encode('n,,')},r=${rnonce}`;
+
+            const fullMessage = `${firstMessage},${authBytes.toString()},${finalMessage}`;
+            const clientSignature = hmac(clientKeyHash, fullMessage, digest);
+            const clientProof = base64Encode(xor(clientKey, clientSignature));
+
+            finalMessage += `,p=${clientProof}`;
+
+            await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(finalMessage) });
+        },
+    });
+
+export const saslScramSha256 = saslScram({ mechanism: 'SCRAM-SHA-256', keyLength: 32, digest: 'sha256' });
+export const saslScramSha512 = saslScram({ mechanism: 'SCRAM-SHA-512', keyLength: 64, digest: 'sha512' });

package/src/broker.ts

@@ -2,20 +2,18 @@ import { TcpSocketConnectOpts } from 'net';
 import { TLSSocketOptions } from 'tls';
 import { API } from './api';
 import { Connection, SendRequest } from './connection';
-import { base64Decode, base64Encode, generateNonce, hash, hmac, saltPassword, xor } from './utils/crypto';
 import { KafkaTSError } from './utils/error';
 import { memo } from './utils/memo';
 
-export type SASLOptions = {
-    mechanism: 'PLAIN' | 'SCRAM-SHA-256';
-    username: string;
-    password: string;
+export type SASLProvider = {
+    mechanism: string;
+    authenticate: (context: { sendRequest: SendRequest }) => Promise<void>;
 };
 
 type BrokerOptions = {
     clientId: string | null;
     options: TcpSocketConnectOpts;
-    sasl: SASLOptions | null;
+    sasl: SASLProvider | null;
     ssl: TLSSocketOptions | null;
 };
 
@@ -69,78 +67,6 @@ export class Broker {
     }
 
     private async saslAuthenticate() {
-        if (!this.options.sasl) {
-            return;
-        }
-
-        const { mechanism } = this.options.sasl;
-        const authenticate = { PLAIN: plainProvider, 'SCRAM-SHA-256': scramSha256Provider }[mechanism as string];
-        if (!authenticate) {
-            throw new KafkaTSError(`SASL mechanism ${mechanism} is not supported`);
-        }
-
-        await authenticate({
-            ...this.options.sasl,
-            sendRequest: this.sendRequest.bind(this),
-        });
+        await this.options.sasl?.authenticate({ sendRequest: this.sendRequest });
     }
 }
-
-const plainProvider = async ({
-    username,
-    password,
-    sendRequest,
-}: {
-    username: string;
-    password: string;
-    sendRequest: SendRequest;
-}) => {
-    const authBytes = [null, username, password].join('\u0000');
-    await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(authBytes) });
-};
-
-const scramSha256Provider = async ({
-    username,
-    password,
-    sendRequest,
-}: {
-    username: string;
-    password: string;
-    sendRequest: SendRequest;
-}) => {
-    const nonce = generateNonce();
-    const firstMessage = `n=${username},r=${nonce}`;
-
-    const { authBytes } = await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(`n,,${firstMessage}`) });
-    if (!authBytes) {
-        throw new KafkaTSError('No auth response');
-    }
-
-    const response = Object.fromEntries(
-        authBytes
-            .toString()
-            .split(',')
-            .map((pair) => pair.split('=')),
-    ) as { r: string; s: string; i: string };
-
-    const rnonce = response.r;
-    if (!rnonce.startsWith(nonce)) {
-        throw new KafkaTSError('Invalid nonce');
-    }
-    const iterations = parseInt(response.i);
-    const salt = base64Decode(response.s);
-
-    const saltedPassword = await saltPassword(password, salt, iterations, 32, 'sha256');
-    const clientKey = hmac(saltedPassword, 'Client Key');
-    const clientKeyHash = hash(clientKey);
-
-    let finalMessage = `c=${base64Encode('n,,')},r=${rnonce}`;
-
-    const fullMessage = `${firstMessage},${authBytes.toString()},${finalMessage}`;
-    const clientSignature = hmac(clientKeyHash, fullMessage);
-    const clientProof = base64Encode(xor(clientKey, clientSignature));
-
-    finalMessage += `,p=${clientProof}`;
-
-    await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(finalMessage) });
-};

package/src/client.ts

@@ -1,6 +1,6 @@
 import { TcpSocketConnectOpts } from 'net';
 import { TLSSocketOptions } from 'tls';
-import { SASLOptions } from './broker';
+import { SASLProvider } from './broker';
 import { Cluster } from './cluster';
 import { Consumer, ConsumerOptions } from './consumer/consumer';
 import { Producer, ProducerOptions } from './producer/producer';
@@ -8,7 +8,7 @@ import { Producer, ProducerOptions } from './producer/producer';
 type ClientOptions = {
     clientId?: string | null;
     bootstrapServers: TcpSocketConnectOpts[];
-    sasl?: SASLOptions | null;
+    sasl?: SASLProvider | null;
     ssl?: TLSSocketOptions | null;
 };
 

package/src/cluster.test.ts

@@ -3,6 +3,7 @@ import { readFileSync } from 'fs';
 import { afterAll, beforeAll, describe, expect, it } from 'vitest';
 import { API } from './api';
 import { KEY_TYPE } from './api/find-coordinator';
+import { saslPlain } from './auth';
 import { createKafkaClient } from './client';
 import { Cluster } from './cluster';
 import { KafkaTSApiError } from './utils/error';
@@ -10,7 +11,7 @@ import { KafkaTSApiError } from './utils/error';
 export const kafka = createKafkaClient({
     clientId: 'kafka-ts',
     bootstrapServers: [{ host: 'localhost', port: 9092 }],
-    sasl: { mechanism: 'PLAIN', username: 'admin', password: 'admin' },
+    sasl: saslPlain({ username: 'admin', password: 'admin' }),
     ssl: { ca: readFileSync('./certs/ca.crt').toString() },
 });
 
@@ -148,12 +149,12 @@ describe.sequential('Request handler', () => {
                                     attributes: 0,
                                     offsetDelta: 0,
                                     timestampDelta: 0n,
-                                    key: 'key',
-                                    value: 'value',
+                                    key: Buffer.from('key'),
+                                    value: Buffer.from('value'),
                                     headers: [
                                         {
-                                            key: 'header-key',
-                                            value: 'header-value',
+                                            key: Buffer.from('header-key'),
+                                            value: Buffer.from('header-value'),
                                         },
                                     ],
                                 },

package/src/cluster.ts

@@ -1,14 +1,14 @@
 import { TcpSocketConnectOpts } from 'net';
 import { TLSSocketOptions } from 'tls';
 import { API } from './api';
-import { Broker, SASLOptions } from './broker';
+import { Broker, SASLProvider } from './broker';
 import { SendRequest } from './connection';
 import { ConnectionError, KafkaTSError } from './utils/error';
 
 type ClusterOptions = {
     clientId: string | null;
     bootstrapServers: TcpSocketConnectOpts[];
-    sasl: SASLOptions | null;
+    sasl: SASLProvider | null;
     ssl: TLSSocketOptions | null;
 };
 

package/src/consumer/consumer.ts

@@ -9,7 +9,7 @@ import { ConnectionError, KafkaTSApiError } from '../utils/error';
 import { defaultRetrier, Retrier } from '../utils/retrier';
 import { ConsumerGroup } from './consumer-group';
 import { ConsumerMetadata } from './consumer-metadata';
-import { FetchManager, Granularity } from './fetch-manager';
+import { FetchManager, BatchGranularity } from './fetch-manager';
 import { OffsetManager } from './offset-manager';
 
 export type ConsumerOptions = {
@@ -27,7 +27,7 @@ export type ConsumerOptions = {
     allowTopicAutoCreation?: boolean;
     fromBeginning?: boolean;
     retrier?: Retrier;
-    granularity?: Granularity;
+    batchGranularity?: BatchGranularity;
     concurrency?: number;
 } & ({ onBatch: (messages: Required<Message>[]) => unknown } | { onMessage: (message: Required<Message>) => unknown });
 
@@ -52,13 +52,13 @@ export class Consumer {
             rebalanceTimeoutMs: options.rebalanceTimeoutMs ?? 60_000,
             maxWaitMs: options.maxWaitMs ?? 5000,
             minBytes: options.minBytes ?? 1,
-            maxBytes: options.maxBytes ?? 1_000_000,
-            partitionMaxBytes: options.partitionMaxBytes ?? 1_000_000,
+            maxBytes: options.maxBytes ?? 1_048_576,
+            partitionMaxBytes: options.partitionMaxBytes ?? 1_048_576,
             isolationLevel: options.isolationLevel ?? IsolationLevel.READ_UNCOMMITTED,
             allowTopicAutoCreation: options.allowTopicAutoCreation ?? false,
             fromBeginning: options.fromBeginning ?? false,
             retrier: options.retrier ?? defaultRetrier,
-            granularity: options.granularity ?? 'broker',
+            batchGranularity: options.batchGranularity ?? 'partition',
             concurrency: options.concurrency ?? 1,
         };
 
@@ -118,7 +118,7 @@ export class Consumer {
     }
 
     private startFetchManager = async () => {
-        const { granularity, concurrency } = this.options;
+        const { batchGranularity, concurrency } = this.options;
 
         while (!this.stopHook) {
             const nodeAssignments = Object.entries(
@@ -137,7 +137,7 @@ export class Consumer {
                 metadata: this.metadata,
                 consumerGroup: this.consumerGroup,
                 nodeAssignments,
-                granularity,
+                batchGranularity,
                 concurrency: numProcessors,
             });
 

package/src/consumer/fetch-manager.ts

@@ -11,7 +11,7 @@ import { Processor } from './processor';
 
 const trace = createTracer('FetchManager');
 
-export type Granularity = 'partition' | 'topic' | 'broker';
+export type BatchGranularity = 'partition' | 'topic' | 'broker';
 
 type FetchManagerOptions = {
     fetch: (nodeId: number, assignment: Assignment) => Promise<ReturnType<(typeof API.FETCH)['response']>>;
@@ -19,7 +19,7 @@ type FetchManagerOptions = {
     metadata: Metadata;
     consumerGroup?: ConsumerGroup;
     nodeAssignments: { nodeId: number; assignment: Assignment }[];
-    granularity: Granularity;
+    batchGranularity: BatchGranularity;
     concurrency: number;
 };
 
@@ -110,9 +110,9 @@ export class FetchManager extends EventEmitter<{ data: []; checkpoint: [number];
     }
 
     private async onResponse(fetcherId: number, response: ReturnType<(typeof API.FETCH)['response']>) {
-        const { metadata, granularity } = this.options;
+        const { metadata, batchGranularity } = this.options;
 
-        const batches = fetchResponseToBatches(response, granularity, metadata);
+        const batches = fetchResponseToBatches(response, batchGranularity, metadata);
         if (batches.length) {
             this.queue.push(...batches);
             this.queue.push({ kind: 'checkpoint', fetcherId });
@@ -138,7 +138,7 @@ export class FetchManager extends EventEmitter<{ data: []; checkpoint: [number];
 
 const fetchResponseToBatches = (
     batch: ReturnType<typeof API.FETCH.response>,
-    granularity: Granularity,
+    batchGranularity: BatchGranularity,
     metadata: Metadata,
 ): Batch[] => {
     const brokerTopics = batch.responses.map(({ topicId, partitions }) =>
@@ -159,7 +159,7 @@ const fetchResponseToBatches = (
         ),
     );
 
-    switch (granularity) {
+    switch (batchGranularity) {
         case 'broker':
             const messages = brokerTopics.flatMap((topicPartition) =>
                 topicPartition.flatMap((partitionMessages) => partitionMessages),
@@ -174,6 +174,6 @@ const fetchResponseToBatches = (
                 topicPartition.map((partitionMessages) => partitionMessages),
             );
         default:
-            throw new KafkaTSError(`Unhandled batch granularity: ${granularity}`);
+            throw new KafkaTSError(`Unhandled batch granularity: ${batchGranularity}`);
     }
 };

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,7 @@
-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';

package/src/producer/producer.ts

@@ -1,6 +1,7 @@
 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';
@@ -9,6 +10,7 @@ import { memo } from '../utils/memo';
 
 export type ProducerOptions = {
     allowTopicAutoCreation?: boolean;
+    partitioner?: Partitioner;
 };
 
 export class Producer {
@@ -17,6 +19,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,8 +28,10 @@ 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[]) {
@@ -39,7 +44,7 @@ export class Producer {
         await this.metadata.fetchMetadataIfNecessary({ topics, allowTopicAutoCreation });
 
         const nodeTopicPartitionMessages = distributeMessagesToTopicPartitionLeaders(
-            messages,
+            messages.map(message => ({ ...message, partition: this.partition(message) })),
             this.metadata.getTopicPartitionLeaderIds(),
         );
 
@@ -83,8 +88,8 @@ export class Producer {
                                     key: message.key,
                                     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

@@ -3,8 +3,8 @@ export type Message = {
     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/crypto.ts

@@ -9,6 +9,7 @@ export const saltPassword = (password: string, salt: string, iterations: number,
 
 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) => createHash('sha256').update(data).digest();
-export const hmac = (key: Buffer, data: Buffer | string) => createHmac('sha256', key).update(data).digest();
+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;
     }

package/src/utils/encoder.ts

@@ -91,11 +91,11 @@ export class Encoder {
         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) {

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;