Extend `mirrord-protocol` to allow concurrent send/receive

[Razz4780]

EDIT: chunking delayed until we move to a simpler protocol negotiation mechanism. Until then, we settle on concurrent send/receive of whole messages

mirrord-protocol is our bincode-based protocol used for network communication between mirrord processes (e.g. mirrord-intproxy and mirrord-agent). It always utilizes a single TCP connection, and messages are always processed in a first-come first-served manner. This behavior causes issues when mirrord components compete for the connection. E.g. outgoing traffic feature is used to enable communication with remote databases, and often produces a lot of traffic. This traffic can significantly increase latency for other components, e.g. ping pong mechanism, which leads to timeouts and failures.

The idea here is to implement fair sharing of a single mirrord-protocol connection between multiple components, so that messages can be transferred concurrently, and each component has an equal share of total network bandwidth.

Essentially:

1. Each ClientMessage and DaemonMessage variant represents a separate message group.
2. Within each message group, messages are processed sequentially.
3. Groups are processed concurrently.

For instance, this could be implemented as follows (but more clever solutions are always welcome):

1. Messages to send are chunked and added to their group's queue
2. p's receive buffer
3. When group's receive buffNext chunk to send is selected from the head of a random group's queue
4. Received chunks are added to their grouer is ready, complete message is received

When implementing, mind that:

1. mirrord-protocol must remain backwards compatible. We normally use ClientMessage::SwitchProtocolVersion request to negotiate mirrord-protocol version with the peer. The negotiated version determines which messages can be sent over the wire.
2. This extension should not introduce any significant time or memory overhead.

[linear]

MBE-1418 Extend `mirrord-protocol` to allow concurrent send/receive

[majorz]

Hi @Razz4780, I wanted to provide a quick update on my progress regarding this issue. I made a proof-of-concept repository that reproduces the behavior described above. It implements a simple TCP server and client using Tokio, sending multiple large TcpOutgoing messages with artificial server-side delays to simulate heavy processing. I am measuring Ping/Pong latency to observe how message queuing and network delays affect responsiveness. In addition tc could be used to simulate further network latency and packet delays.

The setup allows me to experiment easier with the solutions. I am already making good progress implementing the algorithms described in the original description and next I will follow up with a new proof-of-concept that integrates those improvements.

[majorz]

Hi @Razz4780, I extended my proof-of-concept code with all the components and algorithms required for this task. I placed all of them in a chunk.rs module.

Chunk represents a single piece of a larger encoded message. It contains a group_id identifier corresponding to the enum variant it is part of, as well as index and total fields to indicate its position within the sequence and the total number of chunks for the message. The actual data is stored in the payload: Payload member.

DaemonMessage and ClientMessage have been extended with a Chunk variant - DaemonMessage::Chunk(Chunk) and ClientMessage::Chunk(Chunk). All existing enum variants can now be serialized into a series of Chunks. For example, a large DaemonMessage::TcpOutgoing(DaemonTcpOutgoing) will be split into multiple DaemonMessage::Chunk(Chunk) instances. The original variant can then be reconstructed from these chunks once all of them have been received.

FragmentWriter<const FRAGMENT_SIZE: usize> serializes values implementing the bincode::Encode trait into a VecDeque<Bytes>. It implements bincode::enc::write::Writer, allowing it to be used with bincode::encode_into_writer to minimize intermediate allocations. The algorithm uses slice arithmetic rather than indexing calculations, which makes the code cleaner and more idiomatic. This component does not produce Chunks directly, so it can also be used in other scenarios where generic fragmentation is needed.

Chunk serialization and deserialization are handled by ChunkProducer and ChunkReader.

ChunkProducer<const FRAGMENT_SIZE: usize> uses FragmentWriter to generate a sequence of Chunk instances. It is an iterator, which minimizes allocations and provides additional flexibility.

