fix: Fixed decompression with multiple batches. (#123)

Signed-off-by: Paolo Insogna <[email protected]>

Author: ShogunPanda

src/apis/consumer/fetch-v17.ts

@@ -174,15 +174,14 @@ export function parseResponse (
             preferredReadReplica: r.readInt32()
           }
 
-          let recordsSize = r.readUnsignedVarInt()
-
           if (partition.errorCode !== 0) {
             errors.push([`/responses/${i}/partitions/${j}`, partition.errorCode])
           }
 
-          if (recordsSize > 1) {
-            recordsSize--
+          // We need to reduce the size by one to follow the COMPACT_RECORDS specification
+          const recordsSize = r.readUnsignedVarInt() - 1
 
+          if (recordsSize > 0) {
             const recordsBatchesReader = Reader.from(r.buffer.subarray(r.position, r.position + recordsSize))
             partition.records = []
             do {

src/protocol/records.ts

@@ -7,7 +7,7 @@ import {
   compressionsAlgorithmsByBitmask
 } from './compression.ts'
 import { crc32c } from './crc32c.ts'
-import { type NullableString } from './definitions.ts'
+import { INT32_SIZE, INT64_SIZE, type NullableString } from './definitions.ts'
 import { DynamicBuffer } from './dynamic-buffer.ts'
 import { Reader } from './reader.ts'
 import { Writer } from './writer.ts'
@@ -253,6 +253,7 @@ export function createRecordsBatch (
 }
 
 export function readRecordsBatch (reader: Reader): RecordsBatch {
+  const initialPosition = reader.position
   const batch = {
     firstOffset: reader.readInt64(),
     length: reader.readInt32(),
@@ -279,10 +280,13 @@ export function readRecordsBatch (reader: Reader): RecordsBatch {
       throw new UnsupportedCompressionError(`Unsupported compression algorithm with bitmask ${compression}`)
     }
 
-    const buffer = algorithm.decompressSync(reader.buffer.slice(reader.position, reader.buffer.length))
+    // The length of all headers immediately following Length up to the length of the Records array
+    const headersLength = reader.position - initialPosition - INT32_SIZE - INT64_SIZE
+    const compressedDataLen = batch.length - headersLength
+    const buffer = algorithm.decompressSync(reader.buffer.slice(reader.position, reader.position + compressedDataLen))
 
-    // Move the original reader to the end
-    reader.skip(reader.buffer.length - reader.position)
+    // Move the original reader to the end of the compressed data
+    reader.skip(compressedDataLen)
 
     // Replace the reader with the decompressed buffer
     reader = Reader.from(buffer)

test/clients/consumer/consumer.test.ts

@@ -1262,16 +1262,29 @@ test('fetch should retrieve messages from multiple batches', async t => {
   }
 })
 
-for (const compression of Object.values(CompressionAlgorithms)) {
-  test(`fetch should retrieve messages from multiple batches (compressed ${compression})`,
-    { skip: compression === 'snappy' || (compression === 'zstd' && !('zstdCompressSync' in zlib)) }, async t => {
+for (const [name, compression] of Object.entries(CompressionAlgorithms)) {
+  test(
+    `fetch should retrieve messages from multiple batches (compressed with CompressionAlgorithms.${name})`,
+    { skip: compression === 'zstd' && !('zstdCompressSync' in zlib) },
+    async t => {
       const topic = await createTopic(t, true)
       const producer = await createProducer(t)
 
-      const msg: MessageToProduce = { key: Buffer.from('test'), value: Buffer.from('test'), topic }
-      await producer.send({ acks: ProduceAcks.NO_RESPONSE, compression, messages: [msg] })
-      await producer.send({ acks: ProduceAcks.NO_RESPONSE, compression, messages: [msg] })
-      await producer.send({ acks: ProduceAcks.NO_RESPONSE, compression, messages: [msg] })
+      await producer.send({
+        acks: ProduceAcks.NO_RESPONSE,
+        compression,
+        messages: [{ key: Buffer.from('test1'), value: Buffer.from('test1'), topic }]
+      })
+      await producer.send({
+        acks: ProduceAcks.NO_RESPONSE,
+        compression,
+        messages: [{ key: Buffer.from('test2'), value: Buffer.from('test2'), topic }]
+      })
+      await producer.send({
+        acks: ProduceAcks.NO_RESPONSE,
+        compression,
+        messages: [{ key: Buffer.from('test3'), value: Buffer.from('test3'), topic }]
+      })
 
       const consumer = createConsumer(t, {})
 
@@ -1287,15 +1300,18 @@ for (const compression of Object.values(CompressionAlgorithms)) {
 
       let i = 0
       for await (const message of stream) {
+        i++
+
         strictEqual(message.topic, topic)
-        strictEqual(message.key.toString(), 'test')
-        strictEqual(message.value.toString(), 'test')
+        strictEqual(message.key.toString(), 'test' + i)
+        strictEqual(message.value.toString(), 'test' + i)
 
-        if (++i === 3) {
+        if (i === 3) {
           break
         }
       }
-    })
+    }
+  )
 }
 
 test('commit should commit offsets to Kafka and support diagnostic channels', async t => {