test: add coverage for large magic-v2 batch encoding

Author: zmstone

scripts/produce_req_encoding_benchmark.sh

@@ -104,4 +104,4 @@ main() {
 }
 
 # Run main function
-main "$@"
+main "$@"

src/kpro_batch.erl

@@ -37,6 +37,7 @@
 -type headers() :: kpro:headers().
 -type batch_meta() :: kpro:batch_meta().
 -define(NO_META, ?KPRO_NO_BATCH_META).
+-define(CRC32C_NON_DIRTY_BYTES_THRESHOLD, 1024 * 1024).
 
 %% @doc Encode a list of batch inputs into byte stream.
 -spec encode(magic(), batch_input(), compress_option()) -> {non_neg_integer(), iodata()}.
@@ -112,7 +113,7 @@ encode_tx([FirstMsg | _] = Batch, Compression, FirstSequence,
 %% the cost of crc32c computation is negligible comparing to the batch
 %% encoding part, so we do not call dirty-scheduler.
 %% Otherwise we respect iolist and call dirty-scheduler (_d flavor API).
-crc32c(Bytes, IoData) when Bytes =< 1024 * 1024 ->
+crc32c(Bytes, IoData) when Bytes =< ?CRC32C_NON_DIRTY_BYTES_THRESHOLD ->
     crc32cer:nif(IoData);
 crc32c(_, IoData) ->
     crc32cer:nif_iolist_d(IoData).

test/kpro_batch_tests.erl

@@ -46,6 +46,88 @@ encode_decode_test_() ->
    CompressionOpt <- CompressionOpts,
    MagicV <- MagicVersions].
 
+encode_large_batch_test() ->
+  %% Test encoding large batches (>1MB) with different compression options
+  LargeBatch = make_large_batch(),
+
+  %% Test with no compression
+  {SizeNoComp, EncodedNoComp} = kpro_batch:encode(2, LargeBatch, no_compression),
+  ?assert(SizeNoComp > 1024 * 1024), %% Should be >1MB
+  ?assert(is_binary(bin(EncodedNoComp))),
+
+  %% Test with gzip compression
+  {SizeGzip, EncodedGzip} = kpro_batch:encode(2, LargeBatch, gzip),
+  %% Compression may not always be smaller due to overhead, so just verify it's reasonable
+  ?assert(SizeGzip > 0),
+  ?assert(SizeGzip < SizeNoComp * 2), %% Should not be more than 2x larger
+  ?assert(is_binary(bin(EncodedGzip))),
+
+  %% Test with snappy compression
+  {SizeSnappy, EncodedSnappy} = kpro_batch:encode(2, LargeBatch, snappy),
+  %% Compression may not always be smaller due to overhead, so just verify it's reasonable
+  ?assert(SizeSnappy > 0),
+  ?assert(SizeSnappy < SizeNoComp * 2), %% Should not be more than 2x larger
+  ?assert(is_binary(bin(EncodedSnappy))),
+
+  %% Test with lz4 compression (if available)
+  try
+    {SizeLz4, EncodedLz4} = kpro_batch:encode(2, LargeBatch, lz4),
+    ?assert(SizeLz4 > 0),
+    ?assert(SizeLz4 < SizeNoComp * 2), %% Should not be more than 2x larger
+    ?assert(is_binary(bin(EncodedLz4)))
+  catch
+    error:undef ->
+      %% LZ4 not available, skip test
+      ok
+  end,
+
+  %% Test with zstd compression (if available)
+  try
+    {SizeZstd, EncodedZstd} = kpro_batch:encode(2, LargeBatch, zstd),
+    ?assert(SizeZstd > 0),
+    ?assert(SizeZstd < SizeNoComp * 2), %% Should not be more than 2x larger
+    ?assert(is_binary(bin(EncodedZstd)))
+  catch
+    error:undef ->
+      %% ZSTD not available, skip test
+      ok
+  end,
+
+  %% Verify we can decode the uncompressed batch
+  [{_Meta, DecodedMessages}] = kpro_batch:decode(bin(EncodedNoComp)),
+  ?assertEqual(length(LargeBatch), length(DecodedMessages)),
+
+  %% Verify message content matches
+  lists:foreach(fun({Original, Decoded}) ->
+    #kafka_message{key = DecodedKey, value = DecodedValue, ts = DecodedTs} = Decoded,
+    #{key := OriginalKey, value := OriginalValue, ts := OriginalTs} = Original,
+    ?assertEqual(OriginalKey, DecodedKey),
+    ?assertEqual(OriginalValue, DecodedValue),
+    ?assertEqual(OriginalTs, DecodedTs)
+  end, lists:zip(LargeBatch, DecodedMessages)),
+
+  %% Test performance: encoding should complete in reasonable time
+  StartTime = erlang:system_time(microsecond),
+  _ = kpro_batch:encode(2, LargeBatch, no_compression),
+  EndTime = erlang:system_time(microsecond),
+  EncodingTime = EndTime - StartTime,
+  ?assert(EncodingTime < 1000000), %% Should complete in <1 second
+
+  ok.
+
+make_large_batch() ->
+  %% Create a batch that will be >1MB total
+  %% Each message is ~100KB, so 11 messages = ~1.1MB total
+  MessageCount = 11,
+  MessageSize = 100 * 1024, %% 100KB per message
+
+  [#{ ts => kpro_lib:now_ts() + N * 1000 %% Different timestamps
+    , key => <<"large_msg_key_", (integer_to_binary(N))/binary>>
+    , value => crypto:strong_rand_bytes(MessageSize)
+    , headers => [{<<"header_key_", (integer_to_binary(N))/binary>>,
+                   <<"header_value_", (integer_to_binary(N))/binary>>}]
+    } || N <- lists:seq(1, MessageCount)].
+
 provide_compression_test() ->
   kpro:provide_compression([{snappy, ?MODULE}]),
   try