ChunkReader allows sequential reading from multiple Chunks as if they were a single contiguous byte stream. It implements bincode::de::read::Reader, enabling bincode to decode data split across multiple Chunks.

While the original description suggested using a group receive buffer, I decided to gather all chunks in a vector and decode them directly with ChunkReader. Both approaches are valid and good.

Group handling is performed by GroupChunkQueue (on the sender side) and GroupChunkCollector (on the receiver side).

GroupChunkQueue manages multiple groups of chunks, each identified by a group_id. Chunks within a group maintain their order, but when popping, the queue selects a group at random and returns its next chunk. This preserves the order of individual messages while interleaving chunks from multiple groups during transmission. Each group has an equal share of total network bandwidth.

GroupChunkCollector stores incoming chunks by group_id until a complete message is received. Once the last chunk arrives, it returns the full set of chunks, ready for reassembly.

A potential future optimization is to use separate thread(s) for encoding messages into chunks. This could further reduce allocations, as each chunk could be sent immediately after encoding, avoiding double allocation.

All algorithms and components in chunk.rs are generic and do not directly reference DaemonMessage or ClientMessage and can be reused in other contexts.

Next I will explore how to integrate these components and algorithms into mirrord-protocol. No tests have been implemented yet. That will follow later. I will also need to do another round of careful error-handling exploration to properly cover potential error scenarios. I am open to questions and suggestions for improvement.

[Razz4780]

Hi @majorz!

I checked out the repo, and it's looking good so far ^^
I have two suggestions. Let me know what you think.

Chunk definition

In the repo, you have:

#[derive(Encode, Decode, Debug, PartialEq, Eq, Clone)]
pub struct Chunk {
    /// Identifier for the group of chunks that belong to the same message.
    pub group_id: u8,

    /// Index of this chunk within the group.
    pub index: u64,

    /// Total number of chunks in the group.
    pub total: u64,

    /// Data payload contained in this chunk.
    pub payload: Payload,
}

I'd rather represent Chunk like this:

pub enum Chunk<'a> {
    Start {
        message_id: u16,
        /// Total length of data within this message.
        total_length: u64,
        data: Cow<'a, [u8>,
    },
    Data {
        message_id: u16,
        data: Cow<'a, [u8]>,
    },
}

impl<'a> Encode for Chunk<'a> { ... }

impl Decode for Chunk<'static> { ... }

1. This would give smaller overhead on transmitted bytes: 11 bytes for first chunks, 3 bytes for all other chunks (representing the enum discriminant will take 1 byte). There should be no need to transmit the chunk index and the total count of chunks within each message. With this representation, the receiver should be able to reconstruct the original message just by keeping track of the length of already received data.
2. Using u16 for message id would cap us at 64k concurrently streamed messages, which should be way more than enough (even if we later introduce some kind of concurrency within the groups). Since we use TCP as the transport protocol, we can reuse the message id after fully transmitting the previous message.
3. With data stored in Cow<'a, [u8]>, the sender wrapper would not have to allocate multiple Vecs when chunking a message. It could simply encode it all to one Vec, and keep track of next chunk's index. This should give us fewer allocations on the sender side.

Usage with mirrord-protocol

I wouldn't add the chunks to the ClientMessage/DaemonMessage enums. We've already managed to pollute those pretty heavily 😆
Instead, I'd do sth like:

enum ClientMessage {
    ... // previous variants
    /// Client demands switching the protocol.
    /// After receiving [`DaemonMessage::ProtocolSwitched`] the client can recover the raw IO stream and apply the protocol-specific codec.
    /// After sending this message, the client must suspend sending messages until [`DaemonMessage::ProtocolSwitched`] is received.
    SwitchProtocol(ProtocolFlavor),
}

enum ProtocolFlavor {
    Chunked,
}

enum DaemonMessage {
    ... // previous variants
    /// Informs the client that the protocol has been switched according to the [`ClientMessage::SwitchProtocol`] request.
    /// After sending this message, the daemon can recover the raw IO stream and apply the protocol-specific codec.
    ProtocolSwitched,
}

