Allow explicit reload of buffers via `Project::reload_buffers`

Antonio Scandurra created

Change summary

crates/language/src/buffer.rs |  65 +++++++++++--------
crates/language/src/tests.rs  |   8 +
crates/project/src/project.rs |  94 ++++++++++++++++++++++++++++
crates/rpc/proto/zed.proto    |  85 ++++++++++++++-----------
crates/rpc/src/proto.rs       |   4 +
crates/rpc/src/rpc.rs         |   2 
crates/server/src/rpc.rs      | 120 ++++++++++++++++++++++++++++++++++++
7 files changed, 310 insertions(+), 68 deletions(-)

Detailed changes

crates/language/src/buffer.rs 🔗

@@ -498,6 +498,30 @@ impl Buffer {
         cx.notify();
     }
 
+    pub fn reload(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<Option<Transaction>>> {
+        cx.spawn(|this, mut cx| async move {
+            if let Some((new_mtime, new_text)) = this.read_with(&cx, |this, cx| {
+                let file = this.file.as_ref()?.as_local()?;
+                Some((file.mtime(), file.load(cx)))
+            }) {
+                let new_text = new_text.await?;
+                let diff = this
+                    .read_with(&cx, |this, cx| this.diff(new_text.into(), cx))
+                    .await;
+                this.update(&mut cx, |this, cx| {
+                    if let Some(transaction) = this.apply_diff(diff, cx).cloned() {
+                        this.did_reload(this.version(), new_mtime, cx);
+                        Ok(Some(transaction))
+                    } else {
+                        Ok(None)
+                    }
+                })
+            } else {
+                Ok(None)
+            }
+        })
+    }
+
     pub fn did_reload(
         &mut self,
         version: clock::Global,
@@ -543,29 +567,8 @@ impl Buffer {
                 file_changed = true;
 
                 if !self.is_dirty() {
-                    task = cx.spawn(|this, mut cx| {
-                        async move {
-                            let new_text = this.read_with(&cx, |this, cx| {
-                                this.file
-                                    .as_ref()
-                                    .and_then(|file| file.as_local().map(|f| f.load(cx)))
-                            });
-                            if let Some(new_text) = new_text {
-                                let new_text = new_text.await?;
-                                let diff = this
-                                    .read_with(&cx, |this, cx| this.diff(new_text.into(), cx))
-                                    .await;
-                                this.update(&mut cx, |this, cx| {
-                                    if this.apply_diff(diff, cx) {
-                                        this.did_reload(this.version(), new_mtime, cx);
-                                    }
-                                });
-                            }
-                            Ok(())
-                        }
-                        .log_err()
-                        .map(drop)
-                    });
+                    let reload = self.reload(cx).log_err().map(drop);
+                    task = cx.foreground().spawn(reload);
                 }
             }
         }
@@ -902,8 +905,13 @@ impl Buffer {
         })
     }
 
-    pub(crate) fn apply_diff(&mut self, diff: Diff, cx: &mut ModelContext<Self>) -> bool {
+    pub(crate) fn apply_diff(
+        &mut self,
+        diff: Diff,
+        cx: &mut ModelContext<Self>,
+    ) -> Option<&Transaction> {
         if self.version == diff.base_version {
+            self.finalize_last_transaction();
             self.start_transaction();
             let mut offset = diff.start_offset;
             for (tag, len) in diff.changes {
@@ -924,10 +932,13 @@ impl Buffer {
                     }
                 }
             }
-            self.end_transaction(cx);
-            true
+            if self.end_transaction(cx).is_some() {
+                self.finalize_last_transaction()
+            } else {
+                None
+            }
         } else {
-            false
+            None
         }
     }
 

crates/language/src/tests.rs 🔗

