Remove remaining instances of router

Co-Authored-By: Nathan Sobo <nathan@zed.dev>
Co-Authored-By: Max Brunsfeld <max@zed.dev>

server/src/rpc.rs

@@ -29,7 +29,7 @@ use tide::{
 use time::OffsetDateTime;
 use zrpc::{
     auth::random_token,
-    proto::{self, EnvelopedMessage},
+    proto::{self, AnyTypedEnvelope, EnvelopedMessage},
     ConnectionId, Peer, TypedEnvelope,
 };
 
@@ -38,16 +38,12 @@ type ReplicaId = u16;
 type MessageHandler = Box<
     dyn Send
         + Sync
-        + Fn(
+        + Fn(Box<dyn AnyTypedEnvelope>, Arc<Server>) -> BoxFuture<'static, tide::Result<()>>,
-            &mut Option<Box<dyn Any + Send + Sync>>,
-            Arc<Server>,
-        ) -> Option<BoxFuture<'static, tide::Result<()>>>,
 >;
 
 #[derive(Default)]
 struct ServerBuilder {
-    handlers: Vec<MessageHandler>,
+    handlers: HashMap<TypeId, MessageHandler>,
-    handler_types: HashSet<TypeId>,
 }
 
 impl ServerBuilder {
@@ -57,24 +53,17 @@ impl ServerBuilder {
         Fut: 'static + Send + Future<Output = tide::Result<()>>,
         M: EnvelopedMessage,
     {
-        if self.handler_types.insert(TypeId::of::<M>()) {
+        let prev_handler = self.handlers.insert(
+            TypeId::of::<M>(),
+            Box::new(move |envelope, server| {
+                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
+                (handler)(envelope, server).boxed()
+            }),
+        );
+        if prev_handler.is_some() {
             panic!("registered a handler for the same message twice");
         }
 
-        self.handlers
-            .push(Box::new(move |untyped_envelope, server| {
-                if let Some(typed_envelope) = untyped_envelope.take() {
-                    match typed_envelope.downcast::<TypedEnvelope<M>>() {
-                        Ok(typed_envelope) => Some((handler)(typed_envelope, server).boxed()),
-                        Err(envelope) => {
-                            *untyped_envelope = Some(envelope);
-                            None
-                        }
-                    }
-                } else {
-                    None
-                }
-            }));
         self
     }
 
@@ -90,16 +79,17 @@ impl ServerBuilder {
 pub struct Server {
     rpc: Arc<Peer>,
     state: Arc<AppState>,
-    handlers: Vec<MessageHandler>,
+    handlers: HashMap<TypeId, MessageHandler>,
 }
 
 impl Server {
-    pub async fn handle_connection<Conn>(
+    pub fn handle_connection<Conn>(
         self: &Arc<Self>,
         connection: Conn,
         addr: String,
         user_id: UserId,
-    ) where
+    ) -> impl Future<Output = ()>
+    where
         Conn: 'static
             + futures::Sink<WebSocketMessage, Error = WebSocketError>
             + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>
@@ -107,7 +97,9 @@ impl Server {
             + Unpin,
     {
         let this = self.clone();
-        let (connection_id, handle_io, mut incoming_rx) = this.rpc.add_connection(connection).await;
+        async move {
+            let (connection_id, handle_io, mut incoming_rx) =
+                this.rpc.add_connection(connection).await;
             this.state
                 .rpc
                 .write()
@@ -123,21 +115,15 @@ impl Server {
                     message = next_message => {
                         if let Some(message) = message {
                             let start_time = Instant::now();
-                        log::info!("RPC message received");
+                            log::info!("RPC message received: {}", message.payload_type_name());
-                        let mut message = Some(message);
+                            if let Some(handler) = this.handlers.get(&message.payload_type_id()) {
-                        for handler in &this.handlers {
+                                if let Err(err) = (handler)(message, this.clone()).await {
-                            if let Some(future) = (handler)(&mut message, this.clone()) {
-                                 if let Err(err) = future.await {
                                     log::error!("error handling message: {:?}", err);
                                 } else {
                                     log::info!("RPC message handled. duration:{:?}", start_time.elapsed());
                                 }
-                                break;
+                            } else {
-                            }
+                                log::warn!("unhandled message: {}", message.payload_type_name());
-                        }
-
-                        if let Some(message) = message {
-                            log::warn!("unhandled message: {:?}", message);
                             }
                         } else {
                             log::info!("rpc connection closed {:?}", addr);
@@ -157,6 +143,7 @@ impl Server {
                 log::error!("error signing out connection {:?} - {:?}", addr, err);
             }
         }
+    }
 }
 
 #[derive(Default)]

server/src/tests.rs

@@ -1,9 +1,7 @@
 use crate::{
     auth,
     db::{self, UserId},
-    github,
+    github, rpc, AppState, Config,
-    rpc::{self, build_server},
-    AppState, Config,
 };
 use async_std::task;
 use gpui::TestAppContext;
@@ -28,6 +26,8 @@ use zrpc::Peer;
 
 #[gpui::test]
 async fn test_share_worktree(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    tide::log::start();
+
     let (window_b, _) = cx_b.add_window(|_| EmptyView);
     let settings = settings::channel(&cx_b.font_cache()).unwrap().1;
     let lang_registry = Arc::new(LanguageRegistry::new());
@@ -514,9 +514,9 @@ async fn test_basic_chat(mut cx_a: TestAppContext, cx_b: TestAppContext) {
     .await
     .unwrap();
 
-    let channels_a = client_a.get_channels().await;
+    // let channels_a = client_a.get_channels().await;
-    assert_eq!(channels_a.len(), 1);
+    // assert_eq!(channels_a.len(), 1);
-    assert_eq!(channels_a[0].read(&cx_a).name(), "test-channel");
+    // assert_eq!(channels_a[0].read(&cx_a).name(), "test-channel");
 
     // assert_eq!(
     //     db.get_recent_channel_messages(channel_id, 50)
@@ -530,8 +530,8 @@ async fn test_basic_chat(mut cx_a: TestAppContext, cx_b: TestAppContext) {
 struct TestServer {
     peer: Arc<Peer>,
     app_state: Arc<AppState>,
+    server: Arc<rpc::Server>,
     db_name: String,
-    router: Arc<Router>,
 }
 
 impl TestServer {
@@ -540,36 +540,27 @@ impl TestServer {
         let db_name = format!("zed-test-{}", rng.gen::<u128>());
         let app_state = Self::build_app_state(&db_name).await;
         let peer = Peer::new();
-        let mut router = Router::new();
+        let server = rpc::build_server(&app_state, &peer);
-        build_server(&mut router, &app_state, &peer);
         Self {
             peer,
-            router: Arc::new(router),
             app_state,
+            server,
             db_name,
         }
     }
 
     async fn create_client(&mut self, cx: &mut TestAppContext, name: &str) -> (UserId, Client) {
         let user_id = self.app_state.db.create_user(name, false).await.unwrap();
-        let lang_registry = Arc::new(LanguageRegistry::new());
+        let client = Client::new();
-        let client = Client::new(lang_registry.clone());
-        let mut client_router = ForegroundRouter::new();
-        cx.update(|cx| zed::worktree::init(cx, &client, &mut client_router));
-
         let (client_conn, server_conn) = Channel::bidirectional();
         cx.background()
-            .spawn(rpc::handle_connection(
+            .spawn(
-                self.peer.clone(),
+                self.server
-                self.router.clone(),
+                    .handle_connection(server_conn, name.to_string(), user_id),
-                self.app_state.clone(),
+            )
-                name.to_string(),
-                server_conn,
-                user_id,
-            ))
             .detach();
         client
-            .add_connection(client_conn, Arc::new(client_router), cx.to_async())
+            .add_connection(client_conn, cx.to_async())
             .await
             .unwrap();
 

zed/src/channel.rs

@@ -1,6 +1,6 @@
 use crate::rpc::{self, Client};
 use anyhow::Result;
-use gpui::{Entity, ModelContext, Task, WeakModelHandle};
+use gpui::{Entity, ModelContext, WeakModelHandle};
 use std::{
     collections::{HashMap, VecDeque},
     sync::Arc,
@@ -22,7 +22,7 @@ pub struct Channel {
     first_message_id: Option<u64>,
     messages: Option<VecDeque<ChannelMessage>>,
     rpc: Arc<Client>,
-    _message_handler: Task<()>,
+    _subscription: rpc::Subscription,
 }
 
 pub struct ChannelMessage {
@@ -50,20 +50,20 @@ impl Entity for Channel {
 
 impl Channel {
     pub fn new(details: ChannelDetails, rpc: Arc<Client>, cx: &mut ModelContext<Self>) -> Self {
-        let _message_handler = rpc.subscribe_from_model(details.id, cx, Self::handle_message_sent);
+        let _subscription = rpc.subscribe_from_model(details.id, cx, Self::handle_message_sent);
 
         Self {
             details,
             rpc,
             first_message_id: None,
             messages: None,
-            _message_handler,
+            _subscription,
         }
     }
 
     fn handle_message_sent(
         &mut self,
-        message: &TypedEnvelope<ChannelMessageSent>,
+        message: TypedEnvelope<ChannelMessageSent>,
         rpc: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {

zed/src/main.rs

@@ -13,7 +13,6 @@ use zed::{
     workspace::{self, OpenParams},
     AppState,
 };
-use zrpc::ForegroundRouter;
 
 fn main() {
     init_logger();
@@ -31,8 +30,7 @@ fn main() {
         settings_tx: Arc::new(Mutex::new(settings_tx)),
         settings,
         themes,
-        rpc_router: Arc::new(ForegroundRouter::new()),
+        rpc: rpc::Client::new(),
-        rpc: rpc::Client::new(languages),
         fs: Arc::new(RealFs),
     };
 

zed/src/menus.rs

@@ -19,13 +19,13 @@ pub fn menus(state: &Arc<AppState>) -> Vec<Menu<'static>> {
                     name: "Share",
                     keystroke: None,
                     action: "workspace:share_worktree",
-                    arg: Some(Box::new(state.clone())),
+                    arg: None,
                 },
                 MenuItem::Action {
                     name: "Join",
                     keystroke: None,
                     action: "workspace:join_worktree",
-                    arg: Some(Box::new(state.clone())),
+                    arg: None,
                 },
                 MenuItem::Action {
                     name: "Quit",

zed/src/rpc.rs

@@ -1,15 +1,17 @@
-use crate::language::LanguageRegistry;
 use anyhow::{anyhow, Context, Result};
 use async_tungstenite::tungstenite::http::Request;
 use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
-use futures::StreamExt;
 use gpui::{AsyncAppContext, Entity, ModelContext, Task};
 use lazy_static::lazy_static;
-use smol::lock::RwLock;
+use parking_lot::RwLock;
-use std::time::Duration;
+use postage::prelude::Stream;
+use std::any::TypeId;
+use std::collections::HashMap;
+use std::sync::Weak;
+use std::time::{Duration, Instant};
 use std::{convert::TryFrom, future::Future, sync::Arc};
 use surf::Url;
-use zrpc::proto::EntityMessage;
+use zrpc::proto::{AnyTypedEnvelope, EntityMessage};
 pub use zrpc::{proto, ConnectionId, PeerId, TypedEnvelope};
 use zrpc::{
     proto::{EnvelopedMessage, RequestMessage},
@@ -24,22 +26,37 @@ lazy_static! {
 #[derive(Clone)]
 pub struct Client {
     peer: Arc<Peer>,
-    pub state: Arc<RwLock<ClientState>>,
+    state: Arc<RwLock<ClientState>>,
 }
 
+#[derive(Default)]
 pub struct ClientState {
     connection_id: Option<ConnectionId>,
-    pub languages: Arc<LanguageRegistry>,
+    entity_id_extractors: HashMap<TypeId, Box<dyn Send + Sync + Fn(&dyn AnyTypedEnvelope) -> u64>>,
+    model_handlers: HashMap<
+        (TypeId, u64),
+        Box<dyn Send + Sync + FnMut(Box<dyn AnyTypedEnvelope>, &mut AsyncAppContext)>,
+    >,
+}
+
+pub struct Subscription {
+    state: Weak<RwLock<ClientState>>,
+    id: (TypeId, u64),
+}
+
+impl Drop for Subscription {
+    fn drop(&mut self) {
+        if let Some(state) = self.state.upgrade() {
+            let _ = state.write().model_handlers.remove(&self.id).unwrap();
+        }
+    }
 }
 
 impl Client {
-    pub fn new(languages: Arc<LanguageRegistry>) -> Self {
+    pub fn new() -> Self {
         Self {
             peer: Peer::new(),
-            state: Arc::new(RwLock::new(ClientState {
+            state: Default::default(),
-                connection_id: None,
-                languages,
-            })),
         }
     }
 
@@ -48,31 +65,56 @@ impl Client {
         remote_id: u64,
         cx: &mut ModelContext<M>,
         mut handler: F,
-    ) -> Task<()>
+    ) -> Subscription
     where
         T: EntityMessage,
         M: Entity,
-        F: 'static + FnMut(&mut M, &TypedEnvelope<T>, Client, &mut ModelContext<M>) -> Result<()>,
+        F: 'static
+            + Send
+            + Sync
+            + FnMut(&mut M, TypedEnvelope<T>, Client, &mut ModelContext<M>) -> Result<()>,
     {
-        let rpc = self.clone();
+        let subscription_id = (TypeId::of::<T>(), remote_id);
-        let mut incoming = self.peer.subscribe::<T>();
+        let client = self.clone();
-        cx.spawn_weak(|model, mut cx| async move {
+        let mut state = self.state.write();
-            while let Some(envelope) = incoming.next().await {
+        let model = cx.handle().downgrade();
-                if envelope.payload.remote_entity_id() == remote_id {
+        state
-                    if let Some(model) = model.upgrade(&cx) {
+            .entity_id_extractors
-                        model.update(&mut cx, |model, cx| {
+            .entry(subscription_id.0)
-                            if let Err(error) = handler(model, &envelope, rpc.clone(), cx) {
+            .or_insert_with(|| {
+                Box::new(|envelope| {
+                    let envelope = envelope
+                        .as_any()
+                        .downcast_ref::<TypedEnvelope<T>>()
+                        .unwrap();
+                    envelope.payload.remote_entity_id()
+                })
+            });
+        let prev_handler = state.model_handlers.insert(
+            subscription_id,
+            Box::new(move |envelope, cx| {
+                if let Some(model) = model.upgrade(cx) {
+                    let envelope = envelope.into_any().downcast::<TypedEnvelope<T>>().unwrap();
+                    model.update(cx, |model, cx| {
+                        if let Err(error) = handler(model, *envelope, client.clone(), cx) {
                             log::error!("error handling message: {}", error)
                         }
                     });
                 }
+            }),
+        );
+        if prev_handler.is_some() {
+            panic!("registered a handler for the same entity twice")
         }
+
+        Subscription {
+            state: Arc::downgrade(&self.state),
+            id: subscription_id,
         }
-        })
     }
 
     pub async fn log_in_and_connect(&self, cx: AsyncAppContext) -> surf::Result<()> {
-        if self.state.read().await.connection_id.is_some() {
+        if self.state.read().connection_id.is_some() {
             return Ok(());
         }
 
@@ -110,8 +152,39 @@ impl Client {
             + Unpin
             + Send,
     {
-        let (connection_id, handle_io, handle_messages) = self.peer.add_connection(conn).await;
+        let (connection_id, handle_io, mut incoming) = self.peer.add_connection(conn).await;
-        cx.foreground().spawn(handle_messages).detach();
+        {
+            let mut cx = cx.clone();
+            let state = self.state.clone();
+            cx.foreground()
+                .spawn(async move {
+                    while let Some(message) = incoming.recv().await {
+                        let mut state = state.write();
+                        if let Some(extract_entity_id) =
+                            state.entity_id_extractors.get(&message.payload_type_id())
+                        {
+                            let entity_id = (extract_entity_id)(message.as_ref());
+                            if let Some(handler) = state
+                                .model_handlers
+                                .get_mut(&(message.payload_type_id(), entity_id))
+                            {
+                                let start_time = Instant::now();
+                                log::info!("RPC client message {}", message.payload_type_name());
+                                (handler)(message, &mut cx);
+                                log::info!(
+                                    "RPC message handled. duration:{:?}",
+                                    start_time.elapsed()
+                                );
+                            } else {
+                                log::info!("unhandled message {}", message.payload_type_name());
+                            }
+                        } else {
+                            log::info!("unhandled message {}", message.payload_type_name());
+                        }
+                    }
+                })
+                .detach();
+        }
         cx.background()
             .spawn(async move {
                 if let Err(error) = handle_io.await {
@@ -119,7 +192,7 @@ impl Client {
                 }
             })
             .detach();
-        self.state.write().await.connection_id = Some(connection_id);
+        self.state.write().connection_id = Some(connection_id);
         Ok(())
     }
 
@@ -200,27 +273,24 @@ impl Client {
     }
 
     pub async fn disconnect(&self) -> Result<()> {
-        let conn_id = self.connection_id().await?;
+        let conn_id = self.connection_id()?;
         self.peer.disconnect(conn_id).await;
         Ok(())
     }
 
-    async fn connection_id(&self) -> Result<ConnectionId> {
+    fn connection_id(&self) -> Result<ConnectionId> {
         self.state
             .read()
-            .await
             .connection_id
             .ok_or_else(|| anyhow!("not connected"))
     }
 
     pub async fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
-        self.peer.send(self.connection_id().await?, message).await
+        self.peer.send(self.connection_id()?, message).await
     }
 
     pub async fn request<T: RequestMessage>(&self, request: T) -> Result<T::Response> {
-        self.peer
+        self.peer.request(self.connection_id()?, request).await
-            .request(self.connection_id().await?, request)
-            .await
     }
 
     pub fn respond<T: RequestMessage>(

zed/src/test.rs

@@ -162,7 +162,7 @@ pub fn build_app_state(cx: &AppContext) -> Arc<AppState> {
         settings,
         themes,
         languages: languages.clone(),
-        rpc: rpc::Client::new(languages),
+        rpc: rpc::Client::new(),
         fs: Arc::new(RealFs),
     })
 }

zed/src/util.rs

@@ -82,15 +82,13 @@ impl<T: Rng> Iterator for RandomCharIter<T> {
     }
 }
 
-pub async fn log_async_errors<F>(f: F) -> impl Future<Output = ()>
+pub async fn log_async_errors<F>(f: F)
 where
     F: Future<Output = anyhow::Result<()>>,
 {
-    async {
     if let Err(error) = f.await {
         log::error!("{}", error)
     }
-    }
 }
 
 #[cfg(test)]

zed/src/workspace.rs

@@ -108,7 +108,7 @@ fn open_new(app_state: &Arc<AppState>, cx: &mut MutableAppContext) {
 fn join_worktree(app_state: &Arc<AppState>, cx: &mut MutableAppContext) {
     cx.add_window(|cx| {
         let mut view = Workspace::new(app_state.as_ref(), cx);
-        view.join_worktree(&app_state, cx);
+        view.join_worktree(&(), cx);
         view
     });
 }
@@ -725,7 +725,7 @@ impl Workspace {
         };
     }
 
-    fn share_worktree(&mut self, app_state: &Arc<AppState>, cx: &mut ViewContext<Self>) {
+    fn share_worktree(&mut self, _: &(), cx: &mut ViewContext<Self>) {
         let rpc = self.rpc.clone();
         let platform = cx.platform();
 
@@ -757,7 +757,7 @@ impl Workspace {
         .detach();
     }
 
-    fn join_worktree(&mut self, app_state: &Arc<AppState>, cx: &mut ViewContext<Self>) {
+    fn join_worktree(&mut self, _: &(), cx: &mut ViewContext<Self>) {
         let rpc = self.rpc.clone();
         let languages = self.languages.clone();
 

zed/src/worktree.rs

@@ -213,7 +213,7 @@ impl Worktree {
                     .detach();
                 }
 
-                let _message_handlers = vec![
+                let _subscriptions = vec![
                     rpc.subscribe_from_model(remote_id, cx, Self::handle_add_peer),
                     rpc.subscribe_from_model(remote_id, cx, Self::handle_remove_peer),
                     rpc.subscribe_from_model(remote_id, cx, Self::handle_update),
@@ -234,7 +234,7 @@ impl Worktree {
                         .map(|p| (PeerId(p.peer_id), p.replica_id as ReplicaId))
                         .collect(),
                     languages,
-                    _message_handlers,
+                    _subscriptions,
                 })
             })
         });
@@ -282,7 +282,7 @@ impl Worktree {
 
     pub fn handle_add_peer(
         &mut self,
-        envelope: &TypedEnvelope<proto::AddPeer>,
+        envelope: TypedEnvelope<proto::AddPeer>,
         _: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
@@ -294,7 +294,7 @@ impl Worktree {
 
     pub fn handle_remove_peer(
         &mut self,
-        envelope: &TypedEnvelope<proto::RemovePeer>,
+        envelope: TypedEnvelope<proto::RemovePeer>,
         _: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
@@ -306,7 +306,7 @@ impl Worktree {
 
     pub fn handle_update(
         &mut self,
-        envelope: &TypedEnvelope<proto::UpdateWorktree>,
+        envelope: TypedEnvelope<proto::UpdateWorktree>,
         _: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> anyhow::Result<()> {
@@ -317,7 +317,7 @@ impl Worktree {
 
     pub fn handle_open_buffer(
         &mut self,
-        envelope: &TypedEnvelope<proto::OpenBuffer>,
+        envelope: TypedEnvelope<proto::OpenBuffer>,
         rpc: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> anyhow::Result<()> {
@@ -340,7 +340,7 @@ impl Worktree {
 
     pub fn handle_close_buffer(
         &mut self,
-        envelope: &TypedEnvelope<proto::CloseBuffer>,
+        envelope: TypedEnvelope<proto::CloseBuffer>,
         _: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> anyhow::Result<()> {
@@ -396,7 +396,7 @@ impl Worktree {
 
     pub fn handle_update_buffer(
         &mut self,
-        envelope: &TypedEnvelope<proto::UpdateBuffer>,
+        envelope: TypedEnvelope<proto::UpdateBuffer>,
         _: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
@@ -443,7 +443,7 @@ impl Worktree {
 
     pub fn handle_save_buffer(
         &mut self,
-        envelope: &TypedEnvelope<proto::SaveBuffer>,
+        envelope: TypedEnvelope<proto::SaveBuffer>,
         rpc: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
@@ -485,7 +485,7 @@ impl Worktree {
 
     pub fn handle_buffer_saved(
         &mut self,
-        envelope: &TypedEnvelope<proto::BufferSaved>,
+        envelope: TypedEnvelope<proto::BufferSaved>,
         _: rpc::Client,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
@@ -791,7 +791,7 @@ impl LocalWorktree {
 
     pub fn open_remote_buffer(
         &mut self,
-        envelope: &TypedEnvelope<proto::OpenBuffer>,
+        envelope: TypedEnvelope<proto::OpenBuffer>,
         cx: &mut ModelContext<Worktree>,
     ) -> Task<Result<proto::OpenBufferResponse>> {
         let peer_id = envelope.original_sender_id();
@@ -818,11 +818,12 @@ impl LocalWorktree {
 
     pub fn close_remote_buffer(
         &mut self,
-        envelope: &TypedEnvelope<proto::CloseBuffer>,
+        envelope: TypedEnvelope<proto::CloseBuffer>,
-        _: &mut ModelContext<Worktree>,
+        cx: &mut ModelContext<Worktree>,
     ) -> Result<()> {
         if let Some(shared_buffers) = self.shared_buffers.get_mut(&envelope.original_sender_id()?) {
             shared_buffers.remove(&envelope.payload.buffer_id);
+            cx.notify();
         }
 
         Ok(())
@@ -830,7 +831,7 @@ impl LocalWorktree {
 
     pub fn add_peer(
         &mut self,
-        envelope: &TypedEnvelope<proto::AddPeer>,
+        envelope: TypedEnvelope<proto::AddPeer>,
         cx: &mut ModelContext<Worktree>,
     ) -> Result<()> {
         let peer = envelope
@@ -847,7 +848,7 @@ impl LocalWorktree {
 
     pub fn remove_peer(
         &mut self,
-        envelope: &TypedEnvelope<proto::RemovePeer>,
+        envelope: TypedEnvelope<proto::RemovePeer>,
         cx: &mut ModelContext<Worktree>,
     ) -> Result<()> {
         let peer_id = PeerId(envelope.payload.peer_id);
@@ -994,7 +995,7 @@ impl LocalWorktree {
                 .detach();
 
             this.update(&mut cx, |worktree, cx| {
-                let _message_handlers = vec![
+                let _subscriptions = vec![
                     rpc.subscribe_from_model(remote_id, cx, Worktree::handle_add_peer),
                     rpc.subscribe_from_model(remote_id, cx, Worktree::handle_remove_peer),
                     rpc.subscribe_from_model(remote_id, cx, Worktree::handle_open_buffer),
@@ -1008,7 +1009,7 @@ impl LocalWorktree {
                     rpc,
                     remote_id: share_response.worktree_id,
                     snapshots_tx: snapshots_to_send_tx,
-                    _message_handlers,
+                    _subscriptions,
                 });
             });
 
@@ -1068,7 +1069,7 @@ struct ShareState {
     rpc: rpc::Client,
     remote_id: u64,
     snapshots_tx: Sender<Snapshot>,
-    _message_handlers: Vec<Task<()>>,
+    _subscriptions: Vec<rpc::Subscription>,
 }
 
 pub struct RemoteWorktree {
@@ -1081,7 +1082,7 @@ pub struct RemoteWorktree {
     open_buffers: HashMap<usize, RemoteBuffer>,
     peers: HashMap<PeerId, ReplicaId>,
     languages: Arc<LanguageRegistry>,
-    _message_handlers: Vec<Task<()>>,
+    _subscriptions: Vec<rpc::Subscription>,
 }
 
 impl RemoteWorktree {
@@ -1151,7 +1152,7 @@ impl RemoteWorktree {
 
     fn update_from_remote(
         &mut self,
-        envelope: &TypedEnvelope<proto::UpdateWorktree>,
+        envelope: TypedEnvelope<proto::UpdateWorktree>,
         cx: &mut ModelContext<Worktree>,
     ) -> Result<()> {
         let mut tx = self.updates_tx.clone();
@@ -1167,7 +1168,7 @@ impl RemoteWorktree {
 
     pub fn add_peer(
         &mut self,
-        envelope: &TypedEnvelope<proto::AddPeer>,
+        envelope: TypedEnvelope<proto::AddPeer>,
         cx: &mut ModelContext<Worktree>,
     ) -> Result<()> {
         let peer = envelope
@@ -1183,7 +1184,7 @@ impl RemoteWorktree {
 
     pub fn remove_peer(
         &mut self,
-        envelope: &TypedEnvelope<proto::RemovePeer>,
+        envelope: TypedEnvelope<proto::RemovePeer>,
         cx: &mut ModelContext<Worktree>,
     ) -> Result<()> {
         let peer_id = PeerId(envelope.payload.peer_id);
@@ -2761,7 +2762,7 @@ mod tests {
                 replica_id: 1,
                 peers: Vec::new(),
             },
-            rpc::Client::new(Default::default()),
+            rpc::Client::new(),
             Default::default(),
             &mut cx.to_async(),
         )

zrpc/src/peer.rs

@@ -1,4 +1,4 @@
-use crate::proto::{self, EnvelopedMessage, MessageStream, RequestMessage};
+use crate::proto::{self, AnyTypedEnvelope, EnvelopedMessage, MessageStream, RequestMessage};
 use anyhow::{anyhow, Context, Result};
 use async_lock::{Mutex, RwLock};
 use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
@@ -8,7 +8,6 @@ use postage::{
     prelude::{Sink as _, Stream as _},
 };
 use std::{
-    any::Any,
     collections::HashMap,
     fmt,
     future::Future,
@@ -105,7 +104,7 @@ impl Peer {
     ) -> (
         ConnectionId,
         impl Future<Output = anyhow::Result<()>> + Send,
-        mpsc::Receiver<Box<dyn Any + Sync + Send>>,
+        mpsc::Receiver<Box<dyn AnyTypedEnvelope>>,
     )
     where
         Conn: futures::Sink<WebSocketMessage, Error = WebSocketError>
@@ -409,10 +408,11 @@ mod tests {
             client2.disconnect(client1_conn_id).await;
 
             async fn handle_messages(
-                mut messages: mpsc::Receiver<Box<dyn Any + Sync + Send>>,
+                mut messages: mpsc::Receiver<Box<dyn AnyTypedEnvelope>>,
                 peer: Arc<Peer>,
             ) -> Result<()> {
                 while let Some(envelope) = messages.next().await {
+                    let envelope = envelope.into_any();
                     if let Some(envelope) = envelope.downcast_ref::<TypedEnvelope<proto::Ping>>() {
                         let receipt = envelope.receipt();
                         peer.respond(

zrpc/src/proto.rs

@@ -3,7 +3,7 @@ use anyhow::Result;
 use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
 use futures::{SinkExt as _, StreamExt as _};
 use prost::Message;
-use std::any::Any;
+use std::any::{Any, TypeId};
 use std::{
     io,
     time::{Duration, SystemTime, UNIX_EPOCH},
@@ -31,9 +31,34 @@ pub trait RequestMessage: EnvelopedMessage {
     type Response: EnvelopedMessage;
 }
 
+pub trait AnyTypedEnvelope: 'static + Send + Sync {
+    fn payload_type_id(&self) -> TypeId;
+    fn payload_type_name(&self) -> &'static str;
+    fn as_any(&self) -> &dyn Any;
+    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync>;
+}
+
+impl<T: EnvelopedMessage> AnyTypedEnvelope for TypedEnvelope<T> {
+    fn payload_type_id(&self) -> TypeId {
+        TypeId::of::<T>()
+    }
+
+    fn payload_type_name(&self) -> &'static str {
+        T::NAME
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync> {
+        self
+    }
+}
+
 macro_rules! messages {
     ($($name:ident),* $(,)?) => {
-        pub fn build_typed_envelope(sender_id: ConnectionId, envelope: Envelope) -> Option<Box<dyn Any + Send + Sync>> {
+        pub fn build_typed_envelope(sender_id: ConnectionId, envelope: Envelope) -> Option<Box<dyn AnyTypedEnvelope>> {
             match envelope.payload {
                 $(Some(envelope::Payload::$name(payload)) => {
                     Some(Box::new(TypedEnvelope {