[majorz]

Hi @Razz4780 , I started implementing the new Chunk definition in a dedicated chunk-cow branch: chunk.rs. I like how the new structure is much more compact. The algorithms in this branch now correctly assemble and disassemble the new Chunk variants, Chunk::Start and Chunk::Data.

Regarding minimizing allocations I have an idea that could also simplify the Chunk definition. Since all chunks have a fixed maximum size we can manage them using a pool of reusable data fragments of size FRAGMENT_SIZE. When a new chunk is created during encoding, its data buffer can be taken from the pool and after the chunk is transmitted its buffer is returned. This approach can eliminate new allocations entirely during message serialization.

Using such a pool may also allows us to preallocate a large buffer upfront, which can be split into many reusable fragments. This strategy would improve cache locality as well. The bytes crate provides all the necessary functionality to implement this, making it a better choice in my opinion than Cow<'a, [u8]> or Vec<u8>.

The bytes crate can also support the approach you suggested - encoding whole buffer at once and then creating Chunks that reference slices of it via Bytes::slice. I believe Bytes is a cleaner and also more flexible solution as it provides exactly the functionality we need out of the box.

As a potential future optimization if serialization is done in separate thread(s), we would not need to wait for the entire value to be encoded before starting transmission - chunks could be sent incrementally as they are generated.

I think the Chunk structure could adopt the new Chunk::Start and Chunk::Data variants to reduce its footprint, while continue using the bytes crate for buffer management. Please let me know what is your opinion on this?

[Razz4780]

Hey!

Re chunks buffer: my concern is that bincode crate is inherently sync. It does not work with async IO types, and we use async TCP stream to transmit the messages. Therefore, we can't do any lazy encoding - we can't block the async runtime thread. I'm not entirely sure what is suggested, so let me provide a snippet with what I had in mind:

pub enum Chunk<'a> {
    Start {
        id: u16,
        total_length: u64,
        data: Cow<'a, [u8]>,
    },
    Data {
        id: u16,
        data: Cow<'a, [u8]>,
    },
}

pub struct MessageChunks {
    encoded: Vec<u8>,
    next_chunk_offset: usize,
    id: u16,
}

impl MessageChunks {
    const CHUNK_SIZE: usize = 4096;

    pub fn new<M: Encode>(id: u16, message: &M) -> Result<Self, EncodeError> {
        let encoded = encode_to_vec(message, bincode::config::standard())?;
        Ok(Self {
            encoded,
            next_chunk_offset: 0,
            id,
        })
    }

    pub fn next(&mut self) -> Option<Chunk<'_>> {
        if self.next_chunk_offset == self.encoded.len() {
            return None;
        }

        let is_first = self.next_chunk_offset == 0;
        let data = self
            .encoded
            .get(self.next_chunk_offset..self.next_chunk_offset + Self::CHUNK_SIZE)
            .unwrap_or_else(|| self.encoded.get(self.next_chunk_offset..).unwrap());
        self.next_chunk_offset += data.len();

        if is_first {
            Some(Chunk::Start {
                id: self.id,
                total_length: self.encoded.len().try_into().unwrap(),
                data: Cow::Borrowed(data),
            })
        } else {
            Some(Chunk::Data {
                id: self.id,
                data: Cow::Borrowed(data),
            })
        }
    }
}

Also, now that I think of it, we don't need Encode + Decode on Chunk. Using bincode for encoding/decoding Chunks would mean extra temporary vectors (since we can't directly encode/decode to/from the IO stream). But, since the Chunk itself is a very simple message, we can encode/decode it manually.

Reconstructing messages would look like this:

1. Read three bytes for enum discriminant and message id
2. If it's Start, read message length and allocate a buffer.
3. Read data chunk len and data
4. Push data to the buffer
5. If the buffer is filled, decode with bincode

