Initial commit - extracted urpc from monorepo

1 files changed, 124 insertions, 0 deletions

diff --git a/src/client_channel.rs b/src/client_channel.rs
new file mode 100644
index 0000000..5666a16
--- /dev/null
+++ b/src/client_channel.rs
@@ -0,0 +1,124 @@
+use std::{collections::HashMap, sync::Arc};
+
+use bytes::Bytes;
+use futures::{SinkExt, StreamExt};
+use tokio::{
+    sync::{mpsc, oneshot},
+    task::AbortHandle,
+};
+
+use crate::{
+    Channel,
+    protocol::{RpcCall, RpcMessage, RpcResponse},
+};
+
+const CLIENT_CHANNEL_BUFFER_SIZE: usize = 64;
+
+struct ClientChannelMessage {
+    service: String,
+    method: String,
+    arguments: Bytes,
+    responder: oneshot::Sender<std::io::Result<Bytes>>,
+}
+
+struct ClientChannelInner {
+    sender: mpsc::Sender<ClientChannelMessage>,
+    abort_handle: AbortHandle,
+}
+
+impl Drop for ClientChannelInner {
+    fn drop(&mut self) {
+        self.abort_handle.abort();
+    }
+}
+
+#[derive(Clone)]
+pub struct ClientChannel(Arc<ClientChannelInner>);
+
+impl ClientChannel {
+    pub fn new<C: Channel>(channel: C) -> Self {
+        let (tx, rx) = mpsc::channel(CLIENT_CHANNEL_BUFFER_SIZE);
+        let abort_handle = tokio::spawn(client_channel_loop(channel, rx)).abort_handle();
+        Self(Arc::new(ClientChannelInner {
+            sender: tx,
+            abort_handle,
+        }))
+    }
+
+    pub async fn call(
+        &self,
+        service: String,
+        method: String,
+        arguments: Bytes,
+    ) -> std::io::Result<Bytes> {
+        let (tx, rx) = oneshot::channel();
+        self.0
+            .sender
+            .send(ClientChannelMessage {
+                service,
+                method,
+                arguments,
+                responder: tx,
+            })
+            .await
+            .expect("client channel task should never shutdown while a client is alive");
+        rx.await
+            .expect("client channel task should never shutdown while a client is alive")
+    }
+}
+
+async fn client_channel_loop<C: Channel>(
+    mut channel: C,
+    mut rx: mpsc::Receiver<ClientChannelMessage>,
+) {
+    enum Select {
+        RpcMessage(RpcMessage),
+        ClientChannelMessage(ClientChannelMessage),
+    }
+
+    let mut responders = HashMap::<u64, oneshot::Sender<std::io::Result<Bytes>>>::default();
+    let mut rpc_call_id = 0;
+
+    loop {
+        let select = tokio::select! {
+            Some(Ok(v)) = channel.next() => Select::RpcMessage(v),
+            Some(v) = rx.recv() => Select::ClientChannelMessage(v),
+        };
+
+        match select {
+            Select::RpcMessage(RpcMessage::Response(RpcResponse { id, value })) => {
+                if let Some(responder) = responders.remove(&id) {
+                    let _ = responder.send(Ok(value));
+                }
+            }
+            Select::RpcMessage(_) => todo!(),
+            Select::ClientChannelMessage(ClientChannelMessage {
+                service,
+                method,
+                arguments,
+                responder,
+            }) => {
+                let id = rpc_call_id;
+                rpc_call_id += 1;
+
+                let result = channel
+                    .send(RpcMessage::Call(RpcCall {
+                        id,
+                        service,
+                        method,
+                        arguments,
+                    }))
+                    .await;
+
+                match result {
+                    Ok(()) => {
+                        responders.insert(id, responder);
+                    }
+                    Err(err) => {
+                        let _ = responder.send(Err(err));
+                    }
+                }
+            }
+        }
+    }
+}