Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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
Download patch file Download diff file Expand all files Collapse all files

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 {