Buffer messages in Client while no entity is listening to them

This commit is contained in:
Antonio Scandurra 2022-02-16 11:49:37 +01:00
parent 71abea728e
commit 978dae201c
3 changed files with 216 additions and 132 deletions

View file

@ -190,7 +190,7 @@ impl Channel {
rpc: Arc<Client>, rpc: Arc<Client>,
cx: &mut ModelContext<Self>, cx: &mut ModelContext<Self>,
) -> Self { ) -> Self {
let _subscription = rpc.add_model_for_remote_entity(cx.handle(), details.id); let _subscription = rpc.add_model_for_remote_entity(details.id, cx);
{ {
let user_store = user_store.clone(); let user_store = user_store.clone();

View file

@ -13,7 +13,7 @@ use async_tungstenite::tungstenite::{
}; };
use futures::{future::LocalBoxFuture, FutureExt, StreamExt}; use futures::{future::LocalBoxFuture, FutureExt, StreamExt};
use gpui::{ use gpui::{
action, AnyModelHandle, AnyWeakModelHandle, AsyncAppContext, Entity, ModelHandle, action, AnyModelHandle, AnyWeakModelHandle, AsyncAppContext, Entity, ModelContext, ModelHandle,
MutableAppContext, Task, MutableAppContext, Task,
}; };
use http::HttpClient; use http::HttpClient;
@ -140,7 +140,7 @@ struct ClientState {
model_types_by_message_type: HashMap<TypeId, TypeId>, model_types_by_message_type: HashMap<TypeId, TypeId>,
message_handlers: HashMap< message_handlers: HashMap<
TypeId, TypeId,
Box< Arc<
dyn Send dyn Send
+ Sync + Sync
+ Fn( + Fn(
@ -175,16 +175,33 @@ impl Default for ClientState {
} }
} }
pub struct Subscription { pub enum Subscription {
client: Weak<Client>, Entity {
id: (TypeId, u64), client: Weak<Client>,
id: (TypeId, u64),
},
Message {
client: Weak<Client>,
id: TypeId,
},
} }
impl Drop for Subscription { impl Drop for Subscription {
fn drop(&mut self) { fn drop(&mut self) {
if let Some(client) = self.client.upgrade() { match self {
let mut state = client.state.write(); Subscription::Entity { client, id } => {
let _ = state.models_by_entity_type_and_remote_id.remove(&self.id); if let Some(client) = client.upgrade() {
let mut state = client.state.write();
let _ = state.models_by_entity_type_and_remote_id.remove(id);
}
}
Subscription::Message { client, id } => {
if let Some(client) = client.upgrade() {
let mut state = client.state.write();
let _ = state.model_types_by_message_type.remove(id);
let _ = state.message_handlers.remove(id);
}
}
} }
} }
} }
@ -285,21 +302,66 @@ impl Client {
pub fn add_model_for_remote_entity<T: Entity>( pub fn add_model_for_remote_entity<T: Entity>(
self: &Arc<Self>, self: &Arc<Self>,
handle: ModelHandle<T>,
remote_id: u64, remote_id: u64,
cx: &mut ModelContext<T>,
) -> Subscription { ) -> Subscription {
let handle = AnyModelHandle::from(cx.handle());
let mut state = self.state.write(); let mut state = self.state.write();
let id = (TypeId::of::<T>(), remote_id); let id = (TypeId::of::<T>(), remote_id);
state state
.models_by_entity_type_and_remote_id .models_by_entity_type_and_remote_id
.insert(id, AnyModelHandle::from(handle).downgrade()); .insert(id, handle.downgrade());
Subscription { let pending_messages = state.pending_messages.remove(&id);
drop(state);
let client_id = self.id;
for message in pending_messages.into_iter().flatten() {
let type_id = message.payload_type_id();
let type_name = message.payload_type_name();
let state = self.state.read();
if let Some(handler) = state.message_handlers.get(&type_id).cloned() {
let future = (handler)(handle.clone(), message, cx.to_async());
drop(state);
log::debug!(
"deferred rpc message received. client_id:{}, name:{}",
client_id,
type_name
);
cx.foreground()
.spawn(async move {
match future.await {
Ok(()) => {
log::debug!(
"deferred rpc message handled. client_id:{}, name:{}",
client_id,
type_name
);
}
Err(error) => {
log::error!(
"error handling deferred message. client_id:{}, name:{}, {}",
client_id,
type_name,
error
);
}
}
})
.detach();
}
}
Subscription::Entity {
client: Arc::downgrade(self), client: Arc::downgrade(self),
id, id,
} }
} }
pub fn add_message_handler<M, E, H, F>(self: &Arc<Self>, model: ModelHandle<E>, handler: H) pub fn add_message_handler<M, E, H, F>(
self: &Arc<Self>,
model: ModelHandle<E>,
handler: H,
) -> Subscription
where where
M: EnvelopedMessage, M: EnvelopedMessage,
E: Entity, E: Entity,
@ -319,7 +381,7 @@ impl Client {
let prev_handler = state.message_handlers.insert( let prev_handler = state.message_handlers.insert(
message_type_id, message_type_id,
Box::new(move |handle, envelope, cx| { Arc::new(move |handle, envelope, cx| {
let model = handle.downcast::<E>().unwrap(); let model = handle.downcast::<E>().unwrap();
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap(); let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
handler(model, *envelope, client.clone(), cx).boxed_local() handler(model, *envelope, client.clone(), cx).boxed_local()
@ -328,6 +390,11 @@ impl Client {
if prev_handler.is_some() { if prev_handler.is_some() {
panic!("registered handler for the same message twice"); panic!("registered handler for the same message twice");
} }
Subscription::Message {
client: Arc::downgrade(self),
id: message_type_id,
}
} }
pub fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H) pub fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
@ -363,7 +430,7 @@ impl Client {
let prev_handler = state.message_handlers.insert( let prev_handler = state.message_handlers.insert(
message_type_id, message_type_id,
Box::new(move |handle, envelope, cx| { Arc::new(move |handle, envelope, cx| {
let model = handle.downcast::<E>().unwrap(); let model = handle.downcast::<E>().unwrap();
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap(); let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
handler(model, *envelope, client.clone(), cx).boxed_local() handler(model, *envelope, client.clone(), cx).boxed_local()
@ -501,37 +568,54 @@ impl Client {
let mut state = this.state.write(); let mut state = this.state.write();
let payload_type_id = message.payload_type_id(); let payload_type_id = message.payload_type_id();
let type_name = message.payload_type_name(); let type_name = message.payload_type_name();
let model_type_id = state
.model_types_by_message_type
.get(&payload_type_id)
.copied();
let entity_id = state
.entity_id_extractors
.get(&message.payload_type_id())
.map(|extract_entity_id| (extract_entity_id)(message.as_ref()));
let model = state.models_by_message_type.get(&payload_type_id).cloned().or_else(|| { let model = state
let extract_entity_id = state.entity_id_extractors.get(&message.payload_type_id())?; .models_by_message_type
let entity_id = (extract_entity_id)(message.as_ref()); .get(&payload_type_id)
let model_type_id = *state.model_types_by_message_type.get(&payload_type_id)?; .cloned()
.or_else(|| {
// TODO - if we don't have this model yet, then buffer the message let model_type_id = model_type_id?;
let model = state.models_by_entity_type_and_remote_id.get(&(model_type_id, entity_id))?; let entity_id = entity_id?;
let model = state
if let Some(model) = model.upgrade(&cx) { .models_by_entity_type_and_remote_id
Some(model) .get(&(model_type_id, entity_id))?;
} else { if let Some(model) = model.upgrade(&cx) {
state.models_by_entity_type_and_remote_id.remove(&(model_type_id, entity_id)); Some(model)
None } else {
} state
}); .models_by_entity_type_and_remote_id
.remove(&(model_type_id, entity_id));
None
}
});
let model = if let Some(model) = model { let model = if let Some(model) = model {
model model
} else { } else {
log::info!("unhandled message {}", type_name); log::info!("unhandled message {}", type_name);
if let Some((model_type_id, entity_id)) = model_type_id.zip(entity_id) {
state
.pending_messages
.entry((model_type_id, entity_id))
.or_default()
.push(message);
}
continue; continue;
}; };
if let Some(handler) = state.message_handlers.remove(&payload_type_id) { if let Some(handler) = state.message_handlers.get(&payload_type_id).cloned()
{
drop(state); // Avoid deadlocks if the handler interacts with rpc::Client drop(state); // Avoid deadlocks if the handler interacts with rpc::Client
let future = handler(model, message, cx.clone()); let future = handler(model, message, cx.clone());
{
let mut state = this.state.write();
state.message_handlers.insert(payload_type_id, handler);
}
let client_id = this.id; let client_id = this.id;
log::debug!( log::debug!(
@ -551,7 +635,7 @@ impl Client {
} }
Err(error) => { Err(error) => {
log::error!( log::error!(
"error handling rpc message. client_id:{}, name:{}, error:{}", "error handling message. client_id:{}, name:{}, {}",
client_id, client_id,
type_name, type_name,
error error
@ -926,109 +1010,111 @@ mod tests {
assert_eq!(decode_worktree_url("not://the-right-format"), None); assert_eq!(decode_worktree_url("not://the-right-format"), None);
} }
// #[gpui::test] #[gpui::test]
// async fn test_subscribing_to_entity(mut cx: TestAppContext) { async fn test_subscribing_to_entity(mut cx: TestAppContext) {
// cx.foreground().forbid_parking(); cx.foreground().forbid_parking();
// let user_id = 5; let user_id = 5;
// let mut client = Client::new(FakeHttpClient::with_404_response()); let mut client = Client::new(FakeHttpClient::with_404_response());
// let server = FakeServer::for_client(user_id, &mut client, &cx).await; let server = FakeServer::for_client(user_id, &mut client, &cx).await;
// let model = cx.add_model(|_| Model { subscription: None }); let (done_tx1, mut done_rx1) = smol::channel::unbounded();
// let (mut done_tx1, mut done_rx1) = postage::oneshot::channel(); let (done_tx2, mut done_rx2) = smol::channel::unbounded();
// let (mut done_tx2, mut done_rx2) = postage::oneshot::channel(); client.add_entity_message_handler(
// let _subscription1 = model.update(&mut cx, |_, cx| { move |model: ModelHandle<Model>, _: TypedEnvelope<proto::UnshareProject>, _, cx| {
// client.add_entity_message_handler( match model.read_with(&cx, |model, _| model.id) {
// 1, 1 => done_tx1.try_send(()).unwrap(),
// cx, 2 => done_tx2.try_send(()).unwrap(),
// move |_, _: TypedEnvelope<proto::UnshareProject>, _, _| { _ => unreachable!(),
// postage::sink::Sink::try_send(&mut done_tx1, ()).unwrap(); }
// async { Ok(()) } async { Ok(()) }
// }, },
// ) );
// }); let model1 = cx.add_model(|_| Model {
// let _subscription2 = model.update(&mut cx, |_, cx| { id: 1,
// client.add_entity_message_handler( subscription: None,
// 2, });
// cx, let model2 = cx.add_model(|_| Model {
// move |_, _: TypedEnvelope<proto::UnshareProject>, _, _| { id: 2,
// postage::sink::Sink::try_send(&mut done_tx2, ()).unwrap(); subscription: None,
// async { Ok(()) } });
// }, let model3 = cx.add_model(|_| Model {
// ) id: 3,
// }); subscription: None,
});
// // Ensure dropping a subscription for the same entity type still allows receiving of let _subscription1 =
// // messages for other entity IDs of the same type. model1.update(&mut cx, |_, cx| client.add_model_for_remote_entity(1, cx));
// let subscription3 = model.update(&mut cx, |_, cx| { let _subscription2 =
// client.add_entity_message_handler( model2.update(&mut cx, |_, cx| client.add_model_for_remote_entity(2, cx));
// 3, // Ensure dropping a subscription for the same entity type still allows receiving of
// cx, // messages for other entity IDs of the same type.
// |_, _: TypedEnvelope<proto::UnshareProject>, _, _| async { Ok(()) }, let subscription3 =
// ) model3.update(&mut cx, |_, cx| client.add_model_for_remote_entity(3, cx));
// }); drop(subscription3);
// drop(subscription3);
// server.send(proto::UnshareProject { project_id: 1 }); server.send(proto::UnshareProject { project_id: 1 });
// server.send(proto::UnshareProject { project_id: 2 }); server.send(proto::UnshareProject { project_id: 2 });
// done_rx1.next().await.unwrap(); done_rx1.next().await.unwrap();
// done_rx2.next().await.unwrap(); done_rx2.next().await.unwrap();
// } }
// #[gpui::test] #[gpui::test]
// async fn test_subscribing_after_dropping_subscription(mut cx: TestAppContext) { async fn test_subscribing_after_dropping_subscription(mut cx: TestAppContext) {
// cx.foreground().forbid_parking(); cx.foreground().forbid_parking();
// let user_id = 5; let user_id = 5;
// let mut client = Client::new(FakeHttpClient::with_404_response()); let mut client = Client::new(FakeHttpClient::with_404_response());
// let server = FakeServer::for_client(user_id, &mut client, &cx).await; let server = FakeServer::for_client(user_id, &mut client, &cx).await;
// let model = cx.add_model(|_| Model { subscription: None }); let model = cx.add_model(|_| Model::default());
// let (mut done_tx1, _done_rx1) = postage::oneshot::channel(); let (done_tx1, _done_rx1) = smol::channel::unbounded();
// let (mut done_tx2, mut done_rx2) = postage::oneshot::channel(); let (done_tx2, mut done_rx2) = smol::channel::unbounded();
// let subscription1 = model.update(&mut cx, |_, cx| { let subscription1 = client.add_message_handler(
// client.add_message_handler(cx, move |_, _: TypedEnvelope<proto::Ping>, _, _| { model.clone(),
// postage::sink::Sink::try_send(&mut done_tx1, ()).unwrap(); move |_, _: TypedEnvelope<proto::Ping>, _, _| {
// async { Ok(()) } done_tx1.try_send(()).unwrap();
// }) async { Ok(()) }
// }); },
// drop(subscription1); );
// let _subscription2 = model.update(&mut cx, |_, cx| { drop(subscription1);
// client.add_message_handler(cx, move |_, _: TypedEnvelope<proto::Ping>, _, _| { let _subscription2 =
// postage::sink::Sink::try_send(&mut done_tx2, ()).unwrap(); client.add_message_handler(model, move |_, _: TypedEnvelope<proto::Ping>, _, _| {
// async { Ok(()) } done_tx2.try_send(()).unwrap();
// }) async { Ok(()) }
// }); });
// server.send(proto::Ping {}); server.send(proto::Ping {});
// done_rx2.next().await.unwrap(); done_rx2.next().await.unwrap();
// } }
// #[gpui::test] #[gpui::test]
// async fn test_dropping_subscription_in_handler(mut cx: TestAppContext) { async fn test_dropping_subscription_in_handler(mut cx: TestAppContext) {
// cx.foreground().forbid_parking(); cx.foreground().forbid_parking();
// let user_id = 5; let user_id = 5;
// let mut client = Client::new(FakeHttpClient::with_404_response()); let mut client = Client::new(FakeHttpClient::with_404_response());
// let server = FakeServer::for_client(user_id, &mut client, &cx).await; let server = FakeServer::for_client(user_id, &mut client, &cx).await;
// let model = cx.add_model(|_| Model { subscription: None }); let model = cx.add_model(|_| Model::default());
// let (mut done_tx, mut done_rx) = postage::oneshot::channel(); let (done_tx, mut done_rx) = smol::channel::unbounded();
// client.add_message_handler( let subscription = client.add_message_handler(
// model.clone(), model.clone(),
// move |model, _: TypedEnvelope<proto::Ping>, _, mut cx| { move |model, _: TypedEnvelope<proto::Ping>, _, mut cx| {
// model.update(&mut cx, |model, _| model.subscription.take()); model.update(&mut cx, |model, _| model.subscription.take());
// postage::sink::Sink::try_send(&mut done_tx, ()).unwrap(); done_tx.try_send(()).unwrap();
// async { Ok(()) } async { Ok(()) }
// }, },
// ); );
// model.update(&mut cx, |model, cx| { model.update(&mut cx, |model, _| {
// model.subscription = Some(); model.subscription = Some(subscription);
// }); });
// server.send(proto::Ping {}); server.send(proto::Ping {});
// done_rx.next().await.unwrap(); done_rx.next().await.unwrap();
// } }
#[derive(Default)]
struct Model { struct Model {
id: usize,
subscription: Option<Subscription>, subscription: Option<Subscription>,
} }

View file

@ -312,7 +312,7 @@ impl Project {
languages, languages,
user_store, user_store,
fs, fs,
subscriptions: vec![client.add_model_for_remote_entity(cx.handle(), remote_id)], subscriptions: vec![client.add_model_for_remote_entity(remote_id, cx)],
client, client,
client_state: ProjectClientState::Remote { client_state: ProjectClientState::Remote {
sharing_has_stopped: false, sharing_has_stopped: false,
@ -349,10 +349,8 @@ impl Project {
self.subscriptions.clear(); self.subscriptions.clear();
if let Some(remote_id) = remote_id { if let Some(remote_id) = remote_id {
self.subscriptions.push( self.subscriptions
self.client .push(self.client.add_model_for_remote_entity(remote_id, cx));
.add_model_for_remote_entity(cx.handle(), remote_id),
);
} }
} }