zed-industries/zed
1
1
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2024-09-16 17:07:14 +03:00
Code Issues Packages Projects Releases Wiki Activity

Consolidate logic for protobuf message handling between ssh and web socket clients (#17185)

Browse Source 
This is a refactor to prepare for adding LSP support in SSH remote
projects.

Release Notes:

- N/A

---------

Co-authored-by: Mikayla <mikayla@zed.dev>
Co-authored-by: Conrad <conrad@zed.dev>
This commit is contained in:
Max Brunsfeld 2024-09-01 10:14:21 -07:00 committed by GitHub
parent 144793bf16
commit b8e6098f60
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 1002 additions and 963 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.lock generated
View File

@ -8512,6 +8512,8 @@ dependencies = [
"anyhow",
"collections",
"futures 0.3.30",
"gpui",
"parking_lot",
"prost",
"prost-build",
"serde",

2
crates/assistant/src/assistant.rs
View File

@ -209,7 +209,7 @@ pub fn init(
})
.detach();
context_store::init(&client);
context_store::init(&client.clone().into());
prompt_library::init(cx);
init_language_model_settings(cx);
assistant_slash_command::init(cx);

3
crates/assistant/src/context_store.rs
View File

@ -2,6 +2,7 @@ use crate::{
prompts::PromptBuilder, Context, ContextEvent, ContextId, ContextOperation, ContextVersion,
SavedContext, SavedContextMetadata,
};
use ::proto::AnyProtoClient;
use anyhow::{anyhow, Context as _, Result};
use client::{proto, telemetry::Telemetry, Client, TypedEnvelope};
use clock::ReplicaId;
@ -25,7 +26,7 @@ use std::{
};
use util::{ResultExt, TryFutureExt};
pub fn init(client: &Arc<Client>) {
pub fn init(client: &AnyProtoClient) {
client.add_model_message_handler(ContextStore::handle_advertise_contexts);
client.add_model_request_handler(ContextStore::handle_open_context);
client.add_model_request_handler(ContextStore::handle_create_context);

4
crates/channel/src/channel.rs
View File

@ -18,6 +18,6 @@ mod channel_store_tests;
pub fn init(client: &Arc<Client>, user_store: Model<UserStore>, cx: &mut AppContext) {
channel_store::init(client, user_store, cx);
channel_buffer::init(client);
channel_chat::init(client);
channel_buffer::init(&client.clone().into());
channel_chat::init(&client.clone().into());
}

4
crates/channel/src/channel_buffer.rs
View File

@ -5,7 +5,7 @@ use collections::HashMap;
use gpui::{AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, Task};
use language::proto::serialize_version;
use rpc::{
proto::{self, PeerId},
proto::{self, AnyProtoClient, PeerId},
TypedEnvelope,
};
use std::{sync::Arc, time::Duration};
@ -14,7 +14,7 @@ use util::ResultExt;
pub const ACKNOWLEDGE_DEBOUNCE_INTERVAL: Duration = Duration::from_millis(250);
pub(crate) fn init(client: &Arc<Client>) {
pub(crate) fn init(client: &AnyProtoClient) {
client.add_model_message_handler(ChannelBuffer::handle_update_channel_buffer);
client.add_model_message_handler(ChannelBuffer::handle_update_channel_buffer_collaborators);
}

3
crates/channel/src/channel_chat.rs
View File

@ -11,6 +11,7 @@ use gpui::{
AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, Task, WeakModel,
};
use rand::prelude::*;
use rpc::proto::AnyProtoClient;
use std::{
ops::{ControlFlow, Range},
sync::Arc,
@ -95,7 +96,7 @@ pub enum ChannelChatEvent {
}
impl EventEmitter<ChannelChatEvent> for ChannelChat {}
pub fn init(client: &Arc<Client>) {
pub fn init(client: &AnyProtoClient) {
client.add_model_message_handler(ChannelChat::handle_message_sent);
client.add_model_message_handler(ChannelChat::handle_message_removed);
client.add_model_message_handler(ChannelChat::handle_message_updated);

305
crates/client/src/client.rs
View File

@ -14,22 +14,18 @@ use async_tungstenite::tungstenite::{
};
use chrono::{DateTime, Utc};
use clock::SystemClock;
use collections::HashMap;
use futures::{
channel::oneshot,
future::{BoxFuture, LocalBoxFuture},
AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,
};
use gpui::{
actions, AnyModel, AnyWeakModel, AppContext, AsyncAppContext, Global, Model, Task, WeakModel,
channel::oneshot, future::BoxFuture, AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt,
TryFutureExt as _, TryStreamExt,
};
use gpui::{actions, AppContext, AsyncAppContext, Global, Model, Task, WeakModel};
use http_client::{AsyncBody, HttpClient, HttpClientWithUrl};
use parking_lot::RwLock;
use postage::watch;
use proto::ProtoClient;
use proto::{AnyProtoClient, EntityMessageSubscriber, ProtoClient, ProtoMessageHandlerSet};
use rand::prelude::*;
use release_channel::{AppVersion, ReleaseChannel};
use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, PeerId, RequestMessage};
use rpc::proto::{AnyTypedEnvelope, EnvelopedMessage, PeerId, RequestMessage};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use settings::{Settings, SettingsSources};
@ -208,6 +204,7 @@ pub struct Client {
telemetry: Arc<Telemetry>,
credentials_provider: Arc<dyn CredentialsProvider + Send + Sync + 'static>,
state: RwLock<ClientState>,
handler_set: parking_lot::Mutex<ProtoMessageHandlerSet>,
#[allow(clippy::type_complexity)]
#[cfg(any(test, feature = "test-support"))]
@ -304,30 +301,7 @@ impl Status {
struct ClientState {
credentials: Option<Credentials>,
status: (watch::Sender<Status>, watch::Receiver<Status>),
entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
_reconnect_task: Option<Task<()>>,
entities_by_type_and_remote_id: HashMap<(TypeId, u64), WeakSubscriber>,
models_by_message_type: HashMap<TypeId, AnyWeakModel>,
entity_types_by_message_type: HashMap<TypeId, TypeId>,
#[allow(clippy::type_complexity)]
message_handlers: HashMap<
TypeId,
Arc<
dyn Send
+ Sync
+ Fn(
AnyModel,
Box<dyn AnyTypedEnvelope>,
&Arc<Client>,
AsyncAppContext,
) -> LocalBoxFuture<'static, Result<()>>,
>,
>,
}
enum WeakSubscriber {
Entity { handle: AnyWeakModel },
Pending(Vec<Box<dyn AnyTypedEnvelope>>),
}
#[derive(Clone, Debug, Eq, PartialEq)]
@ -379,12 +353,7 @@ impl Default for ClientState {
Self {
credentials: None,
status: watch::channel_with(Status::SignedOut),
entity_id_extractors: Default::default(),
_reconnect_task: None,
models_by_message_type: Default::default(),
entities_by_type_and_remote_id: Default::default(),
entity_types_by_message_type: Default::default(),
message_handlers: Default::default(),
}
}
}
@ -405,13 +374,13 @@ impl Drop for Subscription {
match self {
Subscription::Entity { client, id } => {
if let Some(client) = client.upgrade() {
let mut state = client.state.write();
let mut state = client.handler_set.lock();
let _ = state.entities_by_type_and_remote_id.remove(id);
}
}
Subscription::Message { client, id } => {
if let Some(client) = client.upgrade() {
let mut state = client.state.write();
let mut state = client.handler_set.lock();
let _ = state.entity_types_by_message_type.remove(id);
let _ = state.message_handlers.remove(id);
}
@ -430,21 +399,21 @@ pub struct PendingEntitySubscription<T: 'static> {
impl<T: 'static> PendingEntitySubscription<T> {
pub fn set_model(mut self, model: &Model<T>, cx: &mut AsyncAppContext) -> Subscription {
self.consumed = true;
let mut state = self.client.state.write();
let mut handlers = self.client.handler_set.lock();
let id = (TypeId::of::<T>(), self.remote_id);
let Some(WeakSubscriber::Pending(messages)) =
state.entities_by_type_and_remote_id.remove(&id)
let Some(EntityMessageSubscriber::Pending(messages)) =
handlers.entities_by_type_and_remote_id.remove(&id)
else {
unreachable!()
};
state.entities_by_type_and_remote_id.insert(
handlers.entities_by_type_and_remote_id.insert(
id,
WeakSubscriber::Entity {
EntityMessageSubscriber::Entity {
handle: model.downgrade().into(),
},
);
drop(state);
drop(handlers);
for message in messages {
let client_id = self.client.id();
let type_name = message.payload_type_name();
@ -467,8 +436,8 @@ impl<T: 'static> PendingEntitySubscription<T> {
impl<T: 'static> Drop for PendingEntitySubscription<T> {
fn drop(&mut self) {
if !self.consumed {
let mut state = self.client.state.write();
if let Some(WeakSubscriber::Pending(messages)) = state
let mut state = self.client.handler_set.lock();
if let Some(EntityMessageSubscriber::Pending(messages)) = state
.entities_by_type_and_remote_id
.remove(&(TypeId::of::<T>(), self.remote_id))
{
@ -549,6 +518,7 @@ impl Client {
http,
credentials_provider,
state: Default::default(),
handler_set: Default::default(),
#[cfg(any(test, feature = "test-support"))]
authenticate: Default::default(),
@ -592,10 +562,7 @@ impl Client {
pub fn teardown(&self) {
let mut state = self.state.write();
state._reconnect_task.take();
state.message_handlers.clear();
state.models_by_message_type.clear();
state.entities_by_type_and_remote_id.clear();
state.entity_id_extractors.clear();
self.handler_set.lock().clear();
self.peer.teardown();
}
@ -708,14 +675,14 @@ impl Client {
{
let id = (TypeId::of::<T>(), remote_id);
let mut state = self.state.write();
let mut state = self.handler_set.lock();
if state.entities_by_type_and_remote_id.contains_key(&id) {
return Err(anyhow!("already subscribed to entity"));
}
state
.entities_by_type_and_remote_id
.insert(id, WeakSubscriber::Pending(Default::default()));
.insert(id, EntityMessageSubscriber::Pending(Default::default()));
Ok(PendingEntitySubscription {
client: self.clone(),
@ -752,13 +719,13 @@ impl Client {
E: 'static,
H: 'static
+ Sync
+ Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F
+ Fn(Model<E>, TypedEnvelope<M>, AnyProtoClient, AsyncAppContext) -> F
+ Send
+ Sync,
F: 'static + Future<Output = Result<()>>,
{
let message_type_id = TypeId::of::<M>();
let mut state = self.state.write();
let mut state = self.handler_set.lock();
state
.models_by_message_type
.insert(message_type_id, entity.into());
@ -803,85 +770,18 @@ impl Client {
})
}
pub fn add_model_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
where
M: EntityMessage,
E: 'static,
H: 'static + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F + Send + Sync,
F: 'static + Future<Output = Result<()>>,
{
self.add_entity_message_handler::<M, E, _, _>(move |subscriber, message, _, cx| {
handler(subscriber.downcast::<E>().unwrap(), message, cx)
})
}
fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
where
M: EntityMessage,
E: 'static,
H: 'static + Fn(AnyModel, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F + Send + Sync,
F: 'static + Future<Output = Result<()>>,
{
let model_type_id = TypeId::of::<E>();
let message_type_id = TypeId::of::<M>();
let mut state = self.state.write();
state
.entity_types_by_message_type
.insert(message_type_id, model_type_id);
state
.entity_id_extractors
.entry(message_type_id)
.or_insert_with(|| {
|envelope| {
envelope
.as_any()
.downcast_ref::<TypedEnvelope<M>>()
.unwrap()
.payload
.remote_entity_id()
}
});
let prev_handler = state.message_handlers.insert(
message_type_id,
Arc::new(move |handle, envelope, client, cx| {
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
handler(handle, *envelope, client.clone(), cx).boxed_local()
}),
);
if prev_handler.is_some() {
panic!("registered handler for the same message twice");
}
}
pub fn add_model_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
where
M: EntityMessage + RequestMessage,
E: 'static,
H: 'static + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F + Send + Sync,
F: 'static + Future<Output = Result<M::Response>>,
{
self.add_entity_message_handler::<M, E, _, _>(move |entity, envelope, client, cx| {
Self::respond_to_request::<M, _>(
envelope.receipt(),
handler(entity.downcast::<E>().unwrap(), envelope, cx),
client,
)
})
}
async fn respond_to_request<T: RequestMessage, F: Future<Output = Result<T::Response>>>(
receipt: Receipt<T>,
response: F,
client: Arc<Self>,
client: AnyProtoClient,
) -> Result<()> {
match response.await {
Ok(response) => {
client.respond(receipt, response)?;
client.send_response(receipt.message_id, response)?;
Ok(())
}
Err(error) => {
client.respond_with_error(receipt, error.to_proto())?;
client.send_response(receipt.message_id, error.to_proto())?;
Err(error)
}
}
@ -1541,16 +1441,6 @@ impl Client {
self.peer.send(self.connection_id()?, message)
}
pub fn send_dynamic(
&self,
envelope: proto::Envelope,
message_type: &'static str,
) -> Result<()> {
log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
let connection_id = self.connection_id()?;
self.peer.send_dynamic(connection_id, envelope)
}
pub fn request<T: RequestMessage>(
&self,
request: T,
@ -1632,115 +1522,56 @@ impl Client {
}
}
fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
self.peer.respond(receipt, response)
}
fn respond_with_error<T: RequestMessage>(
&self,
receipt: Receipt<T>,
error: proto::Error,
) -> Result<()> {
log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
self.peer.respond_with_error(receipt, error)
}
fn handle_message(
self: &Arc<Client>,
message: Box<dyn AnyTypedEnvelope>,
cx: &AsyncAppContext,
) {
let mut state = self.state.write();
let sender_id = message.sender_id();
let request_id = message.message_id();
let type_name = message.payload_type_name();
let payload_type_id = message.payload_type_id();
let sender_id = message.original_sender_id();
let original_sender_id = message.original_sender_id();
let mut subscriber = None;
if let Some(handle) = state
.models_by_message_type
.get(&payload_type_id)
.and_then(|handle| handle.upgrade())
{
subscriber = Some(handle);
} else if let Some((extract_entity_id, entity_type_id)) =
state.entity_id_extractors.get(&payload_type_id).zip(
state
.entity_types_by_message_type
.get(&payload_type_id)
.copied(),
)
{
let entity_id = (extract_entity_id)(message.as_ref());
match state
.entities_by_type_and_remote_id
.get_mut(&(entity_type_id, entity_id))
{
Some(WeakSubscriber::Pending(pending)) => {
pending.push(message);
return;
}
Some(weak_subscriber) => match weak_subscriber {
WeakSubscriber::Entity { handle } => {
subscriber = handle.upgrade();
}
WeakSubscriber::Pending(_) => {}
},
_ => {}
}
}
let subscriber = if let Some(subscriber) = subscriber {
subscriber
} else {
log::info!("unhandled message {}", type_name);
self.peer.respond_with_unhandled_message(message).log_err();
return;
};
let handler = state.message_handlers.get(&payload_type_id).cloned();
// Dropping the state prevents deadlocks if the handler interacts with rpc::Client.
// It also ensures we don't hold the lock while yielding back to the executor, as
// that might cause the executor thread driving this future to block indefinitely.
drop(state);
if let Some(handler) = handler {
let future = handler(subscriber, message, self, cx.clone());
if let Some(future) = ProtoMessageHandlerSet::handle_message(
&self.handler_set,
message,
self.clone().into(),
cx.clone(),
) {
let client_id = self.id();
log::debug!(
"rpc message received. client_id:{}, sender_id:{:?}, type:{}",
client_id,
sender_id,
original_sender_id,
type_name
);
cx.spawn(move |_| async move {
match future.await {
Ok(()) => {
log::debug!(
"rpc message handled. client_id:{}, sender_id:{:?}, type:{}",
client_id,
sender_id,
type_name
);
}
Err(error) => {
log::error!(
"error handling message. client_id:{}, sender_id:{:?}, type:{}, error:{:?}",
client_id,
sender_id,
type_name,
error
);
}
match future.await {
Ok(()) => {
log::debug!(
"rpc message handled. client_id:{}, sender_id:{:?}, type:{}",
client_id,
original_sender_id,
type_name
);
}
})
.detach();
Err(error) => {
log::error!(
"error handling message. client_id:{}, sender_id:{:?}, type:{}, error:{:?}",
client_id,
original_sender_id,
type_name,
error
);
}
}
})
.detach();
} else {
log::info!("unhandled message {}", type_name);
self.peer.respond_with_unhandled_message(message).log_err();
self.peer
.respond_with_unhandled_message(sender_id.into(), request_id, type_name)
.log_err();
}
}
@ -1759,7 +1590,23 @@ impl ProtoClient for Client {
}
fn send(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
self.send_dynamic(envelope, message_type)
log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
let connection_id = self.connection_id()?;
self.peer.send_dynamic(connection_id, envelope)
}
fn send_response(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
log::debug!(
"rpc respond. client_id:{}, name:{}",
self.id(),
message_type
);
let connection_id = self.connection_id()?;
self.peer.send_dynamic(connection_id, envelope)
}
fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet> {
&self.handler_set
}
}
@ -2103,7 +1950,7 @@ mod tests {
let (done_tx1, mut done_rx1) = smol::channel::unbounded();
let (done_tx2, mut done_rx2) = smol::channel::unbounded();
client.add_model_message_handler(
AnyProtoClient::from(client.clone()).add_model_message_handler(
move |model: Model<TestModel>, _: TypedEnvelope<proto::JoinProject>, mut cx| {
match model.update(&mut cx, |model, _| model.id).unwrap() {
1 => done_tx1.try_send(()).unwrap(),

2
crates/collab/src/tests/test_server.rs
View File

@ -301,7 +301,7 @@ impl TestServer {
dev_server_projects::init(client.clone(), cx);
settings::KeymapFile::load_asset(os_keymap, cx).unwrap();
language_model::LanguageModelRegistry::test(cx);
assistant::context_store::init(&client);
assistant::context_store::init(&client.clone().into());
});
client

2
crates/project/src/buffer_store.rs
View File

@ -71,7 +71,7 @@ pub struct ProjectTransaction(pub HashMap<Model<Buffer>, language::Transaction>)
impl EventEmitter<BufferStoreEvent> for BufferStore {}
impl BufferStore {
pub fn init(client: &Arc<Client>) {
pub fn init(client: &AnyProtoClient) {
client.add_model_message_handler(Self::handle_buffer_reloaded);
client.add_model_message_handler(Self::handle_buffer_saved);
client.add_model_message_handler(Self::handle_update_buffer_file);

1045
crates/project/src/lsp_store.rs
View File

File diff suppressed because it is too large Load Diff

28
crates/project/src/project.rs
View File

@ -65,7 +65,10 @@ use paths::{
use prettier_support::{DefaultPrettier, PrettierInstance};
use project_settings::{LspSettings, ProjectSettings};
use remote::SshSession;
use rpc::{proto::AnyProtoClient, ErrorCode};
use rpc::{
proto::{AnyProtoClient, SSH_PROJECT_ID},
ErrorCode,
};
use search::{SearchQuery, SearchResult};
use search_history::SearchHistory;
use settings::{watch_config_file, Settings, SettingsLocation, SettingsStore};
@ -574,6 +577,7 @@ impl Project {
connection_manager::init(client.clone(), cx);
Self::init_settings(cx);
let client: AnyProtoClient = client.clone().into();
client.add_model_message_handler(Self::handle_add_collaborator);
client.add_model_message_handler(Self::handle_update_project_collaborator);
client.add_model_message_handler(Self::handle_remove_collaborator);
@ -594,9 +598,9 @@ impl Project {
client.add_model_request_handler(Self::handle_task_templates);
client.add_model_message_handler(Self::handle_create_buffer_for_peer);
WorktreeStore::init(client);
BufferStore::init(client);
LspStore::init(client);
WorktreeStore::init(&client);
BufferStore::init(&client);
LspStore::init(&client);
}
pub fn local(
@ -697,15 +701,19 @@ impl Project {
) -> Model<Self> {
let this = Self::local(client, node, user_store, languages, fs, None, cx);
this.update(cx, |this, cx| {
let buffer_store = this.buffer_store.downgrade();
let client: AnyProtoClient = ssh.clone().into();
this.worktree_store.update(cx, |store, _cx| {
store.set_upstream_client(ssh.clone().into());
store.set_upstream_client(client.clone());
});
ssh.add_message_handler(cx.weak_model(), Self::handle_update_worktree);
ssh.add_message_handler(cx.weak_model(), Self::handle_create_buffer_for_peer);
ssh.add_message_handler(buffer_store.clone(), BufferStore::handle_update_buffer_file);
ssh.add_message_handler(buffer_store.clone(), BufferStore::handle_update_diff_base);
ssh.subscribe_to_entity(SSH_PROJECT_ID, &cx.handle());
ssh.subscribe_to_entity(SSH_PROJECT_ID, &this.buffer_store);
ssh.subscribe_to_entity(SSH_PROJECT_ID, &this.worktree_store);
client.add_model_message_handler(Self::handle_update_worktree);
client.add_model_message_handler(Self::handle_create_buffer_for_peer);
client.add_model_message_handler(BufferStore::handle_update_buffer_file);
client.add_model_message_handler(BufferStore::handle_update_diff_base);
this.ssh_session = Some(ssh);
});

21
crates/project/src/worktree_store.rs
View File

@ -5,7 +5,7 @@ use std::{
};
use anyhow::{anyhow, Context as _, Result};
use client::{Client, DevServerProjectId};
use client::DevServerProjectId;
use collections::{HashMap, HashSet};
use fs::Fs;
use futures::{
@ -17,7 +17,7 @@ use gpui::{
};
use postage::oneshot;
use rpc::{
proto::{self, AnyProtoClient},
proto::{self, AnyProtoClient, SSH_PROJECT_ID},
TypedEnvelope,
};
use smol::{
@ -58,12 +58,12 @@ pub enum WorktreeStoreEvent {
impl EventEmitter<WorktreeStoreEvent> for WorktreeStore {}
impl WorktreeStore {
pub fn init(client: &Arc<Client>) {
client.add_model_request_handler(WorktreeStore::handle_create_project_entry);
client.add_model_request_handler(WorktreeStore::handle_rename_project_entry);
client.add_model_request_handler(WorktreeStore::handle_copy_project_entry);
client.add_model_request_handler(WorktreeStore::handle_delete_project_entry);
client.add_model_request_handler(WorktreeStore::handle_expand_project_entry);
pub fn init(client: &AnyProtoClient) {
client.add_model_request_handler(Self::handle_create_project_entry);
client.add_model_request_handler(Self::handle_rename_project_entry);
client.add_model_request_handler(Self::handle_copy_project_entry);
client.add_model_request_handler(Self::handle_delete_project_entry);
client.add_model_request_handler(Self::handle_expand_project_entry);
}
pub fn new(retain_worktrees: bool, fs: Arc<dyn Fs>) -> Self {
@ -188,7 +188,10 @@ impl WorktreeStore {
let path = abs_path.to_string_lossy().to_string();
cx.spawn(|this, mut cx| async move {
let response = client
.request(proto::AddWorktree { path: path.clone() })
.request(proto::AddWorktree {
project_id: SSH_PROJECT_ID,
path: path.clone(),
})
.await?;
let worktree = cx.update(|cx| {
Worktree::remote(

2
crates/proto/Cargo.toml
View File

@ -20,8 +20,10 @@ doctest = false
anyhow.workspace = true
collections.workspace = true
futures.workspace = true
parking_lot.workspace = true
prost.workspace = true
serde.workspace = true
gpui.workspace = true
[build-dependencies]
prost-build.workspace = true

1
crates/proto/proto/zed.proto
View File

@ -2484,6 +2484,7 @@ message GetLlmTokenResponse {
// Remote FS
message AddWorktree {
uint64 project_id = 2;
string path = 1;
}

55
crates/proto/src/proto.rs
View File

@ -2,14 +2,14 @@
pub mod error;
mod macros;
mod proto_client;
mod typed_envelope;
pub use error::*;
pub use proto_client::*;
pub use typed_envelope::*;
use anyhow::anyhow;
use collections::HashMap;
use futures::{future::BoxFuture, Future};
pub use prost::{DecodeError, Message};
use serde::Serialize;
use std::{
@ -17,12 +17,14 @@ use std::{
cmp,
fmt::{self, Debug},
iter, mem,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
include!(concat!(env!("OUT_DIR"), "/zed.messages.rs"));
pub const SSH_PEER_ID: PeerId = PeerId { owner_id: 0, id: 0 };
pub const SSH_PROJECT_ID: u64 = 0;
pub trait EnvelopedMessage: Clone + Debug + Serialize + Sized + Send + Sync + 'static {
const NAME: &'static str;
const PRIORITY: MessagePriority;
@ -60,51 +62,6 @@ pub enum MessagePriority {
Background,
}
pub trait ProtoClient: Send + Sync {
fn request(
&self,
envelope: Envelope,
request_type: &'static str,
) -> BoxFuture<'static, anyhow::Result<Envelope>>;
fn send(&self, envelope: Envelope, message_type: &'static str) -> anyhow::Result<()>;
}
#[derive(Clone)]
pub struct AnyProtoClient(Arc<dyn ProtoClient>);
impl<T> From<Arc<T>> for AnyProtoClient
where
T: ProtoClient + 'static,
{
fn from(client: Arc<T>) -> Self {
Self(client)
}
}
impl AnyProtoClient {
pub fn new<T: ProtoClient + 'static>(client: Arc<T>) -> Self {
Self(client)
}
pub fn request<T: RequestMessage>(
&self,
request: T,
) -> impl Future<Output = anyhow::Result<T::Response>> {
let envelope = request.into_envelope(0, None, None);
let response = self.0.request(envelope, T::NAME);
async move {
T::Response::from_envelope(response.await?)
.ok_or_else(|| anyhow!("received response of the wrong type"))
}
}
pub fn send<T: EnvelopedMessage>(&self, request: T) -> anyhow::Result<()> {
let envelope = request.into_envelope(0, None, None);
self.0.send(envelope, T::NAME)
}
}
impl<T: EnvelopedMessage> AnyTypedEnvelope for TypedEnvelope<T> {
fn payload_type_id(&self) -> TypeId {
TypeId::of::<T>()
@ -537,11 +494,13 @@ request_messages!(
entity_messages!(
{project_id, ShareProject},
AddProjectCollaborator,
AddWorktree,
ApplyCodeAction,
ApplyCompletionAdditionalEdits,
BlameBuffer,
BufferReloaded,
BufferSaved,
CloseBuffer,
CopyProjectEntry,
CreateBufferForPeer,
CreateProjectEntry,

277
crates/proto/src/proto_client.rs Normal file
View File

@ -0,0 +1,277 @@
use crate::{
error::ErrorExt as _, AnyTypedEnvelope, EntityMessage, Envelope, EnvelopedMessage,
RequestMessage, TypedEnvelope,
};
use anyhow::anyhow;
use collections::HashMap;
use futures::{
future::{BoxFuture, LocalBoxFuture},
Future, FutureExt as _,
};
use gpui::{AnyModel, AnyWeakModel, AsyncAppContext, Model};
pub use prost::Message;
use std::{any::TypeId, sync::Arc};
#[derive(Clone)]
pub struct AnyProtoClient(Arc<dyn ProtoClient>);
pub trait ProtoClient: Send + Sync {
fn request(
&self,
envelope: Envelope,
request_type: &'static str,
) -> BoxFuture<'static, anyhow::Result<Envelope>>;
fn send(&self, envelope: Envelope, message_type: &'static str) -> anyhow::Result<()>;
fn send_response(&self, envelope: Envelope, message_type: &'static str) -> anyhow::Result<()>;
fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet>;
}
#[derive(Default)]
pub struct ProtoMessageHandlerSet {
pub entity_types_by_message_type: HashMap<TypeId, TypeId>,
pub entities_by_type_and_remote_id: HashMap<(TypeId, u64), EntityMessageSubscriber>,
pub entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
pub models_by_message_type: HashMap<TypeId, AnyWeakModel>,
pub message_handlers: HashMap<TypeId, ProtoMessageHandler>,
}
pub type ProtoMessageHandler = Arc<
dyn Send
+ Sync
+ Fn(
AnyModel,
Box<dyn AnyTypedEnvelope>,
AnyProtoClient,
AsyncAppContext,
) -> LocalBoxFuture<'static, anyhow::Result<()>>,
>;
impl ProtoMessageHandlerSet {
pub fn clear(&mut self) {
self.message_handlers.clear();
self.models_by_message_type.clear();
self.entities_by_type_and_remote_id.clear();
self.entity_id_extractors.clear();
}
fn add_message_handler(
&mut self,
message_type_id: TypeId,
model: gpui::AnyWeakModel,
handler: ProtoMessageHandler,
) {
self.models_by_message_type.insert(message_type_id, model);
let prev_handler = self.message_handlers.insert(message_type_id, handler);
if prev_handler.is_some() {
panic!("registered handler for the same message twice");
}
}
fn add_entity_message_handler(
&mut self,
message_type_id: TypeId,
model_type_id: TypeId,
entity_id_extractor: fn(&dyn AnyTypedEnvelope) -> u64,
handler: ProtoMessageHandler,
) {
self.entity_id_extractors
.entry(message_type_id)
.or_insert(entity_id_extractor);
self.entity_types_by_message_type
.insert(message_type_id, model_type_id);
let prev_handler = self.message_handlers.insert(message_type_id, handler);
if prev_handler.is_some() {
panic!("registered handler for the same message twice");
}
}
pub fn handle_message(
this: &parking_lot::Mutex<Self>,
message: Box<dyn AnyTypedEnvelope>,
client: AnyProtoClient,
cx: AsyncAppContext,
) -> Option<LocalBoxFuture<'static, anyhow::Result<()>>> {
let payload_type_id = message.payload_type_id();
let mut this = this.lock();
let handler = this.message_handlers.get(&payload_type_id)?.clone();
let entity = if let Some(entity) = this.models_by_message_type.get(&payload_type_id) {
entity.upgrade()?
} else {
let extract_entity_id = *this.entity_id_extractors.get(&payload_type_id)?;
let entity_type_id = *this.entity_types_by_message_type.get(&payload_type_id)?;
let entity_id = (extract_entity_id)(message.as_ref());
match this
.entities_by_type_and_remote_id
.get_mut(&(entity_type_id, entity_id))?
{
EntityMessageSubscriber::Pending(pending) => {
pending.push(message);
return None;
}
EntityMessageSubscriber::Entity { handle } => handle.upgrade()?,
}
};
drop(this);
Some(handler(entity, message, client, cx))
}
}
pub enum EntityMessageSubscriber {
Entity { handle: AnyWeakModel },
Pending(Vec<Box<dyn AnyTypedEnvelope>>),
}
impl<T> From<Arc<T>> for AnyProtoClient
where
T: ProtoClient + 'static,
{
fn from(client: Arc<T>) -> Self {
Self(client)
}
}
impl AnyProtoClient {
pub fn new<T: ProtoClient + 'static>(client: Arc<T>) -> Self {
Self(client)
}
pub fn request<T: RequestMessage>(
&self,
request: T,
) -> impl Future<Output = anyhow::Result<T::Response>> {
let envelope = request.into_envelope(0, None, None);
let response = self.0.request(envelope, T::NAME);
async move {
T::Response::from_envelope(response.await?)
.ok_or_else(|| anyhow!("received response of the wrong type"))
}
}
pub fn send<T: EnvelopedMessage>(&self, request: T) -> anyhow::Result<()> {
let envelope = request.into_envelope(0, None, None);
self.0.send(envelope, T::NAME)
}
pub fn send_response<T: EnvelopedMessage>(
&self,
request_id: u32,
request: T,
) -> anyhow::Result<()> {
let envelope = request.into_envelope(0, Some(request_id), None);
self.0.send(envelope, T::NAME)
}
pub fn add_request_handler<M, E, H, F>(&self, model: gpui::WeakModel<E>, handler: H)
where
M: RequestMessage,
E: 'static,
H: 'static + Sync + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F + Send + Sync,
F: 'static + Future<Output = anyhow::Result<M::Response>>,
{
self.0.message_handler_set().lock().add_message_handler(
TypeId::of::<M>(),
model.into(),
Arc::new(move |model, envelope, client, cx| {
let model = model.downcast::<E>().unwrap();
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
let request_id = envelope.message_id();
handler(model, *envelope, cx)
.then(move |result| async move {
match result {
Ok(response) => {
client.send_response(request_id, response)?;
Ok(())
}
Err(error) => {
client.send_response(request_id, error.to_proto())?;
Err(error)
}
}
})
.boxed_local()
}),
)
}
pub fn add_model_request_handler<M, E, H, F>(&self, handler: H)
where
M: EnvelopedMessage + RequestMessage + EntityMessage,
E: 'static,
H: 'static + Sync + Send + Fn(gpui::Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
F: 'static + Future<Output = anyhow::Result<M::Response>>,
{
let message_type_id = TypeId::of::<M>();
let model_type_id = TypeId::of::<E>();
let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
envelope
.as_any()
.downcast_ref::<TypedEnvelope<M>>()
.unwrap()
.payload
.remote_entity_id()
};
self.0
.message_handler_set()
.lock()
.add_entity_message_handler(
message_type_id,
model_type_id,
entity_id_extractor,
Arc::new(move |model, envelope, client, cx| {
let model = model.downcast::<E>().unwrap();
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
let request_id = envelope.message_id();
handler(model, *envelope, cx)
.then(move |result| async move {
match result {
Ok(response) => {
client.send_response(request_id, response)?;
Ok(())
}
Err(error) => {
client.send_response(request_id, error.to_proto())?;
Err(error)
}
}
})
.boxed_local()
}),
);
}
pub fn add_model_message_handler<M, E, H, F>(&self, handler: H)
where
M: EnvelopedMessage + EntityMessage,
E: 'static,
H: 'static + Sync + Send + Fn(gpui::Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
F: 'static + Future<Output = anyhow::Result<()>>,
{
let message_type_id = TypeId::of::<M>();
let model_type_id = TypeId::of::<E>();
let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
envelope
.as_any()
.downcast_ref::<TypedEnvelope<M>>()
.unwrap()
.payload
.remote_entity_id()
};
self.0
.message_handler_set()
.lock()
.add_entity_message_handler(
message_type_id,
model_type_id,
entity_id_extractor,
Arc::new(move |model, envelope, _, cx| {
let model = model.downcast::<E>().unwrap();
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
handler(model, *envelope, cx).boxed_local()
}),
);
}
}

133
crates/remote/src/ssh_session.rs
View File

@ -8,17 +8,14 @@ use anyhow::{anyhow, Context as _, Result};
use collections::HashMap;
use futures::{
channel::{mpsc, oneshot},
future::{BoxFuture, LocalBoxFuture},
future::BoxFuture,
select_biased, AsyncReadExt as _, AsyncWriteExt as _, Future, FutureExt as _, StreamExt as _,
};
use gpui::{AppContext, AsyncAppContext, Model, SemanticVersion, WeakModel};
use gpui::{AppContext, AsyncAppContext, Model, SemanticVersion};
use parking_lot::Mutex;
use rpc::{
proto::{
self, build_typed_envelope, AnyTypedEnvelope, Envelope, EnvelopedMessage, PeerId,
ProtoClient, RequestMessage,
},
TypedEnvelope,
use rpc::proto::{
self, build_typed_envelope, EntityMessageSubscriber, Envelope, EnvelopedMessage, PeerId,
ProtoClient, ProtoMessageHandlerSet, RequestMessage,
};
use smol::{
fs,
@ -48,20 +45,7 @@ pub struct SshSession {
outgoing_tx: mpsc::UnboundedSender<Envelope>,
spawn_process_tx: mpsc::UnboundedSender<SpawnRequest>,
client_socket: Option<SshSocket>,
message_handlers: Mutex<
HashMap<
TypeId,
Arc<
dyn Send
+ Sync
+ Fn(
Box<dyn AnyTypedEnvelope>,
Arc<SshSession>,
AsyncAppContext,
) -> Option<LocalBoxFuture<'static, Result<()>>>,
>,
>,
>,
state: Mutex<ProtoMessageHandlerSet>,
}
struct SshClientState {
@ -330,7 +314,7 @@ impl SshSession {
outgoing_tx,
spawn_process_tx,
client_socket,
message_handlers: Default::default(),
state: Default::default(),
});
cx.spawn(|cx| {
@ -351,18 +335,26 @@ impl SshSession {
} else if let Some(envelope) =
build_typed_envelope(peer_id, Instant::now(), incoming)
{
log::debug!(
"ssh message received. name:{}",
envelope.payload_type_name()
);
let type_id = envelope.payload_type_id();
let handler = this.message_handlers.lock().get(&type_id).cloned();
if let Some(handler) = handler {
if let Some(future) = handler(envelope, this.clone(), cx.clone()) {
future.await.ok();
} else {
this.message_handlers.lock().remove(&type_id);
let type_name = envelope.payload_type_name();
if let Some(future) = ProtoMessageHandlerSet::handle_message(
&this.state,
envelope,
this.clone().into(),
cx.clone(),
) {
log::debug!("ssh message received. name:{type_name}");
match future.await {
Ok(_) => {
log::debug!("ssh message handled. name:{type_name}");
}
Err(error) => {
log::error!(
"error handling message. type:{type_name}, error:{error:?}",
);
}
}
} else {
log::error!("unhandled ssh message name:{type_name}");
}
}
}
@ -389,6 +381,7 @@ impl SshSession {
}
pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
log::debug!("ssh send name:{}", T::NAME);
self.send_dynamic(payload.into_envelope(0, None, None))
}
@ -412,6 +405,22 @@ impl SshSession {
Ok(())
}
pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
let id = (TypeId::of::<E>(), remote_id);
let mut state = self.state.lock();
if state.entities_by_type_and_remote_id.contains_key(&id) {
panic!("already subscribed to entity");
}
state.entities_by_type_and_remote_id.insert(
id,
EntityMessageSubscriber::Entity {
handle: entity.downgrade().into(),
},
);
}
pub async fn spawn_process(&self, command: String) -> process::Child {
let (process_tx, process_rx) = oneshot::channel();
self.spawn_process_tx
@ -426,54 +435,6 @@ impl SshSession {
pub fn ssh_args(&self) -> Vec<String> {
self.client_socket.as_ref().unwrap().ssh_args()
}
pub fn add_message_handler<M, E, H, F>(&self, entity: WeakModel<E>, handler: H)
where
M: EnvelopedMessage,
E: 'static,
H: 'static + Sync + Send + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
F: 'static + Future<Output = Result<()>>,
{
let message_type_id = TypeId::of::<M>();
self.message_handlers.lock().insert(
message_type_id,
Arc::new(move |envelope, _, cx| {
let entity = entity.upgrade()?;
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
Some(handler(entity, *envelope, cx).boxed_local())
}),
);
}
pub fn add_request_handler<M, E, H, F>(&self, entity: WeakModel<E>, handler: H)
where
M: EnvelopedMessage + RequestMessage,
E: 'static,
H: 'static + Sync + Send + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
F: 'static + Future<Output = Result<M::Response>>,
{
let message_type_id = TypeId::of::<M>();
self.message_handlers.lock().insert(
message_type_id,
Arc::new(move |envelope, this, cx| {
let entity = entity.upgrade()?;
let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
let request_id = envelope.message_id();
Some(
handler(entity, *envelope, cx)
.then(move |result| async move {
this.outgoing_tx.unbounded_send(result?.into_envelope(
this.next_message_id.fetch_add(1, SeqCst),
Some(request_id),
None,
))?;
Ok(())
})
.boxed_local(),
)
}),
);
}
}
impl ProtoClient for SshSession {
@ -488,6 +449,14 @@ impl ProtoClient for SshSession {
fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
self.send_dynamic(envelope)
}
fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
self.send_dynamic(envelope)
}
fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
&self.state
}
}
impl SshClientState {

61
crates/remote_server/src/headless_project.rs
View File

@ -7,7 +7,7 @@ use project::{
};
use remote::SshSession;
use rpc::{
proto::{self, AnyProtoClient, PeerId},
proto::{self, AnyProtoClient, SSH_PEER_ID, SSH_PROJECT_ID},
TypedEnvelope,
};
use settings::{Settings as _, SettingsStore};
@ -18,9 +18,6 @@ use std::{
};
use worktree::Worktree;
const PEER_ID: PeerId = PeerId { owner_id: 0, id: 0 };
const PROJECT_ID: u64 = 0;
pub struct HeadlessProject {
pub fs: Arc<dyn Fs>,
pub session: AnyProtoClient,
@ -36,48 +33,34 @@ impl HeadlessProject {
}
pub fn new(session: Arc<SshSession>, fs: Arc<dyn Fs>, cx: &mut ModelContext<Self>) -> Self {
let this = cx.weak_model();
let worktree_store = cx.new_model(|_| WorktreeStore::new(true, fs.clone()));
let buffer_store = cx.new_model(|cx| {
let mut buffer_store = BufferStore::new(worktree_store.clone(), Some(PROJECT_ID), cx);
buffer_store.shared(PROJECT_ID, session.clone().into(), cx);
let mut buffer_store =
BufferStore::new(worktree_store.clone(), Some(SSH_PROJECT_ID), cx);
buffer_store.shared(SSH_PROJECT_ID, session.clone().into(), cx);
buffer_store
});
session.add_request_handler(this.clone(), Self::handle_list_remote_directory);
session.add_request_handler(this.clone(), Self::handle_add_worktree);
session.add_request_handler(this.clone(), Self::handle_open_buffer_by_path);
session.add_request_handler(this.clone(), Self::handle_find_search_candidates);
let client: AnyProtoClient = session.clone().into();
session.add_request_handler(buffer_store.downgrade(), BufferStore::handle_blame_buffer);
session.add_request_handler(buffer_store.downgrade(), BufferStore::handle_update_buffer);
session.add_request_handler(buffer_store.downgrade(), BufferStore::handle_save_buffer);
session.add_message_handler(buffer_store.downgrade(), BufferStore::handle_close_buffer);
session.subscribe_to_entity(SSH_PROJECT_ID, &worktree_store);
session.subscribe_to_entity(SSH_PROJECT_ID, &buffer_store);
session.subscribe_to_entity(SSH_PROJECT_ID, &cx.handle());
session.add_request_handler(
worktree_store.downgrade(),
WorktreeStore::handle_create_project_entry,
);
session.add_request_handler(
worktree_store.downgrade(),
WorktreeStore::handle_rename_project_entry,
);
session.add_request_handler(
worktree_store.downgrade(),
WorktreeStore::handle_copy_project_entry,
);
session.add_request_handler(
worktree_store.downgrade(),
WorktreeStore::handle_delete_project_entry,
);
session.add_request_handler(
worktree_store.downgrade(),
WorktreeStore::handle_expand_project_entry,
);
client.add_request_handler(cx.weak_model(), Self::handle_list_remote_directory);
client.add_model_request_handler(Self::handle_add_worktree);
client.add_model_request_handler(Self::handle_open_buffer_by_path);
client.add_model_request_handler(Self::handle_find_search_candidates);
client.add_model_request_handler(BufferStore::handle_update_buffer);
client.add_model_message_handler(BufferStore::handle_close_buffer);
BufferStore::init(&client);
WorktreeStore::init(&client);
HeadlessProject {
session: session.into(),
session: client,
fs,
worktree_store,
buffer_store,
@ -144,7 +127,7 @@ impl HeadlessProject {
let buffer_id = buffer.read_with(&cx, |b, _| b.remote_id())?;
buffer_store.update(&mut cx, |buffer_store, cx| {
buffer_store
.create_buffer_for_peer(&buffer, PEER_ID, cx)
.create_buffer_for_peer(&buffer, SSH_PEER_ID, cx)
.detach_and_log_err(cx);
})?;
@ -181,7 +164,7 @@ impl HeadlessProject {
response.buffer_ids.push(buffer_id.to_proto());
buffer_store
.update(&mut cx, |buffer_store, cx| {
buffer_store.create_buffer_for_peer(&buffer, PEER_ID, cx)
buffer_store.create_buffer_for_peer(&buffer, SSH_PEER_ID, cx)
})?
.await?;
}

2
crates/remote_server/src/remote_editing_tests.rs
View File

@ -17,7 +17,7 @@ use smol::stream::StreamExt;
use std::{path::Path, sync::Arc};
#[gpui::test]
async fn test_remote_editing(cx: &mut TestAppContext, server_cx: &mut TestAppContext) {
async fn test_basic_remote_editing(cx: &mut TestAppContext, server_cx: &mut TestAppContext) {
let (project, _headless, fs) = init_test(cx, server_cx).await;
let (worktree, _) = project
.update(cx, |project, cx| {

13
crates/rpc/src/peer.rs
View File

@ -639,14 +639,13 @@ impl Peer {
pub fn respond_with_unhandled_message(
&self,
envelope: Box<dyn AnyTypedEnvelope>,
sender_id: ConnectionId,
request_message_id: u32,
message_type_name: &'static str,
) -> Result<()> {
let connection = self.connection_state(envelope.sender_id().into())?;
let connection = self.connection_state(sender_id)?;
let response = ErrorCode::Internal
.message(format!(
"message {} was not handled",
envelope.payload_type_name()
))
.message(format!("message {} was not handled", message_type_name))
.to_proto();
let message_id = connection
.next_message_id
@ -655,7 +654,7 @@ impl Peer {
.outgoing_tx
.unbounded_send(proto::Message::Envelope(response.into_envelope(
message_id,
Some(envelope.message_id()),
Some(request_message_id),
None,
)))?;
Ok(())