Start work on RpcClient

zed/src/lib.rs

@@ -1,9 +1,11 @@
 use anyhow::{anyhow, Context, Result};
 use gpui::{AsyncAppContext, MutableAppContext, Task};
-use std::{convert::TryFrom, time::Duration};
+use rpc_client::RpcClient;
+use std::{convert::TryFrom, net::Shutdown, time::Duration};
 use tiny_http::{Header, Response, Server};
 use url::Url;
 use util::SurfResultExt;
+use zed_rpc::{proto, rest::CreateWorktreeResponse};
 
 pub mod assets;
 pub mod editor;
@@ -11,6 +13,7 @@ pub mod file_finder;
 pub mod language;
 pub mod menus;
 mod operation_queue;
+mod rpc_client;
 pub mod settings;
 mod sum_tree;
 #[cfg(test)]
@@ -33,7 +36,9 @@ pub fn init(cx: &mut MutableAppContext) {
 
 fn share_worktree(_: &(), cx: &mut MutableAppContext) {
     let zed_url = std::env::var("ZED_SERVER_URL").unwrap_or("https://zed.dev".to_string());
-    cx.spawn::<_, _, surf::Result<()>>(|cx| async move {
+    let executor = cx.background_executor().clone();
+
+    let task = cx.spawn::<_, _, surf::Result<()>>(|cx| async move {
         let (user_id, access_token) = login(zed_url.clone(), &cx).await?;
 
         let mut response = surf::post(format!("{}/api/worktrees", &zed_url))
@@ -44,28 +49,47 @@ fn share_worktree(_: &(), cx: &mut MutableAppContext) {
             .await
             .context("")?;
 
-        let body = response
-            .body_json::<zed_rpc::rest::CreateWorktreeResponse>()
-            .await?;
+        let CreateWorktreeResponse {
+            worktree_id,
+            rpc_address,
+        } = response.body_json().await?;
+
+        eprintln!("got worktree response: {:?} {:?}", worktree_id, rpc_address);
 
         // TODO - If the `ZED_SERVER_URL` uses https, then wrap this stream in
         // a TLS stream using `native-tls`.
-        let stream = smol::net::TcpStream::connect(body.rpc_address).await?;
+        let stream = smol::net::TcpStream::connect(rpc_address).await?;
 
-        let mut message_stream = zed_rpc::proto::MessageStream::new(stream);
-        message_stream
-            .write_message(&zed_rpc::proto::FromClient {
-                id: 0,
-                variant: Some(zed_rpc::proto::from_client::Variant::Auth(
-                    zed_rpc::proto::from_client::Auth {
-                        user_id: user_id.parse::<i32>()?,
-                        access_token,
-                    },
-                )),
+        let mut rpc_client = RpcClient::new(stream, executor, |stream| {
+            stream.shutdown(Shutdown::Read).ok();
+        });
+
+        let auth_response = rpc_client
+            .request(proto::from_client::Auth {
+                user_id: user_id.parse::<i32>()?,
+                access_token,
+            })
+            .await?;
+        if !auth_response.credentials_valid {
+            Err(anyhow!("failed to authenticate with RPC server"))?;
+        }
+
+        let share_response = rpc_client
+            .request(proto::from_client::ShareWorktree {
+                worktree_id: worktree_id as u64,
+                files: Vec::new(),
             })
             .await?;
 
+        log::info!("sharing worktree {:?}", share_response);
+
         Ok(())
+    });
+
+    cx.spawn(|_| async move {
+        if let Err(e) = task.await {
+            log::error!("sharing failed: {}", e);
+        }
     })
     .detach();
 }

zed/src/rpc_client.rs

@@ -0,0 +1,97 @@
+use anyhow::{anyhow, Result};
+use gpui::executor::Background;
+use parking_lot::Mutex;
+use postage::{
+    oneshot,
+    prelude::{Sink, Stream},
+};
+use smol::prelude::{AsyncRead, AsyncWrite};
+use std::{collections::HashMap, sync::Arc};
+use zed_rpc::proto::{self, MessageStream, RequestMessage, SendMessage, ServerMessage};
+
+pub struct RpcClient<Conn, ShutdownFn>
+where
+    ShutdownFn: FnMut(&mut Conn),
+{
+    stream: MessageStream<Conn>,
+    response_channels: Arc<Mutex<HashMap<i32, oneshot::Sender<proto::from_server::Variant>>>>,
+    next_message_id: i32,
+    shutdown_fn: ShutdownFn,
+}
+
+impl<Conn, ShutdownFn> RpcClient<Conn, ShutdownFn>
+where
+    Conn: Clone + AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    ShutdownFn: FnMut(&mut Conn),
+{
+    pub fn new(conn: Conn, executor: Arc<Background>, shutdown_fn: ShutdownFn) -> Self {
+        let response_channels = Arc::new(Mutex::new(HashMap::new()));
+
+        let result = Self {
+            next_message_id: 0,
+            stream: MessageStream::new(conn.clone()),
+            response_channels: response_channels.clone(),
+            shutdown_fn,
+        };
+
+        executor
+            .spawn::<Result<()>, _>(async move {
+                let mut stream = MessageStream::new(conn);
+                loop {
+                    let message = stream.read_message::<proto::FromServer>().await?;
+                    if let Some(variant) = message.variant {
+                        if let Some(request_id) = message.request_id {
+                            let tx = response_channels.lock().remove(&request_id);
+                            if let Some(mut tx) = tx {
+                                tx.send(variant).await?;
+                            } else {
+                                log::warn!(
+                                    "received RPC response to unknown request id {}",
+                                    request_id
+                                );
+                            }
+                        }
+                    } else {
+                        log::warn!("received RPC message with no content");
+                    }
+                }
+            })
+            .detach();
+
+        result
+    }
+
+    pub async fn request<T: RequestMessage>(&mut self, req: T) -> Result<T::Response> {
+        let message_id = self.next_message_id;
+        self.next_message_id += 1;
+
+        let (tx, mut rx) = oneshot::channel();
+        self.response_channels.lock().insert(message_id, tx);
+
+        self.stream
+            .write_message(&proto::FromClient {
+                id: message_id,
+                variant: Some(req.to_variant()),
+            })
+            .await?;
+        let response = rx
+            .recv()
+            .await
+            .expect("response channel was unexpectedly dropped");
+        T::Response::from_variant(response)
+            .ok_or_else(|| anyhow!("received response of the wrong t"))
+    }
+
+    pub async fn send<T: SendMessage>(_: T) -> Result<()> {
+        todo!()
+    }
+}
+
+impl<Conn, ShutdownFn> Drop for RpcClient<Conn, ShutdownFn>
+where
+    ShutdownFn: FnMut(&mut Conn),
+{
+    fn drop(&mut self) {
+        (self.shutdown_fn)(self.stream.inner_mut())
+    }
+}