Let me know what you think ^^
I'm open to other solutions, but please mind that this protocol will be used in async context

[majorz]

Hey, thank you for providing the example. Sorry that I was not too clear with my ideas - I have thrown too many in a single message. Let me clarify the main one first - in my opinion it is better to use Bytes instead of Cow<'a, [u8]> for the data field in the Chunk enum.

Bytes may also act in a dualistic manner as Cow - it can either own the data or reference it. But it does so in a way that avoids lifetime issues and self-referential structs. It also provides extra functionality that can be useful for us.

Your example code actually would illustrate where we will run into issues when we try to integrate it into a utility like GroupChunkQueue. The chunks are tied to the encoded Vec<u8> and when we attempt to add both of those as members in GroupChunkQueue we will run into a problem that Rust won't allow us to create a self-referential struct.

Now let me illustrate why we need to keep both the Vec<u8> and the Chunk in the GroupChunkQueue. As you can see from the sender_task function below, we are continuously receiving DaemonMessages, splitting them into chunks and sending them randomly. The important detail here is that we do not send all the chunks immediately, but we allow for more messages to be received and split into chunks, while we do the sending. This means that we need to keep the original Vec<u8> around until all chunks have been sent. That exaplains why it has to be part of the GroupChunkQueue.

/// Continuously receives `DaemonMessage`s, splits into chunks, and sends them randomly
async fn sender_task(
    mut rx: mpsc::Receiver<DaemonMessage>,
    mut sink: SplitSink<Framed<TcpStream, DaemonCodec>, DaemonMessage>,
) {
    let notify = Arc::new(Notify::new());
    let mut group_queue = GroupChunkQueue::new();

    loop {
        tokio::select! {
            // Prioritize receiving new messages
            msg = rx.recv() => {
                match msg {
                    Some(msg) => {
                        let chunk_producer = ChunkProducer::<CHUNK_SIZE>::new(&msg).unwrap();

                        drop(msg); // Drop early the original message to free memory

                        for chunk in chunk_producer {
                            group_queue.push_chunk(chunk);
                        }

                        notify.notify_one();
                    }
                    None => break,
                }
            }

            // Wait for a chunk to be ready and send it
            _ = notify.notified(), if !group_queue.is_empty() => {
                if let Some(chunk) = group_queue.pop_random_chunk()
                    && let Err(e) = sink.send(DaemonMessage::Chunk(chunk)).await {
                    eprintln!("Error sending chunk: {:?}", e);
                    break;
                }
            }
        }
    }
}

The classical way to solve this is to actually use Arc<Vec<u8>> instead of Vec<u8> and borrow from that instead. This way the self-referential struct issue will be avoided. But then we are essentially doing what Bytes already does for us - reference counting and cheap cloning. So my suggestion is just to use Bytes directly.

This way we will avoid also attaching the lifetime parameter 'a to the Chunk enum and all the structs that use it, which will make the code cleaner and easier to use.

I may explain the other ideas I had separately as I do not want to saturate again the discussion about whether we need to use Bytes or Cow<'a, [u8]> for the data field in the Chunk enum. I believe that using Bytes is the better choice here. But maybe I am missing something, so please let me know your thoughts.

[Razz4780]

Ah, ok. Indeed Bytes can be a replacement for Cow ^^

[majorz]

I covered the ideas from the "Chunk definition" section in the implementation code, including the additional allocation optimization you suggested. An interesting aspect is that we first encode into a BytesMut and then, when generating a series of Bytes for the Chunks, we use self.buffer.split_to(chunk_size).freeze(). This splits the already encoded BytesMut into two parts (without extra allocation) and then freezes the produced part into a Bytes. The resulting Bytes instances own the data, yet remain contiguous, and there are no additional allocations. Turned out an elegant solution. I started already looking into the ideas in the "Usage with mirrord-protocol" and the new SwitchProtocol(ProtocolFlavor) mechanism.