@@ -136,12 +136,16 @@ async fn test_apply_diff(cx: &mut gpui::TestAppContext) {
 
     let text = "a\nccc\ndddd\nffffff\n";
     let diff = buffer.read_with(cx, |b, cx| b.diff(text.into(), cx)).await;
-    buffer.update(cx, |b, cx| b.apply_diff(diff, cx));
+    buffer.update(cx, |buffer, cx| {
+        buffer.apply_diff(diff, cx).unwrap();
+    });
     cx.read(|cx| assert_eq!(buffer.read(cx).text(), text));
 
     let text = "a\n1\n\nccc\ndd2dd\nffffff\n";
     let diff = buffer.read_with(cx, |b, cx| b.diff(text.into(), cx)).await;
-    buffer.update(cx, |b, cx| b.apply_diff(diff, cx));
+    buffer.update(cx, |buffer, cx| {
+        buffer.apply_diff(diff, cx).unwrap();
+    });
     cx.read(|cx| assert_eq!(buffer.read(cx).text(), text));
 }
 

crates/project/src/project.rs 🔗

@@ -270,6 +270,7 @@ impl Project {
         client.add_model_message_handler(Self::handle_update_worktree);
         client.add_model_request_handler(Self::handle_apply_additional_edits_for_completion);
         client.add_model_request_handler(Self::handle_apply_code_action);
+        client.add_model_request_handler(Self::handle_reload_buffers);
         client.add_model_request_handler(Self::handle_format_buffers);
         client.add_model_request_handler(Self::handle_get_code_actions);
         client.add_model_request_handler(Self::handle_get_completions);
@@ -1973,6 +1974,70 @@ impl Project {
         Ok(())
     }
 
+    pub fn reload_buffers(
+        &self,
+        buffers: HashSet<ModelHandle<Buffer>>,
+        push_to_history: bool,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<ProjectTransaction>> {
+        let mut local_buffers = Vec::new();
+        let mut remote_buffers = None;
+        for buffer_handle in buffers {
+            let buffer = buffer_handle.read(cx);
+            if buffer.is_dirty() {
+                if let Some(file) = File::from_dyn(buffer.file()) {
+                    if file.is_local() {
+                        local_buffers.push(buffer_handle);
+                    } else {
+                        remote_buffers.get_or_insert(Vec::new()).push(buffer_handle);
+                    }
+                }
+            }
+        }
+
+        let remote_buffers = self.remote_id().zip(remote_buffers);
+        let client = self.client.clone();
+
+        cx.spawn(|this, mut cx| async move {
+            let mut project_transaction = ProjectTransaction::default();
+
+            if let Some((project_id, remote_buffers)) = remote_buffers {
+                let response = client
+                    .request(proto::ReloadBuffers {
+                        project_id,
+                        buffer_ids: remote_buffers
+                            .iter()
+                            .map(|buffer| buffer.read_with(&cx, |buffer, _| buffer.remote_id()))
+                            .collect(),
+                    })
+                    .await?
+                    .transaction
+                    .ok_or_else(|| anyhow!("missing transaction"))?;
+                project_transaction = this
+                    .update(&mut cx, |this, cx| {
+                        this.deserialize_project_transaction(response, push_to_history, cx)
+                    })
+                    .await?;
+            }
+
+            for buffer in local_buffers {
+                let transaction = buffer
+                    .update(&mut cx, |buffer, cx| buffer.reload(cx))
+                    .await?;
+                buffer.update(&mut cx, |buffer, cx| {
+                    if let Some(transaction) = transaction {
+                        if !push_to_history {
+                            buffer.forget_transaction(transaction.id);
+                        }
+                        project_transaction.0.insert(cx.handle(), transaction);
+                    }
+                });
+            }
+
+            Ok(project_transaction)
+        })
+    }
+
     pub fn format(
         &self,
         buffers: HashSet<ModelHandle<Buffer>>,
@@ -3667,6 +3732,35 @@ impl Project {
         })
     }
 
+    async fn handle_reload_buffers(
+        this: ModelHandle<Self>,
+        envelope: TypedEnvelope<proto::ReloadBuffers>,
+        _: Arc<Client>,
+        mut cx: AsyncAppContext,
+    ) -> Result<proto::ReloadBuffersResponse> {
+        let sender_id = envelope.original_sender_id()?;
+        let reload = this.update(&mut cx, |this, cx| {
+            let mut buffers = HashSet::default();
+            for buffer_id in &envelope.payload.buffer_ids {
+                buffers.insert(
+                    this.opened_buffers
+                        .get(buffer_id)
+                        .map(|buffer| buffer.upgrade(cx).unwrap())
+                        .ok_or_else(|| anyhow!("unknown buffer id {}", buffer_id))?,
+                );
+            }
+            Ok::<_, anyhow::Error>(this.reload_buffers(buffers, false, cx))
+        })?;
+
+        let project_transaction = reload.await?;
+        let project_transaction = this.update(&mut cx, |this, cx| {
+            this.serialize_project_transaction_for_peer(project_transaction, sender_id, cx)
+        });
+        Ok(proto::ReloadBuffersResponse {
+            transaction: Some(project_transaction),
+        })
+    }
+
     async fn handle_format_buffers(
         this: ModelHandle<Self>,
         envelope: TypedEnvelope<proto::FormatBuffers>,

crates/rpc/proto/zed.proto 🔗

@@ -48,43 +48,45 @@ message Envelope {
         SaveBuffer save_buffer = 40;
         BufferSaved buffer_saved = 41;
         BufferReloaded buffer_reloaded = 42;
-        FormatBuffers format_buffers = 43;
-        FormatBuffersResponse format_buffers_response = 44;
-        GetCompletions get_completions = 45;
-        GetCompletionsResponse get_completions_response = 46;
-        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 47;
-        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 48;
-        GetCodeActions get_code_actions = 49;
-        GetCodeActionsResponse get_code_actions_response = 50;
-        ApplyCodeAction apply_code_action = 51;
-        ApplyCodeActionResponse apply_code_action_response = 52;
-        PrepareRename prepare_rename = 53;
-        PrepareRenameResponse prepare_rename_response = 54;
-        PerformRename perform_rename = 55;
-        PerformRenameResponse perform_rename_response = 56;
-        SearchProject search_project = 57;
-        SearchProjectResponse search_project_response = 58;
-
-        GetChannels get_channels = 59;
-        GetChannelsResponse get_channels_response = 60;
-        JoinChannel join_channel = 61;
-        JoinChannelResponse join_channel_response = 62;
-        LeaveChannel leave_channel = 63;
-        SendChannelMessage send_channel_message = 64;
-        SendChannelMessageResponse send_channel_message_response = 65;
-        ChannelMessageSent channel_message_sent = 66;
-        GetChannelMessages get_channel_messages = 67;
-        GetChannelMessagesResponse get_channel_messages_response = 68;
-
-        UpdateContacts update_contacts = 69;
-
-        GetUsers get_users = 70;
-        GetUsersResponse get_users_response = 71;
-
-        Follow follow = 72;
-        FollowResponse follow_response = 73;
-        UpdateFollowers update_followers = 74;
-        Unfollow unfollow = 75;
+        ReloadBuffers reload_buffers = 43;
+        ReloadBuffersResponse reload_buffers_response = 44;
+        FormatBuffers format_buffers = 45;
+        FormatBuffersResponse format_buffers_response = 46;
+        GetCompletions get_completions = 47;
+        GetCompletionsResponse get_completions_response = 48;
+        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 49;
+        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 50;
+        GetCodeActions get_code_actions = 51;
+        GetCodeActionsResponse get_code_actions_response = 52;
+        ApplyCodeAction apply_code_action = 53;
+        ApplyCodeActionResponse apply_code_action_response = 54;
+        PrepareRename prepare_rename = 55;
+        PrepareRenameResponse prepare_rename_response = 56;
+        PerformRename perform_rename = 57;
+        PerformRenameResponse perform_rename_response = 58;
+        SearchProject search_project = 59;
+        SearchProjectResponse search_project_response = 60;
+
+        GetChannels get_channels = 61;
+        GetChannelsResponse get_channels_response = 62;
+        JoinChannel join_channel = 63;
+        JoinChannelResponse join_channel_response = 64;
+        LeaveChannel leave_channel = 65;
+        SendChannelMessage send_channel_message = 66;
+        SendChannelMessageResponse send_channel_message_response = 67;
+        ChannelMessageSent channel_message_sent = 68;
+        GetChannelMessages get_channel_messages = 69;
+        GetChannelMessagesResponse get_channel_messages_response = 70;
+
+        UpdateContacts update_contacts = 71;
+
+        GetUsers get_users = 72;
+        GetUsersResponse get_users_response = 73;
+
+        Follow follow = 74;
+        FollowResponse follow_response = 75;
+        UpdateFollowers update_followers = 76;
+        Unfollow unfollow = 77;
     }
 }
 
@@ -299,6 +301,15 @@ message BufferReloaded {
     Timestamp mtime = 4;
 }
 
+message ReloadBuffers {
+    uint64 project_id = 1;
+    repeated uint64 buffer_ids = 2;
+}
+
+message ReloadBuffersResponse {
+    ProjectTransaction transaction = 1;
+}
+
 message FormatBuffers {
     uint64 project_id = 1;
     repeated uint64 buffer_ids = 2;

crates/rpc/src/proto.rs 🔗

@@ -190,6 +190,8 @@ messages!(
     (Ping, Foreground),
     (RegisterProject, Foreground),
     (RegisterWorktree, Foreground),
+    (ReloadBuffers, Foreground),
+    (ReloadBuffersResponse, Foreground),
     (RemoveProjectCollaborator, Foreground),
     (SaveBuffer, Foreground),
     (SearchProject, Background),
@@ -237,6 +239,7 @@ request_messages!(
     (PrepareRename, PrepareRenameResponse),
     (RegisterProject, RegisterProjectResponse),
     (RegisterWorktree, Ack),
+    (ReloadBuffers, ReloadBuffersResponse),
     (SaveBuffer, BufferSaved),
     (SearchProject, SearchProjectResponse),
     (SendChannelMessage, SendChannelMessageResponse),
@@ -268,6 +271,7 @@ entity_messages!(
     OpenBufferForSymbol,
     PerformRename,
     PrepareRename,
+    ReloadBuffers,
     RemoveProjectCollaborator,
     SaveBuffer,
     SearchProject,

crates/rpc/src/rpc.rs 🔗

@@ -5,4 +5,4 @@ pub mod proto;
 pub use conn::Connection;
 pub use peer::*;
 
-pub const PROTOCOL_VERSION: u32 = 12;
+pub const PROTOCOL_VERSION: u32 = 13;

crates/server/src/rpc.rs 🔗

@@ -102,6 +102,7 @@ impl Server {
             .add_request_handler(Server::forward_project_request::<proto::ApplyCodeAction>)
             .add_request_handler(Server::forward_project_request::<proto::PrepareRename>)
             .add_request_handler(Server::forward_project_request::<proto::PerformRename>)
+            .add_request_handler(Server::forward_project_request::<proto::ReloadBuffers>)
             .add_request_handler(Server::forward_project_request::<proto::FormatBuffers>)
             .add_request_handler(Server::update_buffer)
             .add_message_handler(Server::update_buffer_file)
@@ -1089,7 +1090,7 @@ mod tests {
     use gpui::{executor, geometry::vector::vec2f, ModelHandle, TestAppContext, ViewHandle};
     use language::{
         tree_sitter_rust, Diagnostic, DiagnosticEntry, Language, LanguageConfig, LanguageRegistry,
-        LanguageServerConfig, OffsetRangeExt, Point, ToLspPosition,
+        LanguageServerConfig, OffsetRangeExt, Point, Rope, ToLspPosition,
     };
     use lsp;
     use parking_lot::Mutex;
@@ -2460,6 +2461,123 @@ mod tests {
             .await;
     }
 
+    #[gpui::test(iterations = 10)]
+    async fn test_reloading_buffer_manually(cx_a: &mut TestAppContext, cx_b: &mut TestAppContext) {
+        cx_a.foreground().forbid_parking();
+        let lang_registry = Arc::new(LanguageRegistry::test());
+        let fs = FakeFs::new(cx_a.background());
+
+        // Connect to a server as 2 clients.
+        let mut server = TestServer::start(cx_a.foreground(), cx_a.background()).await;
+        let client_a = server.create_client(cx_a, "user_a").await;
+        let client_b = server.create_client(cx_b, "user_b").await;
+
+        // Share a project as client A
+        fs.insert_tree(
+            "/a",
+            json!({
+                ".zed.toml": r#"collaborators = ["user_b"]"#,
+                "a.rs": "let one = 1;",
+            }),
+        )
+        .await;
+        let project_a = cx_a.update(|cx| {
+            Project::local(
+                client_a.clone(),
+                client_a.user_store.clone(),
+                lang_registry.clone(),
+                fs.clone(),
+                cx,
+            )
+        });
+        let (worktree_a, _) = project_a
+            .update(cx_a, |p, cx| {
+                p.find_or_create_local_worktree("/a", true, cx)
+            })
+            .await
+            .unwrap();
+        worktree_a
+            .read_with(cx_a, |tree, _| tree.as_local().unwrap().scan_complete())
+            .await;
+        let project_id = project_a.update(cx_a, |p, _| p.next_remote_id()).await;
+        let worktree_id = worktree_a.read_with(cx_a, |tree, _| tree.id());
+        project_a.update(cx_a, |p, cx| p.share(cx)).await.unwrap();
+        let buffer_a = project_a
+            .update(cx_a, |p, cx| p.open_buffer((worktree_id, "a.rs"), cx))
+            .await
+            .unwrap();
+
+        // Join the worktree as client B.
+        let project_b = Project::remote(
+            project_id,
+            client_b.clone(),
+            client_b.user_store.clone(),
+            lang_registry.clone(),
+            fs.clone(),
+            &mut cx_b.to_async(),
+        )
+        .await
+        .unwrap();
+
+        let buffer_b = cx_b
+            .background()
+            .spawn(project_b.update(cx_b, |p, cx| p.open_buffer((worktree_id, "a.rs"), cx)))
+            .await
+            .unwrap();
+        buffer_b.update(cx_b, |buffer, cx| {
+            buffer.edit([4..7], "six", cx);
+            buffer.edit([10..11], "6", cx);
+            assert_eq!(buffer.text(), "let six = 6;");
+            assert!(buffer.is_dirty());
+            assert!(!buffer.has_conflict());
+        });
+        buffer_a
+            .condition(cx_a, |buffer, _| buffer.text() == "let six = 6;")
+            .await;
+
+        fs.save(Path::new("/a/a.rs"), &Rope::from("let seven = 7;"))
+            .await
+            .unwrap();
+        buffer_a
+            .condition(cx_a, |buffer, _| buffer.has_conflict())
+            .await;
+        buffer_b
+            .condition(cx_b, |buffer, _| buffer.has_conflict())
+            .await;
+
+        project_b
+            .update(cx_b, |project, cx| {
+                project.reload_buffers(HashSet::from_iter([buffer_b.clone()]), true, cx)
+            })
+            .await
+            .unwrap();
+        buffer_a.read_with(cx_a, |buffer, _| {
+            assert_eq!(buffer.text(), "let seven = 7;");
+            assert!(!buffer.is_dirty());
+            assert!(!buffer.has_conflict());
+        });
+        buffer_b.read_with(cx_b, |buffer, _| {
+            assert_eq!(buffer.text(), "let seven = 7;");
+            assert!(!buffer.is_dirty());
+            assert!(!buffer.has_conflict());
+        });
+
+        buffer_a.update(cx_a, |buffer, cx| {
+            // Undoing on the host is a no-op when the reload was initiated by the guest.
+            buffer.undo(cx);
+            assert_eq!(buffer.text(), "let seven = 7;");
+            assert!(!buffer.is_dirty());
+            assert!(!buffer.has_conflict());
+        });
+        buffer_b.update(cx_b, |buffer, cx| {
+            // Undoing on the guest rolls back the buffer to before it was reloaded but the conflict gets cleared.
+            buffer.undo(cx);
+            assert_eq!(buffer.text(), "let six = 6;");
+            assert!(buffer.is_dirty());
+            assert!(!buffer.has_conflict());
+        });
+    }
+
     #[gpui::test(iterations = 10)]
     async fn test_formatting_buffer(cx_a: &mut TestAppContext, cx_b: &mut TestAppContext) {
         cx_a.foreground().forbid_parking();