rama 0.4.0

modular service framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

pub(crate) mod fixture;

use pin_project_lite::pin_project;
use rama::telemetry::tracing;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll, ready};
use tokio::io::{AsyncRead, ReadBuf};
use tokio::sync::mpsc;

// Common read half shared by both stream types
pin_project! {
    #[derive(Debug)]
    pub(crate) struct StreamReadHalf {
        #[pin]
        read_rx: mpsc::UnboundedReceiver<Vec<u8>>,
        read_buffer: Vec<u8>,
    }
}

impl StreamReadHalf {
    pub(crate) fn new(read_rx: mpsc::UnboundedReceiver<Vec<u8>>) -> Self {
        Self {
            read_rx,
            read_buffer: Vec::new(),
        }
    }
}

impl AsyncRead for StreamReadHalf {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let mut this = self.as_mut().project();

        // First, try to satisfy the read request from the internal buffer
        if !this.read_buffer.is_empty() {
            let to_read = std::cmp::min(this.read_buffer.len(), buf.remaining());
            // Copy data from internal buffer to the read buffer
            buf.put_slice(&this.read_buffer[..to_read]);
            // Remove the consumed data from the internal buffer
            this.read_buffer.drain(..to_read);
            return Poll::Ready(Ok(()));
        }

        // If internal buffer is empty, try to get data from the channel
        match this.read_rx.as_mut().get_mut().try_recv() {
            Ok(data) => {
                // Copy as much data as we can fit in the buffer
                let to_read = std::cmp::min(data.len(), buf.remaining());
                buf.put_slice(&data[..to_read]);

                // Store any remaining data in the internal buffer for next time
                if to_read < data.len() {
                    let remaining = &data[to_read..];
                    this.read_buffer.extend_from_slice(remaining);
                }
                Poll::Ready(Ok(()))
            }
            Err(mpsc::error::TryRecvError::Empty) => {
                match ready!(this.read_rx.poll_recv(cx)) {
                    Some(data) => {
                        // Copy as much data as we can fit in the buffer
                        let to_read = std::cmp::min(data.len(), buf.remaining());
                        buf.put_slice(&data[..to_read]);

                        // Store any remaining data in the internal buffer for next time
                        if to_read < data.len() {
                            let remaining = &data[to_read..];
                            this.read_buffer.extend_from_slice(remaining);
                        }
                        Poll::Ready(Ok(()))
                    }
                    None => Poll::Ready(Ok(())),
                }
            }
            Err(mpsc::error::TryRecvError::Disconnected) => {
                // Channel closed, return EOF
                Poll::Ready(Ok(()))
            }
        }
    }
}

pub(crate) fn init_tracing() {
    use std::sync::Once;
    static INIT: Once = Once::new();
    INIT.call_once(|| {
        drop(
            tracing_subscriber::fmt()
                .with_max_level(tracing::Level::INFO)
                .with_target(true)
                .with_thread_ids(true)
                .with_thread_names(true)
                .try_init(),
        );
    });
}