pipeworks_core/

channel.rs

use std::{
    any::TypeId,
    sync::{Arc, RwLock},
    time::SystemTime,
};

use tokio::sync::{broadcast, mpsc};
use tracing::{error, warn};

use crate::{
    event::{AnyBusEvent, BusEvent},
    node_id::NODE_ID,
    reg::{BusType, TypeReg},
};

/// A type-erased broadcast channel for a single type, with optional support for out-of-band
/// (de)serialization to/from another internal broadcast channel.
pub struct Channel {
    /// The type registration used to create this channel.
    pub type_reg: TypeReg,

    /// The broadcast sender for local type-erased events. This bus will optionally feed the
    /// `serde` bus if serialization is enabled.
    local: broadcast::Sender<AnyBusEvent>,

    /// The serialized bus.
    /// - Remote events pushed to this broadcast will result in the message being deserialized and
    ///   pushed to the `local` broadcast.
    /// - Non-remote events pushed to `local` will be (lazily) serialized and pushed to this
    ///   broadcast.
    serde: broadcast::Sender<BusEvent<SerdePayload>>,

    /// The latest value (if any) received by the channel.
    latest: Arc<RwLock<Option<AnyBusEvent>>>,
}

#[derive(Clone)]
pub enum SerdePayload {
    Bitcode(Vec<u8>),
}

impl Channel {
    pub fn from_type_reg(type_reg: TypeReg) -> Self {
        let (local_tx, mut local_rx) = broadcast::channel::<AnyBusEvent>(type_reg.buffer_cap);
        let (serde_tx, mut serde_rx) =
            broadcast::channel::<BusEvent<SerdePayload>>(type_reg.buffer_cap);
        let latest: Arc<RwLock<Option<AnyBusEvent>>> = Arc::new(RwLock::new(None));

        if let Some((to_bytes, from_bytes)) = type_reg.bitcode_support {
            // Spawn the local -> serde task
            let serde_tx_clone = serde_tx.clone();
            tokio::spawn(async move {
                loop {
                    let Ok(AnyBusEvent { source, time, msg }) = local_rx.recv().await else {
                        warn!(
                            "TODO(Send to CAS) Serde serialize was unable to keep up with local bus ingres {}",
                            type_reg.type_name
                        );
                        continue;
                    };

                    // Never serialize remote events, or events when there is no receiver
                    if !source.is_me() || serde_tx_clone.receiver_count() == 0 {
                        continue;
                    }

                    let bytes = to_bytes(msg);

                    // Ignore errors (it just means there is no receiver, which can happen if it was
                    // dropped after the short-circuit check above).
                    let _ = serde_tx_clone.send(BusEvent {
                        source,
                        time,
                        msg: SerdePayload::Bitcode(bytes),
                    });
                }
            });

            // Spawn the serde -> local task
            let local_tx_clone = local_tx.clone();
            let latest_clone = latest.clone();
            tokio::spawn(async move {
                loop {
                    let Ok(BusEvent { source, time, msg }) = serde_rx.recv().await else {
                        warn!(
                            "TODO(Send to CAS) Serde deserialize was unable to keep up with serde ingres {}",
                            type_reg.type_name
                        );
                        continue;
                    };
                    // Never deserialize local events
                    if source.is_me() {
                        continue;
                    }

                    let msg = match msg {
                        SerdePayload::Bitcode(bytes) => {
                            match from_bytes(&bytes) {
                                Ok(v) => v,
                                Err(err) => {
                                    // TODO: Send to some kind of a CAS
                                    error!(
                                        "TODO(push to CAS): Failed to deserialize bytes into bus channel type {}: {}",
                                        type_reg.type_name, err
                                    );
                                    continue;
                                }
                            }
                        }
                    };

                    let event = AnyBusEvent { source, time, msg };
                    *latest_clone.write().unwrap() = Some(event.clone());

                    // Ignore errors. It just means no one is listening.
                    let _ = local_tx_clone.send(event);
                }
            });
        }

        Self {
            type_reg,
            local: local_tx,
            serde: serde_tx,
            latest,
        }
    }

    /// Panics if type T is not the same as the underlying erased type.
    pub fn send_event<T: BusType>(&self, event: BusEvent<T>) {
        assert!(TypeId::of::<T>() == self.type_reg.type_id);
        let any_event = event.type_erase();
        *self.latest.write().unwrap() = Some(any_event.clone());
        let _ = self.local.send(any_event);
    }

    pub fn send_serde_event(&self, event: BusEvent<SerdePayload>) {
        let _ = self.serde.send(event);
    }

    pub fn update_latest<T: Default + BusType>(&self, f: impl FnOnce(T) -> T) {
        assert!(TypeId::of::<T>() == self.type_reg.type_id);

        let mut latest_guard = self.latest.write().unwrap();
        let value = latest_guard
            .as_ref()
            .map(|e| e.downcast_cloned::<T>().msg)
            .unwrap_or_default();

        let event = BusEvent {
            source: *NODE_ID,
            time: SystemTime::now(),
            msg: f(value),
        };

        let any_event = event.type_erase();
        *latest_guard = Some(any_event.clone());
        let _ = self.local.send(any_event);
    }

    pub fn subscribe<T: BusType>(&self, prefix_latest: bool) -> mpsc::Receiver<BusEvent<T>> {
        assert!(TypeId::of::<T>() == self.type_reg.type_id);
        let mut any_event_rx = self.local.subscribe();
        let type_reg = self.type_reg.clone();

        let (tx, rx) = mpsc::channel(1);

        // This can safely be queried now, as we already started the subscription to the broadcast
        // above.
        let push_latest = if prefix_latest {
            self.get_latest::<T>()
        } else {
            None
        };

        tokio::spawn(async move {
            if let Some(latest) = push_latest {
                if let Err(_) = tx.send(latest).await {
                    // Receiver was dropped.
                    return;
                }
            }

            loop {
                match any_event_rx.recv().await {
                    Ok(any_event) => {
                        if let Err(_) = tx.send(any_event.downcast_cloned::<T>()).await {
                            // Receiver was dropped.
                            break;
                        }
                    }
                    Err(err) => {
                        match err {
                            // The whole channel was dropped.
                            broadcast::error::RecvError::Closed => break,
                            // The consumer could not keep up.
                            broadcast::error::RecvError::Lagged(count) => {
                                warn!(
                                    "TODO(push to CAS) Channel receiver {} lagged {} events",
                                    type_reg.type_name, count
                                );
                            }
                        }
                    }
                }
            }
        });

        rx
    }

    pub fn subscribe_serde(&self) -> broadcast::Receiver<BusEvent<SerdePayload>> {
        self.serde.subscribe()
    }

    pub fn get_latest<T: BusType>(&self) -> Option<BusEvent<T>> {
        self.latest
            .read()
            .unwrap()
            .as_ref()
            .map(|any_event| any_event.downcast_cloned::<T>())
    }
}