tidaldb/tidal/src/replication/in_process.rs

//! In-process transport for WAL segment shipping via crossbeam channels.
//!
//! Used for integration tests and single-process multi-shard deployments.
//! Each shard gets a unique receiver and a clone of all senders, enabling
//! any-to-any communication within the process.

use std::collections::HashMap;

use crossbeam::channel::{Receiver, Sender, bounded};

use crate::replication::{
    shard::ShardId,
    transport::{Transport, TransportError, WalSegmentPayload},
};

/// Maximum payload size in bytes (64 MB). Segments larger than this are
/// rejected to prevent unbounded memory growth in the channel buffer.
const MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024;

/// Per-shard channel wiring: a map from each local shard to its
/// `(all_senders, own_receiver)` pair. One entry per shard; the sender map
/// inside each pair carries a clone of every shard's sender so an endpoint can
/// ship to any peer.
pub(crate) type ChannelEndpoints<P> = HashMap<ShardId, (HashMap<ShardId, Sender<P>>, Receiver<P>)>;

/// Build the channel wiring for a set of any-to-any in-process transport
/// endpoints: one bounded channel per shard, where every endpoint owns a clone
/// of *all* senders (so it can ship to any peer) and its own unique receiver.
///
/// Returns, per shard, the `(all_senders, own_receiver)` pair from which a
/// concrete endpoint struct is built. Shared by
/// [`InProcessTransportFactory::build`] (WAL segments) and
/// [`super::session_bridge::InProcessSessionTransportFactory::build`] (session
/// payloads) so the wiring lives in one place; each factory only differs in its
/// payload type `P` and the endpoint struct it wraps the parts in.
pub(crate) fn build_channel_endpoints<P>(
    shards: &[ShardId],
    capacity: usize,
) -> ChannelEndpoints<P> {
    let mut senders: HashMap<ShardId, Sender<P>> = HashMap::new();
    let mut receivers: HashMap<ShardId, Receiver<P>> = HashMap::new();

    for &shard_id in shards {
        let (tx, rx) = bounded(capacity);
        senders.insert(shard_id, tx);
        receivers.insert(shard_id, rx);
    }

    shards
        .iter()
        .map(|&local_shard| {
            // Each receiver is consumed exactly once (the map was populated from
            // the same shard list), so `remove` is always `Some`.
            let receiver = receivers
                .remove(&local_shard)
                .expect("receiver missing for shard that was just created");
            (local_shard, (senders.clone(), receiver))
        })
        .collect()
}

/// Factory for building a set of [`InProcessTransport`] endpoints, one per shard.
///
/// Each transport owns all senders (one per shard) and its own unique receiver.
/// The channel capacity is 64 segments per shard.
pub struct InProcessTransportFactory {
    shards: Vec<ShardId>,
}

impl InProcessTransportFactory {
    /// Create a factory for the given shard IDs.
    #[must_use]
    pub fn new(shards: &[ShardId]) -> Self {
        Self {
            shards: shards.to_vec(),
        }
    }

    /// Build one `InProcessTransport` per shard.
    ///
    /// Each transport can send to any shard and receives on its own channel.
    ///
    /// # Panics
    ///
    /// Panics if a receiver is missing for a shard -- this is unreachable
    /// because the receiver map is populated from the same shard list.
    #[must_use]
    pub fn build(self) -> HashMap<ShardId, InProcessTransport> {
        build_channel_endpoints::<WalSegmentPayload>(&self.shards, 64)
            .into_iter()
            .map(|(local_shard, (senders, receiver))| {
                (
                    local_shard,
                    InProcessTransport {
                        local_shard,
                        senders,
                        receiver,
                    },
                )
            })
            .collect()
    }
}

/// An in-process transport endpoint backed by crossbeam bounded channels.
///
/// Thread-safe: `Sender` and `Receiver` from crossbeam are `Send + Sync`.
pub struct InProcessTransport {
    local_shard: ShardId,
    senders: HashMap<ShardId, Sender<WalSegmentPayload>>,
    receiver: Receiver<WalSegmentPayload>,
}

