sosistab 0.5.43

An obfuscated datagram transport for horrible networks
Documentation
use rtrb::{Consumer, Producer};
use smol::future::Future;
use smol::prelude::*;
use std::{
    io::Read,
    io::Write,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    task::Context,
    task::Poll,
};

/// Create a "bipe". Use async_dup's methods if you want something cloneable/shareable
pub fn bipe(capacity: usize) -> (BipeWriter, BipeReader) {
    let buffer = rtrb::RingBuffer::new(capacity);
    let (send_buf, recv_buf) = buffer.split();
    let write_ready = Arc::new(event_listener::Event::new());
    let read_ready = Arc::new(event_listener::Event::new());
    let closed = Arc::new(AtomicBool::new(false));
    (
        BipeWriter {
            queue: send_buf,
            signal: write_ready.clone(),
            signal_reader: read_ready.clone(),
            listener: write_ready.listen(),
            closed: closed.clone(),
        },
        BipeReader {
            queue: recv_buf,
            signal: read_ready.clone(),
            signal_writer: write_ready.clone(),
            listener: read_ready.listen(),
            closed,
        },
    )
}

/// Writing end of a byte pipe.
pub struct BipeWriter {
    queue: Producer<u8>,
    signal: Arc<event_listener::Event>,
    signal_reader: Arc<event_listener::Event>,
    listener: event_listener::EventListener,
    closed: Arc<AtomicBool>,
}

impl Drop for BipeWriter {
    fn drop(&mut self) {
        self.closed.store(true, Ordering::SeqCst);
        self.signal_reader.notify(1);
    }
}

fn broken_pipe() -> std::io::Error {
    std::io::Error::new(std::io::ErrorKind::ConnectionReset, "broken pipe")
}

impl AsyncWrite for BipeWriter {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        loop {
            if self.closed.load(Ordering::SeqCst) {
                return Poll::Ready(Err(broken_pipe()));
            }
            // if there's room in the buffer then it's fine
            {
                if let Ok(n) = self.queue.write(buf) {
                    // if n > 0 {
                    self.signal_reader.notify(1);
                    return Poll::Ready(Ok(n));
                    // }
                }
            }
            let listen_capacity = &mut self.listener;
            smol::pin!(listen_capacity);
            // there's no room, so we try again later
            smol::ready!(listen_capacity.poll(cx));
            self.listener = self.signal.listen()
        }
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        self.closed.store(true, Ordering::Relaxed);
        self.signal_reader.notify(1);
        Poll::Ready(Ok(()))
    }
}

/// Read end of a byte pipe.
pub struct BipeReader {
    queue: Consumer<u8>,
    signal: Arc<event_listener::Event>,
    signal_writer: Arc<event_listener::Event>,
    listener: event_listener::EventListener,
    closed: Arc<AtomicBool>,
}

impl AsyncRead for BipeReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<std::io::Result<usize>> {
        loop {
            if let Ok(n) = self.queue.read(buf) {
                if n > 0 {
                    self.signal_writer.notify(1);
                    return Poll::Ready(Ok(n));
                }
            }
            if self.closed.load(Ordering::Relaxed) {
                return Poll::Ready(Err(broken_pipe()));
            }
            let listen_new_data = &mut self.listener;
            smol::pin!(listen_new_data);
            smol::ready!(listen_new_data.poll(cx));
            self.listener = self.signal.listen();
        }
    }
}