feat: enhance TopicWriter and TopicReader with compression options and improve handling of payloads

Signed-off-by: Vladislav Polyakov <polRk@ydb.tech>

Author: polRk

.changeset/pre.json

@@ -36,6 +36,7 @@
     "sweet-carpets-talk",
     "tall-colts-follow",
     "tender-symbols-drive",
+    "violet-words-wave",
     "wicked-carrots-own",
     "yellow-banks-burn"
   ]

.changeset/violet-words-wave.md

@@ -0,0 +1,5 @@
+---
+'@ydbjs/topic': patch
+---
+
+enhance TopicWriter and TopicReader with compression options

package-lock.json

@@ -18,7 +18,7 @@
 		"publish-packages": "turbo run clean build test attw && changeset version && changeset publish"
 	},
 	"dependencies": {
-		"@types/node": "^22.13.10"
+		"@types/node": "^22.15.29"
 	},
 	"devDependencies": {
 		"oxlint": "^0.16.4",

package.json

@@ -1,5 +1,11 @@
 # @ydbjs/topic
 
+## 6.0.0-alpha.11
+
+### Patch Changes
+
+- Enhance TopicWriter and TopicReader with compression options
+
 ## 6.0.0-alpha.10
 
 ### Patch Changes

packages/topic/CHANGELOG.md

@@ -1,6 +1,6 @@
 {
 	"name": "@ydbjs/topic",
-	"version": "6.0.0-alpha.10",
+	"version": "6.0.0-alpha.11",
 	"description": "",
 	"license": "Apache-2.0",
 	"type": "module",

packages/topic/package.json

@@ -5,7 +5,7 @@ import { nextTick } from "node:process";
 import { create, protoInt64, toJson } from "@bufbuild/protobuf";
 import { type Duration, DurationSchema, type Timestamp, timestampDate, timestampFromDate } from "@bufbuild/protobuf/wkt";
 import { StatusIds_StatusCode } from "@ydbjs/api/operation";
-import { type OffsetsRange, OffsetsRangeSchema, type StreamReadMessage_CommitOffsetRequest_PartitionCommitOffset, StreamReadMessage_CommitOffsetRequest_PartitionCommitOffsetSchema, type StreamReadMessage_FromClient, StreamReadMessage_FromClientSchema, type StreamReadMessage_FromServer, StreamReadMessage_FromServerSchema, type StreamReadMessage_InitRequest_TopicReadSettings, StreamReadMessage_InitRequest_TopicReadSettingsSchema, type StreamReadMessage_ReadResponse, TopicServiceDefinition } from "@ydbjs/api/topic";
+import { Codec, type OffsetsRange, OffsetsRangeSchema, type StreamReadMessage_CommitOffsetRequest_PartitionCommitOffset, StreamReadMessage_CommitOffsetRequest_PartitionCommitOffsetSchema, type StreamReadMessage_FromClient, StreamReadMessage_FromClientSchema, type StreamReadMessage_FromServer, StreamReadMessage_FromServerSchema, type StreamReadMessage_InitRequest_TopicReadSettings, StreamReadMessage_InitRequest_TopicReadSettingsSchema, type StreamReadMessage_ReadResponse, TopicServiceDefinition } from "@ydbjs/api/topic";
 import type { Driver } from "@ydbjs/core";
 import { YDBError } from "@ydbjs/error";
 import { type RetryConfig, retry } from "@ydbjs/retry";
@@ -91,15 +91,33 @@ type TopicCommitPromise = {
 }
 
 export type TopicReaderOptions<Payload = Uint8Array> = {
+	// Topic path or an array of topic sources.
 	topic: string | TopicReaderSource | TopicReaderSource[]
+	// Consumer name.
 	consumer: string
+	// Maximum size of the internal buffer in bytes.
+	// If not provided, the default is 1MB.
 	maxBufferBytes?: bigint
+	// How often to update the token in milliseconds.
 	updateTokenIntervalMs?: number
-
+	// Compression options for the payload.
+	compression?: {
+		// Custom decompression function that can be used to decompress the payload before emitting it.
+		decompress?(codec: Codec, payload: Uint8Array): Uint8Array | Promise<Uint8Array>
+	}
+	// Custom decode function to decode the payload.
+	// If not provided, the payload will be returned as is.
+	// Decode function calls after decompression, if compression is used.
 	decode?(payload: Uint8Array): Payload
-
+	// Hooks for partition session events.
+	// Called when a partition session is started.
+	// It can be used to initialize the partition session, for example, to set the read offset.
 	onPartitionSessionStart?: onPartitionSessionStartCallback
+	// Called when a partition session is stopped.
+	// It can be used to commit the offsets for the partition session.
 	onPartitionSessionStop?: onPartitionSessionStopCallback
+	// Called when receive commit offset response from server.
+	// This callback is called after the offsets are committed to the server.
 	onCommittedOffset?: onCommittedOffsetCallback
 }
 
@@ -250,6 +268,11 @@ export class TopicReader<Payload = Uint8Array> implements Disposable {
 					}
 				}
 
+				// If there are no pending commits for this partition session, remove it from the map.
+				if (pendingCommits && pendingCommits.length === 0) {
+					this.#pendingCommits.delete(partitionSession.partitionSessionId);
+				}
+
 				this.#fromClientEmitter.emit('message', create(StreamReadMessage_FromClientSchema, {
 					clientMessage: {
 						case: 'startPartitionSessionResponse',
@@ -576,7 +599,7 @@ export class TopicReader<Payload = Uint8Array> implements Disposable {
 	read(options: { limit?: number, waitMs?: number, signal?: AbortSignal } = {}): AsyncIterable<TopicMessage<Payload>[]> {
 		let limit = options.limit || Infinity,
 			signal = options.signal,
-			waitMs = options.waitMs || 60000;
+			waitMs = options.waitMs || 60_000;
 
 		// Check if the reader has been disposed, cannot read with disposed reader
 		if (this.#disposed) {
@@ -750,14 +773,34 @@ export class TopicReader<Payload = Uint8Array> implements Disposable {
 											break;
 										}
 
+										let data = msg.data;
+										if (batch.codec !== Codec.UNSPECIFIED) {
+											if (!this.#options.compression || !this.#options.compression.decompress) {
+												dbg('error: cannot decompress message with codec %s, no decompression function provided', batch.codec);
+
+												throw new Error(`Cannot decompress message with codec ${batch.codec}, no decompression function provided`);
+											}
+
+											// Decompress the message data using the provided decompress function
+											try {
+												// eslint-disable-next-line no-await-in-loop
+												data = await this.#options.compression.decompress(batch.codec, msg.data);
+											} catch (err) {
+												dbg('error: decompression failed for message with codec %s: %O', batch.codec, err);
+
+												throw new Error(`Decompression failed for message with codec ${batch.codec}`, { cause: err });
+											}
+										}
+
 										// Process the message
 										let message: TopicMessage<Payload> = {
 											partitionSessionId: partitionSession.partitionSessionId,
 											partitionId: partitionSession.partitionId,
 											producerId: batch.producerId,
 											seqNo: msg.seqNo,
 											offset: msg.offset,
-											payload: this.#options.decode!(msg.data),
+											payload: this.#options.decode?.(data) || data as Payload,
+											uncompressedSize: msg.uncompressedSize,
 											createdAt: msg.createdAt ? timestampDate(msg.createdAt) : undefined,
 											writtenAt: batch.writtenAt ? timestampDate(batch.writtenAt) : undefined,
 											metadataItems: msg.metadataItems ? Object.fromEntries(msg.metadataItems.map(item => [item.key, item.value])) : undefined,
@@ -940,10 +983,14 @@ export class TopicReader<Payload = Uint8Array> implements Disposable {
 	 * This method should be called when the reader is no longer needed to prevent memory leaks.
 	 */
 	dispose() {
-		dbg('disposing TopicReader for consumer %s', this.#options.consumer)
+		if (this.#disposed) {
+			return; // Already disposed, nothing to do
+		}
+		this.#disposed = true;
+		dbg('disposing TopicReader for consumer %s', this.#options.consumer);
 
-		this.#disposed = true
 		this.#buffer.length = 0 // Clear the buffer to release memory
+		this.#freeBufferSize = this.#maxBufferSize; // Reset free buffer size to max buffer size
 
 		for (let partitionSession of this.#partitionSessions.values()) {
 			// Stop all partition sessions gracefully
@@ -963,11 +1010,11 @@ export class TopicReader<Payload = Uint8Array> implements Disposable {
 
 		this.#pendingCommits.clear() // Clear pending commits to release memory
 
-		this.#controller.abort()
 		this.#fromClientEmitter.removeAllListeners()
 		this.#fromServerEmitter.removeAllListeners()
 
 		clearInterval(this.#updateTokenTicker)
+		this.#controller.abort()
 	}
 
 	[Symbol.dispose]() {

packages/topic/src/reader.ts

@@ -90,7 +90,6 @@ export type TopicWriterOptions<Payload = Uint8Array> = {
 	// Compression options for the payload.
 	compression?: {
 		codec: Codec,
-		// TODO: Codec options, like compression level, etc.
 		// Minimum raw size to compress, if the payload is smaller than this size, it will not be compressed.
 		// This is useful to avoid compressing small payloads that do not benefit from compression.
 		// Default is 1024 bytes.
@@ -100,7 +99,7 @@ export type TopicWriterOptions<Payload = Uint8Array> = {
 		// If not provided, the default compression function will be used.
 		// The default compression function will use the codec specified in the options.
 		// If the codec is Codec.RAW, the payload will not be compressed.
-		compress?(payload: Uint8Array): Uint8Array
+		compress?(payload: Uint8Array): Uint8Array | Promise<Uint8Array>
 	}
 	// Custom encoding function that can be used to transform the payload before compressing or sending it to the topic.
 	encode?(payload: Payload): Uint8Array
@@ -484,7 +483,7 @@ export class TopicWriter<Payload = Uint8Array> implements Disposable, AsyncDispo
 			if (this.#options.compression.codec !== Codec.RAW && data.length >= (this.#options.compression.minRawSize || MIN_RAW_SIZE)) {
 				// Use custom compression function if provided, otherwise use default compression
 				data = this.#options.compression.compress
-					? this.#options.compression.compress(data)
+					? await this.#options.compression.compress(data)
 					: data; // Default to raw if no compression is applied
 			} else {
 				// If the payload is smaller than the minimum size, do not compress it