From afb0fd609bc9fa10aff4ef30385e4fb4f8c26a78 Mon Sep 17 00:00:00 2001
From: Conrad Irwin <conrad.irwin@gmail.com>
Date: Wed, 26 Feb 2025 12:38:16 -0700
Subject: [PATCH] Chunk git status entries (#25627)

Prevents us trying to write 5M untracked files to postgres in one commit

Release Notes:

- N/A *or* Added/Fixed/Improved ...
---
 crates/proto/Cargo.toml   |  1 -
 crates/proto/src/proto.rs | 56 +++++++++++++++++++++++----------------
 2 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/crates/proto/Cargo.toml b/crates/proto/Cargo.toml
index 9af37d9eb2ecda..5cb898d7fb2cb9 100644
--- a/crates/proto/Cargo.toml
+++ b/crates/proto/Cargo.toml
@@ -18,7 +18,6 @@ doctest = false
 
 [dependencies]
 anyhow.workspace = true
-collections.workspace = true
 prost.workspace = true
 serde.workspace = true
 
diff --git a/crates/proto/src/proto.rs b/crates/proto/src/proto.rs
index 13c5229ec3f094..d3384ea1208d5b 100644
--- a/crates/proto/src/proto.rs
+++ b/crates/proto/src/proto.rs
@@ -7,7 +7,6 @@ mod typed_envelope;
 pub use error::*;
 pub use typed_envelope::*;
 
-use collections::HashMap;
 pub use prost::{DecodeError, Message};
 use serde::Serialize;
 use std::{
@@ -746,16 +745,10 @@ pub const MAX_WORKTREE_UPDATE_MAX_CHUNK_SIZE: usize = 2;
 pub const MAX_WORKTREE_UPDATE_MAX_CHUNK_SIZE: usize = 256;
 
 pub fn split_worktree_update(mut message: UpdateWorktree) -> impl Iterator<Item = UpdateWorktree> {
-    let mut done_files = false;
-
-    let mut repository_map = message
-        .updated_repositories
-        .into_iter()
-        .map(|repo| (repo.work_directory_id, repo))
-        .collect::<HashMap<_, _>>();
+    let mut done = false;
 
     iter::from_fn(move || {
-        if done_files {
+        if done {
             return None;
         }
 
@@ -777,28 +770,45 @@ pub fn split_worktree_update(mut message: UpdateWorktree) -> impl Iterator<Item
             .drain(..removed_entries_chunk_size)
             .collect();
 
-        done_files = message.updated_entries.is_empty() && message.removed_entries.is_empty();
-
         let mut updated_repositories = Vec::new();
-
-        if !repository_map.is_empty() {
-            for entry in &updated_entries {
-                if let Some(repo) = repository_map.remove(&entry.id) {
-                    updated_repositories.push(repo);
-                }
+        let mut limit = MAX_WORKTREE_UPDATE_MAX_CHUNK_SIZE;
+        while let Some(repo) = message.updated_repositories.first_mut() {
+            let updated_statuses_limit = cmp::min(repo.updated_statuses.len(), limit);
+            let removed_statuses_limit = cmp::min(repo.removed_statuses.len(), limit);
+
+            updated_repositories.push(RepositoryEntry {
+                work_directory_id: repo.work_directory_id,
+                branch: repo.branch.clone(),
+                branch_summary: repo.branch_summary.clone(),
+                updated_statuses: repo
+                    .updated_statuses
+                    .drain(..updated_statuses_limit)
+                    .collect(),
+                removed_statuses: repo
+                    .removed_statuses
+                    .drain(..removed_statuses_limit)
+                    .collect(),
+                current_merge_conflicts: repo.current_merge_conflicts.clone(),
+            });
+            if repo.removed_statuses.is_empty() && repo.updated_statuses.is_empty() {
+                message.updated_repositories.remove(0);
+            }
+            limit = limit.saturating_sub(removed_statuses_limit + updated_statuses_limit);
+            if limit == 0 {
+                break;
             }
         }
 
-        let removed_repositories = if done_files {
+        done = message.updated_entries.is_empty()
+            && message.removed_entries.is_empty()
+            && message.updated_repositories.is_empty();
+
+        let removed_repositories = if done {
             mem::take(&mut message.removed_repositories)
         } else {
             Default::default()
         };
 
-        if done_files {
-            updated_repositories.extend(mem::take(&mut repository_map).into_values());
-        }
-
         Some(UpdateWorktree {
             project_id: message.project_id,
             worktree_id: message.worktree_id,
@@ -807,7 +817,7 @@ pub fn split_worktree_update(mut message: UpdateWorktree) -> impl Iterator<Item
             updated_entries,
             removed_entries,
             scan_id: message.scan_id,
-            is_last_update: done_files && message.is_last_update,
+            is_last_update: done && message.is_last_update,
             updated_repositories,
             removed_repositories,
         })