Improving efficiency of encoding ProduceRequests

[eblocha]

Hello!

I'm wondering if there's any way to avoid a double-copy I'm having to do when encoding produce requests. Currently I am creating a new BytesMut to pass to RecordBatchEncoder::encode, then passing that into PartitionProduceData. At the tcp sink, those bytes then get encoded again into the sink's internal buffer.

let mut records = BytesMut::new();

RecordBatchEncoder::encode(
    &mut records,
    prepared_records.iter(), // Vec<Record>
    ...
);

// error handling...

// this gets added to the ProduceRequest
let partition_data = PartitionProduceData::default()
    .with_index(partition)
    .with_records(Some(records.into()));

However, when this request is going to be sent to the IO sink, it gets encoded into a FramedWrite's internal buffer (from tokio_util). After encoding, the BytesMut I made is dropped, and I need to allocate a new one.

I'm wondering if there'd be a way to avoid this initial encode step into the intermediate byte array in PartitionProduceData, and do RecordBatchEncoder::encode on the buffer for FramedWrite, skipping the extra allocation.

[eblocha]

Here's a flamegraph of a profile I took that sparked this question:

When building the produce request, I have to copy the bytes into a buffer to put in the request, then copy from the request into the IO sink buffer.

[rukai]

I agree that is a problem.
I'm not sure how to solve it though.
For my own use case I use kafka-protocol for a kafka proxy and I need the records to remain as raw bytes since it would be expensive and unneeded to parse the records.

Maybe we could store the records as something like:

enum RecordBatch {
   RawBytes(Bytes),
   Parsed(RecordBatchParsed),
}

A while back I had the idea to extend this project to generate 3 separate crates.

• one optimized for clients
• one optimized for kafka reimplementations
• one optimized for proxies

This has lots of benefits to performance and ease of use.
But in this particular case the first two crates in this scenario could avoid the complexities of the above enum solution and just directly store the RecordBatchParsed.
Unfortunately this is a huge undertaking that I cant afford to take on.

[eblocha]

On the decode side, you'd have to pick only one variant to decode into. Perhaps if PartitionProduceData was generic, it could defer the Encodable and Decodable implementations to the generic's implementation. Of course, this would be a breaking change.

Avoiding a breaking change would probably mean creating a whole new type.

[rukai]

we are currently pretty flexible with breaking changes in this project, a bump to 0.15 is not an issue.

You make a good point about decoding.
The simplest solution is to just take the performance hit for decoding in your use case and decode to raw bytes every time.
I would like to avoid generics if possible because of the complexities they bring.
But if you could do it in a way that is not very disruptive it could work.

Another alternative is to have an alternate decode method on Decodable that decodes as the parsed form.

[eblocha]

More complexity with this is that the compact byte format uses a dynamically-sized header to record the size of the message, and we don't know the size of the bytes we're about to write.

I'm not sure if there's a way to reserve enough space for the maximum-size length header, write the records, go back and write the length, then copy_within the records backwards to be in the correct position. That might defeat the whole optimization to begin with, not sure.