Fix race conditions in updating buffer diffs on git changes (#26409)

Release Notes:

- N/A

---------

Co-authored-by: Cole Miller <m@cole-miller.net>

crates/buffer_diff/src/buffer_diff.rs

@@ -837,8 +837,8 @@ impl BufferDiff {
         language: Option<Arc<Language>>,
         language_registry: Option<Arc<LanguageRegistry>>,
         cx: &mut AsyncApp,
-    ) -> anyhow::Result<Option<Range<Anchor>>> {
+    ) -> anyhow::Result<BufferDiffSnapshot> {
-        let snapshot = if base_text_changed || language_changed {
+        let inner = if base_text_changed || language_changed {
             cx.update(|cx| {
                 Self::build(
                     buffer.clone(),
@@ -860,18 +860,45 @@ impl BufferDiff {
             })?
             .await
         };
-
+        Ok(BufferDiffSnapshot {
-        this.update(cx, |this, _| this.set_state(snapshot, &buffer))
+            inner,
+            secondary_diff: None,
+        })
     }
 
-    pub fn update_diff_from(
+    pub fn set_snapshot(
         &mut self,
         buffer: &text::BufferSnapshot,
-        other: &Entity<Self>,
+        new_snapshot: BufferDiffSnapshot,
+        language_changed: bool,
+        secondary_changed_range: Option<Range<Anchor>>,
         cx: &mut Context<Self>,
     ) -> Option<Range<Anchor>> {
-        let other = other.read(cx).inner.clone();
+        let changed_range = self.set_state(new_snapshot.inner, buffer);
-        self.set_state(other, buffer)
+        if language_changed {
+            cx.emit(BufferDiffEvent::LanguageChanged);
+        }
+
+        let changed_range = match (secondary_changed_range, changed_range) {
+            (None, None) => None,
+            (Some(unstaged_range), None) => self.range_to_hunk_range(unstaged_range, &buffer, cx),
+            (None, Some(uncommitted_range)) => Some(uncommitted_range),
+            (Some(unstaged_range), Some(uncommitted_range)) => {
+                let mut start = uncommitted_range.start;
+                let mut end = uncommitted_range.end;
+                if let Some(unstaged_range) = self.range_to_hunk_range(unstaged_range, &buffer, cx)
+                {
+                    start = unstaged_range.start.min(&uncommitted_range.start, &buffer);
+                    end = unstaged_range.end.max(&uncommitted_range.end, &buffer);
+                }
+                Some(start..end)
+            }
+        };
+
+        cx.emit(BufferDiffEvent::DiffChanged {
+            changed_range: changed_range.clone(),
+        });
+        changed_range
     }
 
     fn set_state(

crates/project/src/buffer_store.rs

@@ -6,7 +6,7 @@ use crate::{
 };
 use ::git::{parse_git_remote_url, BuildPermalinkParams, GitHostingProviderRegistry};
 use anyhow::{anyhow, bail, Context as _, Result};
-use buffer_diff::{BufferDiff, BufferDiffEvent};
+use buffer_diff::BufferDiff;
 use client::Client;
 use collections::{hash_map, HashMap, HashSet};
 use fs::Fs;
@@ -217,9 +217,10 @@ impl BufferDiffState {
             _ => false,
         };
         self.recalculate_diff_task = Some(cx.spawn(|this, mut cx| async move {
-            let mut unstaged_changed_range = None;
+            let mut new_unstaged_diff = None;
             if let Some(unstaged_diff) = &unstaged_diff {
-                unstaged_changed_range = BufferDiff::update_diff(
+                new_unstaged_diff = Some(
+                    BufferDiff::update_diff(
                         unstaged_diff.clone(),
                         buffer.clone(),
                         index,
@@ -229,27 +230,16 @@ impl BufferDiffState {
                         language_registry.clone(),
                         &mut cx,
                     )
-                .await?;
+                    .await?,
-
+                );
-                unstaged_diff.update(&mut cx, |_, cx| {
-                    if language_changed {
-                        cx.emit(BufferDiffEvent::LanguageChanged);
-                    }
-                    if let Some(changed_range) = unstaged_changed_range.clone() {
-                        cx.emit(BufferDiffEvent::DiffChanged {
-                            changed_range: Some(changed_range),
-                        })
-                    }
-                })?;
             }
 
+            let mut new_uncommitted_diff = None;
             if let Some(uncommitted_diff) = &uncommitted_diff {
-                let uncommitted_changed_range =
+                new_uncommitted_diff = if index_matches_head {
-                    if let (Some(unstaged_diff), true) = (&unstaged_diff, index_matches_head) {
+                    new_unstaged_diff.clone()
-                        uncommitted_diff.update(&mut cx, |uncommitted_diff, cx| {
-                            uncommitted_diff.update_diff_from(&buffer, unstaged_diff, cx)
-                        })?
                 } else {
+                    Some(
                         BufferDiff::update_diff(
                             uncommitted_diff.clone(),
                             buffer.clone(),
@@ -260,32 +250,32 @@ impl BufferDiffState {
                             language_registry.clone(),
                             &mut cx,
                         )
-                        .await?
+                        .await?,
+                    )
+                }
+            }
+
+            let unstaged_changed_range = if let Some((unstaged_diff, new_unstaged_diff)) =
+                unstaged_diff.as_ref().zip(new_unstaged_diff.clone())
+            {
+                unstaged_diff.update(&mut cx, |diff, cx| {
+                    diff.set_snapshot(&buffer, new_unstaged_diff, language_changed, None, cx)
+                })?
+            } else {
+                None
             };
 
-                uncommitted_diff.update(&mut cx, |uncommitted_diff, cx| {
+            if let Some((uncommitted_diff, new_uncommitted_diff)) =
-                    if language_changed {
+                uncommitted_diff.as_ref().zip(new_uncommitted_diff.clone())
-                        cx.emit(BufferDiffEvent::LanguageChanged);
-                    }
-                    let changed_range = match (unstaged_changed_range, uncommitted_changed_range) {
-                        (None, None) => None,
-                        (Some(unstaged_range), None) => {
-                            uncommitted_diff.range_to_hunk_range(unstaged_range, &buffer, cx)
-                        }
-                        (None, Some(uncommitted_range)) => Some(uncommitted_range),
-                        (Some(unstaged_range), Some(uncommitted_range)) => {
-                            let mut start = uncommitted_range.start;
-                            let mut end = uncommitted_range.end;
-                            if let Some(unstaged_range) =
-                                uncommitted_diff.range_to_hunk_range(unstaged_range, &buffer, cx)
             {
-                                start = unstaged_range.start.min(&uncommitted_range.start, &buffer);
+                uncommitted_diff.update(&mut cx, |uncommitted_diff, cx| {
-                                end = unstaged_range.end.max(&uncommitted_range.end, &buffer);
+                    uncommitted_diff.set_snapshot(
-                            }
+                        &buffer,
-                            Some(start..end)
+                        new_uncommitted_diff,
-                        }
+                        language_changed,
-                    };
+                        unstaged_changed_range,
-                    cx.emit(BufferDiffEvent::DiffChanged { changed_range });
+                        cx,
+                    );
                 })?;
             }
 
@@ -813,8 +803,7 @@ impl LocalBufferStore {
             let Some(buffer) = buffer.upgrade() else {
                 continue;
             };
-            let buffer = buffer.read(cx);
+            let Some(file) = File::from_dyn(buffer.read(cx).file()) else {
-            let Some(file) = File::from_dyn(buffer.file()) else {
                 continue;
             };
             if file.worktree != worktree_handle {
@@ -825,7 +814,6 @@ impl LocalBufferStore {
                 .iter()
                 .any(|(work_dir, _)| file.path.starts_with(work_dir))
             {
-                let snapshot = buffer.text_snapshot();
                 let has_unstaged_diff = diff_state
                     .unstaged_diff
                     .as_ref()
@@ -835,7 +823,7 @@ impl LocalBufferStore {
                     .as_ref()
                     .is_some_and(|set| set.is_upgradable());
                 diff_state_updates.push((
-                    snapshot.clone(),
+                    buffer,
                     file.path.clone(),
                     has_unstaged_diff.then(|| diff_state.index_text.clone()),
                     has_uncommitted_diff.then(|| diff_state.head_text.clone()),
@@ -854,8 +842,7 @@ impl LocalBufferStore {
                 .background_spawn(async move {
                     diff_state_updates
                         .into_iter()
-                        .filter_map(
+                        .filter_map(|(buffer, path, current_index_text, current_head_text)| {
-                            |(buffer_snapshot, path, current_index_text, current_head_text)| {
                             let local_repo = snapshot.local_repo_for_path(&path)?;
                             let relative_path = local_repo.relativize(&path).ok()?;
                             let index_text = if current_index_text.is_some() {
@@ -880,10 +867,8 @@ impl LocalBufferStore {
                                 }
                             }
 
-                                let diff_bases_change = match (
+                            let diff_bases_change =
-                                    current_index_text.is_some(),
+                                match (current_index_text.is_some(), current_head_text.is_some()) {
-                                    current_head_text.is_some(),
-                                ) {
                                     (true, true) => Some(if index_text == head_text {
                                         DiffBasesChange::SetBoth(head_text)
                                     } else {
@@ -896,17 +881,17 @@ impl LocalBufferStore {
                                     (false, true) => Some(DiffBasesChange::SetHead(head_text)),
                                     (false, false) => None,
                                 };
-                                Some((buffer_snapshot, diff_bases_change))
+
-                            },
+                            Some((buffer, diff_bases_change))
-                        )
+                        })
                         .collect::<Vec<_>>()
                 })
                 .await;
 
             this.update(&mut cx, |this, cx| {
-                for (buffer_snapshot, diff_bases_change) in diff_bases_changes_by_buffer {
+                for (buffer, diff_bases_change) in diff_bases_changes_by_buffer {
                     let Some(OpenBuffer::Complete { diff_state, .. }) =
-                        this.opened_buffers.get_mut(&buffer_snapshot.remote_id())
+                        this.opened_buffers.get_mut(&buffer.read(cx).remote_id())
                     else {
                         continue;
                     };
@@ -917,8 +902,9 @@ impl LocalBufferStore {
                     diff_state.update(cx, |diff_state, cx| {
                         use proto::update_diff_bases::Mode;
 
+                        let buffer = buffer.read(cx);
                         if let Some((client, project_id)) = this.downstream_client.as_ref() {
-                            let buffer_id = buffer_snapshot.remote_id().to_proto();
+                            let buffer_id = buffer.remote_id().to_proto();
                             let (staged_text, committed_text, mode) = match diff_bases_change
                                 .clone()
                             {
@@ -942,8 +928,11 @@ impl LocalBufferStore {
                             client.send(message).log_err();
                         }
 
-                        let _ =
+                        let _ = diff_state.diff_bases_changed(
-                            diff_state.diff_bases_changed(buffer_snapshot, diff_bases_change, cx);
+                            buffer.text_snapshot(),
+                            diff_bases_change,
+                            cx,
+                        );
                     });
                 }
             })