protosocket_server/
connection_server.rs

1use std::future::Future;
2use std::io::Error;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::task::Context;
6use std::task::Poll;
7
8use protosocket::Connection;
9use protosocket::ConnectionBindings;
10use protosocket::Serializer;
11use tokio::sync::mpsc;
12
13pub trait ServerConnector: Unpin {
14    type Bindings: ConnectionBindings;
15
16    fn serializer(&self) -> <Self::Bindings as ConnectionBindings>::Serializer;
17    fn deserializer(&self) -> <Self::Bindings as ConnectionBindings>::Deserializer;
18
19    fn new_reactor(
20        &self,
21        optional_outbound: mpsc::Sender<
22            <<Self::Bindings as ConnectionBindings>::Serializer as Serializer>::Message,
23        >,
24    ) -> <Self::Bindings as ConnectionBindings>::Reactor;
25
26    fn maximum_message_length(&self) -> usize {
27        4 * (2 << 20)
28    }
29
30    fn max_queued_outbound_messages(&self) -> usize {
31        256
32    }
33}
34
35/// A `protosocket::Connection` is an IO driver. It directly uses tokio's io wrapper of mio to poll
36/// the OS's io primitives, manages read and write buffers, and vends messages to & from connections.
37/// Connections send messages to the ConnectionServer through an mpsc channel, and they receive
38/// inbound messages via a reactor callback.
39///
40/// Protosockets are monomorphic messages: You can only have 1 kind of message per service.
41/// The expected way to work with this is to use prost and protocol buffers to encode messages.
42/// Of course you can do whatever you want, as the telnet example shows.
43///
44/// Protosocket messages are not opinionated about request & reply. If you are, you will need
45/// to implement such a thing. This allows you freely choose whether you want to send
46/// fire-&-forget messages sometimes; however it requires you to write your protocol's rules.
47/// You get an inbound iterable of <MessageIn> batches and an outbound stream of <MessageOut> per
48/// connection - you decide what those mean for you!
49///
50/// A ProtosocketServer is a future: You spawn it and it runs forever.
51pub struct ProtosocketServer<Connector: ServerConnector> {
52    connector: Connector,
53    listener: tokio::net::TcpListener,
54    max_buffer_length: usize,
55    max_queued_outbound_messages: usize,
56    runtime: tokio::runtime::Handle,
57}
58
59impl<Connector: ServerConnector> ProtosocketServer<Connector> {
60    /// Construct a new `ProtosocketServer` listening on the provided address.
61    /// The address will be bound and listened upon with `SO_REUSEADDR` set.
62    /// The server will use the provided runtime to spawn new tcp connections as `protosocket::Connection`s.
63    pub async fn new(
64        address: std::net::SocketAddr,
65        runtime: tokio::runtime::Handle,
66        connector: Connector,
67    ) -> crate::Result<Self> {
68        let listener = tokio::net::TcpListener::bind(address)
69            .await
70            .map_err(Arc::new)?;
71        Ok(Self {
72            connector,
73            listener,
74            max_buffer_length: 16 * (2 << 20),
75            max_queued_outbound_messages: 128,
76            runtime,
77        })
78    }
79
80    /// Set the maximum buffer length for connections created by this server after the setting is applied.
81    pub fn set_max_buffer_length(&mut self, max_buffer_length: usize) {
82        self.max_buffer_length = max_buffer_length;
83    }
84
85    /// Set the maximum queued outbound messages for connections created by this server after the setting is applied.
86    pub fn set_max_queued_outbound_messages(&mut self, max_queued_outbound_messages: usize) {
87        self.max_queued_outbound_messages = max_queued_outbound_messages;
88    }
89}
90
91impl<Connector: ServerConnector> Future for ProtosocketServer<Connector> {
92    type Output = Result<(), Error>;
93
94    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
95        loop {
96            break match self.listener.poll_accept(context) {
97                Poll::Ready(result) => match result {
98                    Ok((stream, address)) => {
99                        stream.set_nodelay(true)?;
100                        let (outbound_submission_queue, outbound_messages) =
101                            mpsc::channel(self.max_queued_outbound_messages);
102                        let reactor = self
103                            .connector
104                            .new_reactor(outbound_submission_queue.clone());
105                        let connection: Connection<Connector::Bindings> = Connection::new(
106                            stream,
107                            address,
108                            self.connector.deserializer(),
109                            self.connector.serializer(),
110                            self.max_buffer_length,
111                            self.max_queued_outbound_messages,
112                            outbound_messages,
113                            reactor,
114                        );
115                        self.runtime.spawn(connection);
116                        continue;
117                    }
118                    Err(e) => {
119                        log::error!("failed to accept connection: {e:?}");
120                        continue;
121                    }
122                },
123                Poll::Pending => Poll::Pending,
124            };
125        }
126    }
127}