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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use std::cell::Cell;

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
};

use crate::{
    endpoint_proto::PlaneHandshake,
    rt::WorkerGroup,
    EndpointAddr,
    HeapBufferPool,
    InterfaceRole,
    PlaneConfig,
    RoleStartFn,
    RuntimeHandle,
    TransportHandle,
    TcpTransport,
    TopicChannels,
    WorkerId,
};

pub struct WebSocketServer {
    next_plane_id: Cell<u32>,
}

impl WebSocketServer {
    pub fn new() -> Self {
        Self {
            next_plane_id: Cell::new(0),
        }
    }

    pub async fn accept_local<Role: InterfaceRole>(
        &self,
        rt: &RuntimeHandle,
        buffer_pool: HeapBufferPool,
        tcp_stream: TcpStream,
        start_fn: impl RoleStartFn<Role>,
        init: &Role::Init,
    ) -> std::io::Result<Role::Hooks>
        where Role::Init: Clone + Send + Sync,
    {
        let plane_id = self.next_plane_id.get();
        self.next_plane_id.set(self.next_plane_id.get() + 1);

        let Ok(mut websocket) = tokio_tungstenite::accept_async(tcp_stream).await else {
            todo!();
        };

        Self::handshake(&mut stream, plane_id, EndpointAddr { endpoint: 1 }, init).await?;

        let transport = rt.add_transport(WebSocketTransport {
            websocket,
            buffer_pool,
        })
        .await;

        let mut plane_builder = rt.start_plane::<Role>(PlaneConfig {
            plane_id,
            endpoint_addr: EndpointAddr { endpoint: 0 },
            transport: &transport,
            topic_channels: TopicChannels::SingleChannel { channel_id: 0 },
            init,
        });

        let transport_shutdown_signal = transport.shutdown_signal.clone();
        let role_handle = plane_builder.local(|cx| {
            // Shutdown the plane when the transport is shutdown.
            let role_shutdown_signal = cx.shutdown_signal.clone();
            cx.spawner.spawn(async move {
                transport_shutdown_signal.wait().await;
                role_shutdown_signal.notify();
            });

            start_fn(cx);
        }).await;

        Ok(role_handle)
    }

    async fn handshake<Init: mproto::Owned>(
        stream: &mut TcpStream,
        plane_id: u32,
        endpoint_addr: EndpointAddr,
        init: &Init,
    ) -> std::io::Result<()> {
        // Send connect respose
        let payload = PlaneHandshake { plane_id, endpoint_addr, init };
        let payload_len = mproto::encoded_len(&payload);
        let mut payload_buf = vec![0u8; 2 + payload_len];
        payload_buf[..2].copy_from_slice(&(payload_len as u16).to_le_bytes());
        mproto::encode_value(payload, &mut payload_buf[2..]);
        stream.write_all(&payload_buf[..]).await?;

        Ok(())
    }
}

pub async fn tcp_connect<Role: InterfaceRole>(
    rt: &RuntimeHandle,
    buffer_pool: HeapBufferPool,
    worker_id: WorkerId,
    mut stream: TcpStream,
) -> std::io::Result<(TransportHandle, Role::Hooks)>
    where Role::Init: Clone + std::fmt::Debug + Send + Sync,
{
    stream.set_nodelay(true).unwrap();

    let mut payload_len_bytes = [0u8; 2];
    stream.read_exact(&mut payload_len_bytes).await?;
    let payload_len = u16::from_le_bytes(payload_len_bytes);
    let mut payload_bytes = vec![0u8; payload_len as usize];
    stream.read_exact(&mut payload_bytes).await?;
    let plane_handshake: PlaneHandshake<Role::Init> = mproto::decode_value(&payload_bytes)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;

    let transport = rt.add_transport(TcpTransport {
        worker_id,
        stream,
        buffer_pool,
    })
    .await;

    let transport_shutdown_signal = transport.shutdown_signal.clone();

    let role_handle = rt.start_plane::<Role>(PlaneConfig {
        plane_id: plane_handshake.plane_id,
        endpoint_addr: plane_handshake.endpoint_addr,
        transport: &transport,
        topic_channels: TopicChannels::SingleChannel { channel_id: 0 },
        init: &plane_handshake.init,
    })
    .local(|cx| {
        // Shutdown the plane when the transport is shutdown.
        let role_shutdown_signal = cx.shutdown_signal.clone();
        cx.spawner.spawn(async move {
            transport_shutdown_signal.wait().await;
            role_shutdown_signal.notify();
        });
    })
    .await;

    Ok((transport, role_handle))
}