impl Transport for InProcessTransport {
    fn send_segment(&self, to: ShardId, payload: WalSegmentPayload) -> Result<(), TransportError> {
        if payload.bytes.len() > MAX_PAYLOAD_BYTES {
            return Err(TransportError::PayloadTooLarge {
                size: payload.bytes.len(),
                max: MAX_PAYLOAD_BYTES,
            });
        }
        let sender = self
            .senders
            .get(&to)
            .ok_or(TransportError::UnknownPeer(to))?;
        sender.send(payload).map_err(|_| TransportError::Closed)
    }

    fn recv_segment(&self) -> Option<WalSegmentPayload> {
        self.receiver.recv().ok()
    }

    fn local_shard(&self) -> ShardId {
        self.local_shard
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::replication::{segment_id::WalSegmentId, shard::RegionId};

    fn make_payload(shard: ShardId, seqno: u64) -> WalSegmentPayload {
        WalSegmentPayload {
            id: WalSegmentId::new(RegionId::SINGLE, shard, seqno),
            bytes: vec![0xAB; 100],
            event_count: 5,
            leader_last_seq: seqno,
        }
    }

    #[test]
    fn factory_creates_one_transport_per_shard() {
        let shards = [ShardId(0), ShardId(1), ShardId(2)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();
        assert_eq!(transports.len(), 3);
        assert!(transports.contains_key(&ShardId(0)));
        assert!(transports.contains_key(&ShardId(1)));
        assert!(transports.contains_key(&ShardId(2)));
    }

    #[test]
    fn send_and_receive_between_shards() {
        let shards = [ShardId(0), ShardId(1)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();

        let t0 = &transports[&ShardId(0)];
        let t1 = &transports[&ShardId(1)];

        // Shard 0 sends to Shard 1.
        let payload = make_payload(ShardId(0), 42);
        t0.send_segment(ShardId(1), payload).unwrap();

        // Shard 1 receives it.
        let received = t1.recv_segment().unwrap();
        assert_eq!(received.id.seqno, 42);
        assert_eq!(received.event_count, 5);
        assert_eq!(received.bytes.len(), 100);
    }

    #[test]
    fn send_to_unknown_peer_fails() {
        let shards = [ShardId(0)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();

        let t0 = &transports[&ShardId(0)];
        let result = t0.send_segment(ShardId(99), make_payload(ShardId(0), 1));
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(matches!(err, TransportError::UnknownPeer(ShardId(99))));
    }

    #[test]
    fn payload_too_large_rejected() {
        let shards = [ShardId(0), ShardId(1)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();

        let t0 = &transports[&ShardId(0)];
        let payload = WalSegmentPayload {
            id: WalSegmentId::single_node(1),
            bytes: vec![0u8; MAX_PAYLOAD_BYTES + 1],
            event_count: 0,
            leader_last_seq: 0,
        };
        let result = t0.send_segment(ShardId(1), payload);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            TransportError::PayloadTooLarge { .. }
        ));
    }

    #[test]
    fn local_shard_returns_correct_id() {
        let shards = [ShardId(7), ShardId(8)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();

        assert_eq!(transports[&ShardId(7)].local_shard(), ShardId(7));
        assert_eq!(transports[&ShardId(8)].local_shard(), ShardId(8));
    }

    #[test]
    fn recv_returns_none_when_all_transports_dropped() {
        let shards = [ShardId(0), ShardId(1)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();

        // Extract shard 1's receiver before dropping everything.
        // We need to drop ALL transports (including t1 itself) because
        // each transport holds clones of all senders -- including senders
        // to its own channel.
        let t1_receiver = {
            let t1 = &transports[&ShardId(1)];
            t1.receiver.clone()
        };

        // Drop all transports -- this drops all senders.
        drop(transports);

        // With all senders dropped, recv should return None.
        assert!(t1_receiver.recv().is_err());
    }

    #[test]
    fn multiple_segments_fifo_order() {
        let shards = [ShardId(0), ShardId(1)];
        let factory = InProcessTransportFactory::new(&shards);
        let transports = factory.build();

        let t0 = &transports[&ShardId(0)];
        let t1 = &transports[&ShardId(1)];

        for seq in 1..=5u64 {
            t0.send_segment(ShardId(1), make_payload(ShardId(0), seq))
                .unwrap();
        }

        for expected_seq in 1..=5u64 {
            let received = t1.recv_segment().unwrap();
            assert_eq!(received.id.seqno, expected_seq);
        }
    }
}