From 90744fb77018062cc38f94098d838ce3ec164b24 Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Fri, 3 Nov 2023 23:02:50 -0700
Subject: [PATCH] working copy: read files ahead when updating

If the commit backend has high latency, it can make a big difference
to read files concurrently. This patch updates the working copy code
to do that in the update code (when reading files from the backend to
write to the working copy). Because our backend at Google reads files
from a local daemon process that already does a lot of prefetching,
this patch doesn't actually help us. I think it's still the right
thing to do for backends that don't do the same kind of
prefetching. It speeds up `jj sparse set --add` by >10x when I disable
the prefetching in our daemon (our `Backend::concurrency()` is 100).
---
 lib/src/local_working_copy.rs | 29 ++++++++++++++++++++---------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/lib/src/local_working_copy.rs b/lib/src/local_working_copy.rs
index 13fa2be7f..7aa9822f3 100644
--- a/lib/src/local_working_copy.rs
+++ b/lib/src/local_working_copy.rs
@@ -1301,22 +1301,34 @@ impl TreeState {
         };
         let mut changed_file_states = Vec::new();
         let mut deleted_files = HashSet::new();
-        let mut diff_stream = old_tree.diff_stream(new_tree, matcher);
-        while let Some((path, diff)) = diff_stream.next().await {
-            let (before, after) = diff?;
+        let mut diff_stream = Box::pin(
+            old_tree
+                .diff_stream(new_tree, matcher)
+                .map(|(path, diff)| async {
+                    match diff {
+                        Ok((before, after)) => {
+                            let result = materialize_tree_value(&self.store, &path, after).await;
+                            (path, result.map(|value| (before.is_present(), value)))
+                        }
+                        Err(err) => (path, Err(err)),
+                    }
+                })
+                .buffered(self.store.concurrency().max(1)),
+        );
+        while let Some((path, data)) = diff_stream.next().await {
+            let (present_before, after) = data?;
             if after.is_absent() {
                 stats.removed_files += 1;
-            } else if before.is_absent() {
+            } else if !present_before {
                 stats.added_files += 1;
             } else {
                 stats.updated_files += 1;
             }
             let disk_path = path.to_fs_path(&self.working_copy_path);
 
-            if before.is_present() {
+            if present_before {
                 fs::remove_file(&disk_path).ok();
-            }
-            if before.is_absent() && disk_path.exists() {
+            } else if disk_path.exists() {
                 changed_file_states.push((path, FileState::placeholder()));
                 stats.skipped_files += 1;
                 continue;
@@ -1330,8 +1342,7 @@ impl TreeState {
                 }
             }
             // TODO: Check that the file has not changed before overwriting/removing it.
-            let materialized = materialize_tree_value(&self.store, &path, after).await?;
-            let file_state = match materialized {
+            let file_state = match after {
                 MaterializedTreeValue::Absent => {
                     let mut parent_dir = disk_path.parent().unwrap();
                     loop {