feat(multiplexer): Allow fine-grained control of concurrency strategy (#106)

This commit is contained in:
Santiago Carmuega 2022-06-03 21:37:38 -03:00 committed by GitHub
parent 07fb2ca016
commit e4784f444d
23 changed files with 822 additions and 559 deletions

View file

@ -0,0 +1,103 @@
//! Interface to interact with the multiplexer as an agent
use crate::Payload;
use pallas_codec::{minicbor, Fragment};
#[derive(Debug)]
pub enum ChannelError {
NotConnected(Option<Payload>),
Encoding(String),
Decoding(String),
}
/// A raw link to the ingress / egress of the multiplexer
pub trait Channel {
fn enqueue_chunk(&mut self, chunk: Payload) -> Result<(), ChannelError>;
fn dequeue_chunk(&mut self) -> Result<Payload, ChannelError>;
}
/// Protocol value that defines max segment length
pub const MAX_SEGMENT_PAYLOAD_LENGTH: usize = 65535;
enum Decoding<M> {
Done(M, usize),
NotEnoughData,
UnexpectedError(Box<dyn std::error::Error>),
}
fn try_decode_message<M>(buffer: &[u8]) -> Decoding<M>
where
M: Fragment,
{
let mut decoder = minicbor::Decoder::new(buffer);
let maybe_msg = decoder.decode();
match maybe_msg {
Ok(msg) => Decoding::Done(msg, decoder.position()),
Err(err) if err.is_end_of_input() => Decoding::NotEnoughData,
Err(err) => Decoding::UnexpectedError(Box::new(err)),
}
}
/// A channel abstraction to hide the complexity of partial payloads
pub struct ChannelBuffer<'c, C: Channel> {
channel: &'c mut C,
temp: Vec<u8>,
}
impl<'c, C: Channel> ChannelBuffer<'c, C> {
pub fn new(channel: &'c mut C) -> Self {
Self {
channel,
temp: Vec::new(),
}
}
/// Enqueues a msg as a sequence payload chunks
pub fn send_msg_chunks<M>(&mut self, msg: &M) -> Result<(), ChannelError>
where
M: Fragment,
{
let mut payload = Vec::new();
minicbor::encode(&msg, &mut payload)
.map_err(|err| ChannelError::Encoding(err.to_string()))?;
let chunks = payload.chunks(MAX_SEGMENT_PAYLOAD_LENGTH);
for chunk in chunks {
self.channel.enqueue_chunk(Vec::from(chunk))?;
}
Ok(())
}
/// Reads from the channel until a complete message is found
pub fn recv_full_msg<M>(&mut self) -> Result<M, ChannelError>
where
M: Fragment,
{
// do an eager reading if buffer is empty, no point in going through the error
// handling
if self.temp.is_empty() {
let chunk = self.channel.dequeue_chunk()?;
self.temp.extend(chunk);
}
let decoding = try_decode_message::<M>(&self.temp);
match decoding {
Decoding::Done(msg, pos) => {
self.temp.drain(0..pos);
Ok(msg)
}
Decoding::UnexpectedError(err) => Err(ChannelError::Decoding(err.to_string())),
Decoding::NotEnoughData => {
let chunk = self.channel.dequeue_chunk()?;
self.temp.extend(chunk);
self.recv_full_msg()
}
}
}
}

View file

