modrpc 0.0.3

Runtime for modrpc - a modular RPC 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114

use futures_util::{SinkExt, StreamExt};
use tokio::net::TcpStream;
use tokio_tungstenite::{WebSocketStream, tungstenite::protocol::Message};

use futures_lite::future;
use mproto::BaseLen;

use crate::{
    HeapBufferPool, Packet, PacketBundle, TransportBuilder, TransportContext, TransportHandle,
    transport::WriterConfig, ws_ingress::WebSocketIngress,
};

pub struct WebSocketTransport {
    pub buffer_pool: HeapBufferPool,
    pub websocket: WebSocketStream<TcpStream>,
}

impl TransportBuilder for WebSocketTransport {
    async fn start_transport(self, cx: TransportContext<'_>) -> TransportHandle {
        let shutdown_signal = bab::SignalTree::new();

        let Some(worker_cx) = cx.rt.local_worker_context() else {
            panic!("tokio_tungstenite WebSocket modrpc transport requires a local worker.");
        };

        let (mut ws_tx, ws_rx) = self.websocket.split();

        let (writer_flush_sender, mut writer_flush_receiver) = bab::new_writer_flusher();

        // Spawn task to flush egress packets
        worker_cx.spawn({
            let buffer_pool = self.buffer_pool.clone();
            let shutdown_signal = shutdown_signal.clone();
            let shutdown_notifier = shutdown_signal.clone();
            future::or(
                async move {
                    'flush_loop: loop {
                        for flush in writer_flush_receiver.flush().await {
                            if flush.len() > 0 {
                                let mut buf = vec![0u8; PacketBundle::BASE_LEN + flush.len()];

                                // Fill bundle header
                                mproto::encode_value(
                                    PacketBundle {
                                        channel_id: flush.writer_id() as u32,
                                        length: flush.len() as u16,
                                    },
                                    &mut buf[..PacketBundle::BASE_LEN],
                                );
                                // Copy payload
                                buf[PacketBundle::BASE_LEN..].copy_from_slice(&flush);

                                // Write to socket
                                if let Err(_) = ws_tx.send(Message::Binary(buf.into())).await {
                                    break 'flush_loop;
                                }
                            }
                        }
                    }
                    shutdown_notifier.notify();
                },
                async move {
                    let _buffer_pool_thread_guard = buffer_pool.register_thread();
                    shutdown_signal.wait().await;
                    //stream.shutdown(std::net::Shutdown::Both).unwrap();
                },
            )
        });

        // Spawn task to receive ingress packets
        let mut ws_ingress = WebSocketIngress::new(
            ws_rx,
            self.buffer_pool.clone(),
            self.buffer_pool.buffer_size(),
        );
        worker_cx.spawn({
            let shutdown_notifier = shutdown_signal.clone();
            let shutdown_waiter = shutdown_signal.clone();
            let process_packet_fn = worker_cx.get_packet_processor();
            future::or(
                async move {
                    use core::mem::MaybeUninit;

                    let mut shatter_offsets: Vec<usize> = Vec::new();
                    let mut shatter_out_packets: Vec<MaybeUninit<Packet>> = Vec::new();

                    while let Ok(packet_bundle) = ws_ingress.receive().await {
                        let _header = crate::shatter_packet_bundle(
                            packet_bundle,
                            &mut shatter_offsets,
                            &mut shatter_out_packets,
                        );

                        for packet in shatter_out_packets.drain(..) {
                            let packet = unsafe { packet.assume_init() };
                            process_packet_fn(&packet).await;
                        }
                    }

                    shutdown_notifier.notify();
                },
                shutdown_waiter.wait_owned(),
            )
        });

        TransportHandle {
            shutdown_signal,
            buffer_pool: self.buffer_pool,
            writer_config: WriterConfig::LocalFlush {
                writer_flush_sender,
            },
        }
    }
}