Fix deadlock when comparing `FileHandle`s when handles are the same

Antonio Scandurra created

Change summary

zed-rpc/src/peer.rs  | 10 ++++++++--
zed-rpc/src/proto.rs |  6 +++++-
zed/src/rpc.rs       |  4 ++--
zed/src/workspace.rs | 28 +++++++++++++++++++++-------
zed/src/worktree.rs  | 42 +++++++++++++++++++++++++++---------------
5 files changed, 63 insertions(+), 27 deletions(-)

Detailed changes

zed-rpc/src/peer.rs 🔗

@@ -430,14 +430,20 @@ mod tests {
             let response2 = proto::AuthResponse {
                 credentials_valid: false,
             };
-            let request3 = proto::OpenBuffer { id: 2 };
+            let request3 = proto::OpenBuffer {
+                worktree_id: 1,
+                id: 2,
+            };
             let response3 = proto::OpenBufferResponse {
                 buffer: Some(proto::Buffer {
                     content: "path/two content".to_string(),
                     history: vec![],
                 }),
             };
-            let request4 = proto::OpenBuffer { id: 1 };
+            let request4 = proto::OpenBuffer {
+                worktree_id: 2,
+                id: 1,
+            };
             let response4 = proto::OpenBufferResponse {
                 buffer: Some(proto::Buffer {
                     content: "path/one content".to_string(),

zed-rpc/src/proto.rs 🔗

@@ -145,7 +145,11 @@ mod tests {
             }
             .into_envelope(3, None, None);
 
-            let message2 = OpenBuffer { id: 1 }.into_envelope(5, None, None);
+            let message2 = OpenBuffer {
+                worktree_id: 0,
+                id: 1,
+            }
+            .into_envelope(5, None, None);
 
             let mut message_stream = MessageStream::new(byte_stream);
             message_stream.write_message(&message1).await.unwrap();

zed/src/rpc.rs 🔗

@@ -7,7 +7,7 @@ use gpui::{AsyncAppContext, ModelHandle, Task};
 use lazy_static::lazy_static;
 use postage::prelude::Stream;
 use smol::lock::Mutex;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::time::Duration;
 use std::{convert::TryFrom, future::Future, sync::Arc};
 use surf::Url;
@@ -31,7 +31,7 @@ pub struct Client {
 #[derive(Default)]
 pub struct ClientState {
     connection_id: Option<ConnectionId>,
-    pub shared_worktrees: HashSet<ModelHandle<Worktree>>,
+    pub shared_worktrees: HashMap<u64, ModelHandle<Worktree>>,
     pub shared_files: HashMap<FileHandle, HashMap<PeerId, usize>>,
 }
 

zed/src/workspace.rs 🔗

@@ -124,7 +124,7 @@ mod remote {
         let mut state = rpc.state.lock().await;
         let worktree = state
             .shared_worktrees
-            .get(&(message.worktree_id as usize))
+            .get(&message.worktree_id)
             .ok_or_else(|| anyhow!("worktree {} not found", message.worktree_id))?
             .clone();
 
@@ -158,7 +158,7 @@ mod remote {
         if let Some((_, ref_counts)) = state
             .shared_files
             .iter_mut()
-            .find(|(file, _)| file.id() as u64 == message.id)
+            .find(|(file, _)| file.id() == message.id)
         {
             if let Some(count) = ref_counts.get_mut(&peer_id) {
                 *count -= 1;
@@ -176,11 +176,25 @@ mod remote {
         rpc: &rpc::Client,
         cx: &mut AsyncAppContext,
     ) -> anyhow::Result<()> {
-        rpc.respond(
-            request.receipt(),
-            proto::OpenBufferResponse { buffer: None },
-        )
-        .await?;
+        let message = &request.payload;
+        let handle = {
+            let state = rpc.state.lock().await;
+            let mut files = state.shared_files.keys();
+            files.find(|file| file.id() == message.id).cloned()
+        };
+        let buffer = if let Some(handle) = handle {
+            let history = cx.read(|cx| handle.load_history(cx)).await?;
+            Some(proto::Buffer {
+                content: history.base_text.to_string(),
+                history: Vec::new(),
+            })
+        } else {
+            None
+        };
+
+        rpc.respond(request.receipt(), proto::OpenBufferResponse { buffer })
+            .await?;
+
         Ok(())
     }
 }

zed/src/worktree.rs 🔗

@@ -142,7 +142,7 @@ struct FileHandleState {
     path: Arc<Path>,
     is_deleted: bool,
     mtime: Duration,
-    worktree_id: usize,
+    remote_worktree_id: usize,
     id: u64,
     rpc: Option<(ConnectionId, rpc::Client)>,
 }
@@ -151,7 +151,7 @@ impl Drop for FileHandleState {
     fn drop(&mut self) {
         if let Some((connection_id, rpc)) = self.rpc.take() {
             let id = self.id;
-            let worktree_id = self.worktree_id as u64;
+            let worktree_id = self.remote_worktree_id as u64;
             smol::spawn(async move {
                 if let Err(error) = rpc
                     .send(connection_id, proto::CloseFile { worktree_id, id })
@@ -354,7 +354,12 @@ impl LocalWorktree {
                 )
                 .await?;
 
-            client.state.lock().await.shared_worktrees.insert(handle);
+            client
+                .state
+                .lock()
+                .await
+                .shared_worktrees
+                .insert(share_response.worktree_id, handle);
 
             log::info!("sharing worktree {:?}", share_response);
             Ok((share_response.worktree_id, share_response.access_token))
@@ -377,7 +382,7 @@ impl fmt::Debug for LocalWorktree {
 }
 
 pub struct RemoteWorktree {
-    id: usize,
+    remote_id: usize,
     snapshot: Snapshot,
     handles: Arc<Mutex<HashMap<Arc<Path>, Arc<AsyncMutex<Weak<Mutex<FileHandleState>>>>>>>,
     rpc: rpc::Client,
@@ -386,7 +391,7 @@ pub struct RemoteWorktree {
 
 impl RemoteWorktree {
     fn new(
-        id: usize,
+        remote_id: usize,
         worktree: proto::Worktree,
         rpc: rpc::Client,
         connection_id: ConnectionId,
@@ -418,7 +423,7 @@ impl RemoteWorktree {
             &(),
         );
         let snapshot = Snapshot {
-            id,
+            id: cx.model_id(),
             scan_id: 0,
             abs_path: Path::new("").into(),
             root_name: worktree.root_name,
@@ -426,7 +431,7 @@ impl RemoteWorktree {
             entries,
         };
         Self {
-            id,
+            remote_id,
             snapshot,
             handles: Default::default(),
             rpc,
@@ -663,7 +668,7 @@ impl FileHandle {
             Worktree::Remote(worktree) => {
                 let state = self.state.lock();
                 let id = state.id;
-                let worktree_id = worktree.id as u64;
+                let worktree_id = worktree.remote_id as u64;
                 let (connection_id, rpc) = state.rpc.clone().unwrap();
                 cx.background_executor().spawn(async move {
                     let response = rpc
@@ -723,7 +728,14 @@ impl FileHandle {
 
 impl PartialEq for FileHandle {
     fn eq(&self, other: &Self) -> bool {
-        self.worktree == other.worktree && self.state.lock().id == other.state.lock().id
+        if Arc::ptr_eq(&self.state, &other.state) {
+            true
+        } else {
+            let self_state = self.state.lock();
+            let other_state = other.state.lock();
+            self_state.remote_worktree_id == other_state.remote_worktree_id
+                && self_state.id == other_state.id
+        }
     }
 }
 
@@ -1488,7 +1500,7 @@ impl WorktreeHandle for ModelHandle<Worktree> {
                                         path: entry.path().clone(),
                                         is_deleted: false,
                                         mtime,
-                                        worktree_id,
+                                        remote_worktree_id: worktree_id,
                                         id,
                                         rpc: None,
                                     }
@@ -1497,7 +1509,7 @@ impl WorktreeHandle for ModelHandle<Worktree> {
                                         path: path.clone(),
                                         is_deleted: !tree.path_is_pending(&path),
                                         mtime,
-                                        worktree_id,
+                                        remote_worktree_id: worktree_id,
                                         id,
                                         rpc: None,
                                     }
@@ -1515,7 +1527,7 @@ impl WorktreeHandle for ModelHandle<Worktree> {
                 })
             }
             Worktree::Remote(tree) => {
-                let worktree_id = tree.id;
+                let remote_worktree_id = tree.remote_id;
                 let connection_id = tree.connection_id;
                 let rpc = tree.rpc.clone();
                 let handles = tree.handles.clone();
@@ -1537,19 +1549,19 @@ impl WorktreeHandle for ModelHandle<Worktree> {
                             .request(
                                 connection_id,
                                 proto::OpenFile {
-                                    worktree_id: worktree_id as u64,
+                                    worktree_id: remote_worktree_id as u64,
                                     path: path.to_string_lossy().to_string(),
                                 },
                             )
                             .await?;
                         let is_deleted = handle.read_with(&cx, |tree, _| {
-                            tree.entry_for_path(&path).is_some() || !tree.path_is_pending(&path)
+                            tree.entry_for_path(&path).is_none() && !tree.path_is_pending(&path)
                         });
                         let new_state = Arc::new(Mutex::new(FileHandleState {
                             path,
                             is_deleted,
                             mtime: Duration::from_secs(response.mtime),
-                            worktree_id,
+                            remote_worktree_id,
                             id: response.id,
                             rpc: Some((connection_id, rpc)),
                         }));