@ -5,35 +5,62 @@ use std::io::{Read, Write};
use std::os::unix::net::UnixStream;
use std::{net::TcpStream, time::Instant};
use crate::{Bearer, Payload};
use crate::Payload;
pub struct Segment {
pub protocol: u16,
pub timestamp: u32,
pub payload: Payload,
}
pub trait Bearer: Read + Write + Send + Sync + Sized {
type Error: std::error::Error;
fn read_segment(&mut self) -> Result<Option<Segment>, Self::Error>;
fn write_segment(&mut self, segment: Segment) -> Result<(), Self::Error>;
fn clone(&self) -> Self;
}
impl Segment {
pub fn new(clock: Instant, protocol: u16, payload: Payload) -> Self {
Segment {
timestamp: clock.elapsed().as_micros() as u32,
protocol,
payload,
}
}
}
fn write_segment(writer: &mut impl Write, segment: Segment) -> Result<(), std::io::Error> {
let Segment {
timestamp,
protocol,
payload,
} = segment;
fn write_segment(
writer: &mut impl Write,
clock: Instant,
protocol_id: u16,
payload: &[u8],
) -> Result<(), std::io::Error> {
let mut msg = Vec::new();
msg.write_u32::<NetworkEndian>(clock.elapsed().as_micros() as u32)?;
msg.write_u16::<NetworkEndian>(protocol_id)?;
msg.write_u32::<NetworkEndian>(timestamp)?;
msg.write_u16::<NetworkEndian>(protocol)?;
msg.write_u16::<NetworkEndian>(payload.len() as u16)?;
if log_enabled!(log::Level::Trace) {
trace!(
"sending segment, header {:?}, protocol id: {}, payload length: {}",
hex::encode(&msg),
protocol_id,
protocol,
payload.len()
);
}
msg.write_all(payload)?;
msg.write_all(&payload)?;
writer.write_all(&msg)?;
writer.flush()
}
fn read_segment(reader: &mut impl Read) -> Result<(u16, u32, Payload), std::io::Error> {
fn read_segment(reader: &mut impl Read) -> Result<Segment, std::io::Error> {
let mut header = [0u8; 8];
reader.read_exact(&mut header)?;
@ -43,12 +70,12 @@ fn read_segment(reader: &mut impl Read) -> Result<(u16, u32, Payload), std::io::
}
let length = NetworkEndian::read_u16(&header[6..]) as usize;
let id = NetworkEndian::read_u16(&header[4..6]) as usize ^ 0x8000;
let ts = NetworkEndian::read_u32(&header[0..4]);
let protocol = NetworkEndian::read_u16(&header[4..6]) as usize ^ 0x8000;
let timestamp = NetworkEndian::read_u32(&header[0..4]);
debug!(
"parsed inbound msg, protocol id: {}, ts: {}, payload length: {}",
id, ts, length
protocol, timestamp, length
);
let mut payload = vec![0u8; length];
@ -58,44 +85,54 @@ fn read_segment(reader: &mut impl Read) -> Result<(u16, u32, Payload), std::io::
trace!("read segment payload: {:?}", hex::encode(&payload));
}
Ok((id as u16, ts, payload))
Ok(Segment {
protocol: protocol as u16,
timestamp,
payload,
})
}
fn read_segment_with_timeout(reader: &mut impl Read) -> Result<Option<Segment>, std::io::Error> {
match read_segment(reader) {
Ok(s) => Ok(Some(s)),
Err(err) => match err.kind() {
std::io::ErrorKind::WouldBlock => Ok(None),
std::io::ErrorKind::TimedOut => Ok(None),
std::io::ErrorKind::Interrupted => Ok(None),
_ => todo!(),
},
}
}
impl Bearer for TcpStream {
type Error = std::io::Error;
fn clone(&self) -> Self {
self.try_clone().expect("error cloning tcp stream")
}
fn read_segment(&mut self) -> Result<(u16, u32, Payload), std::io::Error> {
read_segment(self)
fn read_segment(&mut self) -> Result<Option<Segment>, std::io::Error> {
read_segment_with_timeout(self)
}
fn write_segment(
&mut self,
clock: Instant,
protocol_id: u16,
partial_payload: &[u8],
) -> Result<(), std::io::Error> {
write_segment(self, clock, protocol_id, partial_payload)
fn write_segment(&mut self, segment: Segment) -> Result<(), std::io::Error> {
write_segment(self, segment)
}
}
#[cfg(target_family = "unix")]
impl Bearer for UnixStream {
type Error = std::io::Error;
fn clone(&self) -> Self {
self.try_clone().expect("error cloning unix stream")
}
fn read_segment(&mut self) -> Result<(u16, u32, Payload), std::io::Error> {
read_segment(self)
fn read_segment(&mut self) -> Result<Option<Segment>, std::io::Error> {
read_segment_with_timeout(self)
}
fn write_segment(
&mut self,
clock: Instant,
protocol_id: u16,
partial_payload: &[u8],
) -> Result<(), std::io::Error> {
write_segment(self, clock, protocol_id, partial_payload)
fn write_segment(&mut self, segment: Segment) -> Result<(), std::io::Error> {
write_segment(self, segment)
}
}

View file

@ -0,0 +1,83 @@
use std::collections::HashMap;
use crate::{bearers::Bearer, std::Cancel, Payload};
pub struct EgressError(pub Payload);
pub trait Egress {
fn send(&self, payload: Payload) -> Result<(), EgressError>;
}
pub enum DemuxError<B: Bearer> {
BearerError(B::Error),
EgressDisconnected(u16, Payload),
EgressUnknown(u16, Payload),
}
pub enum TickOutcome {
Busy,
Idle,
}
/// A demuxer that reads from a bearer into the corresponding egress
pub struct Demuxer<B, E> {
bearer: B,
egress: HashMap<u16, E>,
}
impl<B, E> Demuxer<B, E>
where
B: Bearer,
E: Egress,
{
pub fn new(bearer: B) -> Self {
Demuxer {
bearer,
egress: Default::default(),
}
}
pub fn register(&mut self, id: u16, tx: E) {
self.egress.insert(id, tx);
}
fn dispatch(&self, protocol: u16, payload: Payload) -> Result<(), DemuxError<B>> {
match self.egress.get(&protocol) {
Some(tx) => match tx.send(payload) {
Err(EgressError(p)) => Err(DemuxError::EgressDisconnected(protocol, p)),
Ok(_) => Ok(()),
},
None => Err(DemuxError::EgressUnknown(protocol, payload)),
}
}
pub fn tick(&mut self) -> Result<TickOutcome, DemuxError<B>> {
match self.bearer.read_segment() {
Err(err) => Err(DemuxError::BearerError(err)),
Ok(None) => Ok(TickOutcome::Idle),
Ok(Some(segment)) => match self.dispatch(segment.protocol, segment.payload) {
Err(err) => Err(err),
Ok(()) => Ok(TickOutcome::Busy),
},
}
}
pub fn block(&mut self, cancel: Cancel) -> Result<(), B::Error> {
loop {
match self.tick() {
Ok(TickOutcome::Busy) => (),
Ok(TickOutcome::Idle) => match cancel.is_set() {
true => break Ok(()),
false => (),
},
Err(DemuxError::BearerError(err)) => return Err(err),
Err(DemuxError::EgressDisconnected(id, _)) => {
log::warn!("disconnected protocol {}", id)
}
Err(DemuxError::EgressUnknown(id, _)) => {
log::warn!("unknown protocol {}", id)
}
}
}
}
}

View file

@ -1,184 +1,41 @@
mod bearers;
use std::{
collections::HashMap,
io::{Read, Write},
sync::mpsc::{self, Receiver, Sender, TryRecvError},
thread::{self, JoinHandle},
time::{Duration, Instant},
};
use log::{debug, error, warn};
pub trait Bearer: Read + Write + Send + Sync + Sized {
fn read_segment(&mut self) -> Result<(u16, u32, Payload), std::io::Error>;
fn write_segment(
&mut self,
clock: Instant,
protocol_id: u16,
partial_payload: &[u8],
) -> Result<(), std::io::Error>;
fn clone(&self) -> Self;
}
const MAX_SEGMENT_PAYLOAD_LENGTH: usize = 65535;
pub mod agents;
pub mod bearers;
pub mod demux;
pub mod mux;
pub type Payload = Vec<u8>;
enum TxStepError {
BearerError(std::io::Error),
IngressDisconnected,
IngressEmpty,
pub struct Multiplexer<B, I, E>
where
B: bearers::Bearer,
I: mux::Ingress,
E: demux::Egress,
{
pub muxer: mux::Muxer<B, I>,
pub demuxer: demux::Demuxer<B, E>,
}
fn tx_step<TBearer>(
bearer: &mut TBearer,
ingress_id: u16,
ingress_rx: &mut Receiver<Payload>,
clock: Instant,
) -> Result<(), TxStepError>
impl<B, I, E> Multiplexer<B, I, E>
where
TBearer: Bearer,
B: bearers::Bearer,
I: mux::Ingress,
E: demux::Egress,
{
match ingress_rx.try_recv() {
Ok(payload) => {
let chunks = payload.chunks(MAX_SEGMENT_PAYLOAD_LENGTH);
for chunk in chunks {
bearer
.write_segment(clock, ingress_id, chunk)
.map_err(TxStepError::BearerError)?;
}
Ok(())
}
Err(TryRecvError::Disconnected) => Err(TxStepError::IngressDisconnected),
Err(TryRecvError::Empty) => Err(TxStepError::IngressEmpty),
}
}
fn tx_loop<TBearer>(bearer: &mut TBearer, ingress: MuxIngress)
where
TBearer: Bearer,
{
let mut rx_map: HashMap<_, _> = ingress.into_iter().collect();
loop {
let clock = Instant::now();
rx_map.retain(|id, rx| match tx_step(bearer, *id, rx, clock) {
Err(TxStepError::BearerError(err)) => {
error!("{:?}", err);
panic!();
}
Err(TxStepError::IngressDisconnected) => {
warn!("protocol handle {} disconnected", id);
false
}
Err(TxStepError::IngressEmpty) => {
thread::sleep(Duration::from_millis(10));
true
}
Ok(_) => true,
});
}
}
fn rx_loop<TBearer>(bearer: &mut TBearer, egress: DemuxerEgress)
where
TBearer: Bearer,
{
let mut tx_map: HashMap<_, _> = egress.into_iter().collect();
loop {
match bearer.read_segment() {
Err(err) => {
error!("{:?}", err);
panic!();
}
Ok(segment) => {
let (id, _ts, payload) = segment;
match tx_map.get(&id) {
Some(tx) => match tx.send(payload) {
Err(err) => {
error!("error sending egress tx to protocol, removing protocol from egress output. {:?}", err);
tx_map.remove(&id);
}
Ok(_) => {
debug!("successful tx to egress protocol");
}
},
None => warn!("received segment for protocol id not being demuxed {}", id),
}
}
pub fn new(bearer: B) -> Self {
Multiplexer {
muxer: mux::Muxer::new(bearer.clone()),
demuxer: demux::Demuxer::new(bearer.clone()),
}
}
}
pub struct Channel(pub Sender<Payload>, pub Receiver<Payload>);
type ChannelProtocolHandle = (u16, Channel);
type ChannelIngressHandle = (u16, Receiver<Payload>);
type ChannelEgressHandle = (u16, Sender<Payload>);
type MuxIngress = Vec<ChannelIngressHandle>;
type DemuxerEgress = Vec<ChannelEgressHandle>;
pub struct Multiplexer {
tx_thread: JoinHandle<()>,
rx_thread: JoinHandle<()>,
io_handles: HashMap<u16, Channel>,
}
impl Multiplexer {
pub fn setup<TBearer>(
bearer: TBearer,
protocols: &[u16],
) -> Result<Multiplexer, Box<dyn std::error::Error>>
where
TBearer: Bearer + 'static,
{
let handles = protocols.iter().map(|id| {
let (demux_tx, demux_rx) = mpsc::channel::<Payload>();
let (mux_tx, mux_rx) = mpsc::channel::<Payload>();
let channel = Channel(mux_tx, demux_rx);
let protocol_handle: ChannelProtocolHandle = (*id, channel);
let ingress_handle: ChannelIngressHandle = (*id, mux_rx);
let egress_handle: ChannelEgressHandle = (*id, demux_tx);
(protocol_handle, (ingress_handle, egress_handle))
});
let (protocol_handles, multiplex_handles): (Vec<_>, Vec<_>) = handles.into_iter().unzip();
let (ingress, egress): (Vec<_>, Vec<_>) = multiplex_handles.into_iter().unzip();
let mut tx_bearer = bearer.clone();
let tx_thread = thread::spawn(move || tx_loop(&mut tx_bearer, ingress));
let mut rx_bearer = bearer.clone();
let rx_thread = thread::spawn(move || rx_loop(&mut rx_bearer, egress));
let io_handles: HashMap<u16, Channel> = protocol_handles.into_iter().collect();
Ok(Multiplexer {
io_handles,
tx_thread,
rx_thread,
})
}
pub fn use_channel(&mut self, protocol_id: u16) -> Channel {
self.io_handles
.remove(&protocol_id)
.expect("requested channel not found in multiplexer")
}
pub fn join(self) {
self.tx_thread.join().expect("error joining tx loop thread");
self.rx_thread.join().expect("error joining rx loop thread");
pub fn register_channel(&mut self, protocol: u16, ingress: I, egress: E) {
self.muxer.register(protocol, ingress);
self.demuxer.register(protocol, egress);
}
}
#[cfg(feature = "std")]
mod std;
#[cfg(feature = "std")]
pub use crate::std::*;

View file

@ -0,0 +1,122 @@
use std::{collections::HashMap, time::Instant};
use rand::seq::SliceRandom;
use rand::thread_rng;
use crate::{
bearers::{Bearer, Segment},
std::Cancel,
Payload,
};
pub enum IngressError {
Disconnected,
Empty,
}
/// Source of payloads for a particular protocol
///
/// To be implemented by any mechanism that allows to submit a payloads from a
/// particular protocol that need to be muxed by the multiplexer.
pub trait Ingress {
fn try_recv(&mut self) -> Result<Payload, IngressError>;
}
type Message = (u16, Payload);
pub enum TickOutcome<TBearer>
where
TBearer: Bearer,
{
BearerError(TBearer::Error),
Idle,
Busy,
}
pub struct Muxer<B, I> {
bearer: B,
ingress: HashMap<u16, I>,
clock: Instant,
}
impl<B, I> Muxer<B, I>
where
B: Bearer,
I: Ingress,
{
pub fn new(bearer: B) -> Self {
Self {
bearer,
ingress: Default::default(),
clock: Instant::now(),
}
}
/// Register the receiver end of an ingress channel
pub fn register(&mut self, id: u16, rx: I) {
self.ingress.insert(id, rx);
}
/// Remove a protocol from the ingress
///
/// Meant to be used after a receive error in a previous tick
pub fn deregister(&mut self, id: u16) {
self.ingress.remove(&id);
}
#[inline]
fn randomize_ids(&self) -> Vec<u16> {
let mut rng = thread_rng();
let mut keys: Vec<_> = self.ingress.keys().cloned().collect();
keys.shuffle(&mut rng);
keys
}
/// Select the next segment to be muxed
///
/// This method iterates over the existing receivers checking for the first
/// available message. The order of the checks is random to ensure a fair
/// use of the multiplexer amongst all protocols.
pub fn select(&mut self) -> Option<Message> {
for id in self.randomize_ids() {
let rx = self.ingress.get_mut(&id).unwrap();
match rx.try_recv() {
Ok(payload) => return Some((id, payload)),
Err(IngressError::Disconnected) => {
self.deregister(id);
}
_ => (),
};
}
None
}
pub fn tick(&mut self) -> TickOutcome<B> {
match self.select() {
Some((id, payload)) => {
let segment = Segment::new(self.clock, id, payload);
match self.bearer.write_segment(segment) {
Err(err) => TickOutcome::BearerError(err),
_ => TickOutcome::Busy,
}
}
None => TickOutcome::Idle,
}
}
pub fn block(&mut self, cancel: Cancel) -> Result<(), B::Error> {
loop {
match self.tick() {
TickOutcome::BearerError(err) => return Err(err),
TickOutcome::Idle => match cancel.is_set() {
true => break Ok(()),
false => std::thread::yield_now(),
},
TickOutcome::Busy => (),
}
}
}
}

View file

@ -0,0 +1,123 @@
use crate::{agents, bearers::Bearer, demux, mux, Payload};
use std::{
sync::{
atomic::{AtomicBool, Ordering},
mpsc::{channel, Receiver, SendError, Sender, TryRecvError},
Arc,
},
thread::{spawn, JoinHandle},
};
pub type StdIngress = Receiver<Payload>;
impl mux::Ingress for StdIngress {
fn try_recv(&mut self) -> Result<Payload, mux::IngressError> {
match Receiver::try_recv(self) {
Ok(x) => Ok(x),
Err(TryRecvError::Disconnected) => Err(mux::IngressError::Disconnected),
Err(TryRecvError::Empty) => Err(mux::IngressError::Empty),
}
}
}
pub type StdEgress = Sender<Payload>;
impl demux::Egress for StdEgress {
fn send(&self, payload: Payload) -> Result<(), demux::EgressError> {
match Sender::send(self, payload) {
Ok(_) => Ok(()),
Err(SendError(p)) => Err(demux::EgressError(p)),
}
}
}
pub type StdPlexer<B> = crate::Multiplexer<B, StdIngress, StdEgress>;
pub type StdChannel = (Sender<Payload>, Receiver<Payload>);
impl agents::Channel for StdChannel {
fn enqueue_chunk(&mut self, payload: Payload) -> Result<(), agents::ChannelError> {
match self.0.send(payload) {
Ok(_) => Ok(()),
Err(SendError(payload)) => Err(agents::ChannelError::NotConnected(Some(payload))),
}
}
fn dequeue_chunk(&mut self) -> Result<Payload, agents::ChannelError> {
match self.1.recv() {
Ok(payload) => Ok(payload),
Err(_) => Err(agents::ChannelError::NotConnected(None)),
}
}
}
pub fn use_channel<B: Bearer>(plexer: &mut StdPlexer<B>, protocol: u16) -> StdChannel {
let (demux_tx, demux_rx) = channel::<Payload>();
let (mux_tx, mux_rx) = channel::<Payload>();
plexer.register_channel(protocol, mux_rx, demux_tx);
(mux_tx, demux_rx)
}
#[derive(Clone, Debug, Default)]
pub struct Cancel(Arc<AtomicBool>);
impl Cancel {
pub fn set(&self) {
self.0.store(true, Ordering::SeqCst);
}
pub fn is_set(&self) -> bool {
self.0.load(Ordering::SeqCst)
}
}
#[derive(Debug)]
pub struct Loop<B>
where
B: Bearer,
{
cancel: Cancel,
thread: JoinHandle<Result<(), B::Error>>,
}
impl<B> Loop<B>
where
B: Bearer,
{
pub fn cancel(&self) {
self.cancel.set();
}
pub fn join(self) -> Result<(), B::Error> {
self.thread.join().unwrap()
}
}
pub fn spawn_muxer<B, I>(mut muxer: mux::Muxer<B, I>) -> Loop<B>
where
B: Bearer + 'static,
B::Error: Send,
I: mux::Ingress + Send + 'static,
{
let cancel = Cancel::default();
let cancel2 = cancel.clone();
let thread = spawn(move || muxer.block(cancel2));
Loop { cancel, thread }
}
pub fn spawn_demuxer<B, E>(mut demuxer: demux::Demuxer<B, E>) -> Loop<B>
where
B: Bearer + 'static,
B::Error: Send,
E: demux::Egress + Send + 'static,
{
let cancel = Cancel::default();
let cancel2 = cancel.clone();
let thread = spawn(move || demuxer.block(cancel2));
Loop { cancel, thread }
}