From ecb2847ef5deb05a7b5784cb6c6d8c07c6b67389 Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Tue, 3 Sep 2024 21:52:54 -0700
Subject: [PATCH] store: make `write_tree()` async

---
 Cargo.lock              | 1 +
 lib/src/merged_tree.rs  | 9 +++++----
 lib/src/store.rs        | 4 ++--
 lib/src/tree.rs         | 3 ++-
 lib/src/tree_builder.rs | 7 +++++--
 5 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4bde3c46f..729ab738f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3102,6 +3102,7 @@ dependencies = [
  "hex",
  "itertools 0.13.0",
  "jj-lib",
+ "pollster",
  "rand",
  "tempfile",
 ]
diff --git a/lib/src/merged_tree.rs b/lib/src/merged_tree.rs
index c2b0dc8af..a5eb4d901 100644
--- a/lib/src/merged_tree.rs
+++ b/lib/src/merged_tree.rs
@@ -35,6 +35,7 @@ use futures::Stream;
 use futures::TryStreamExt;
 use itertools::EitherOrBoth;
 use itertools::Itertools;
+use pollster::FutureExt;
 
 use crate::backend;
 use crate::backend::BackendResult;
@@ -463,7 +464,7 @@ fn merge_trees(merge: &Merge<Tree>) -> BackendResult<Merge<Tree>> {
         };
     }
     if conflicts.is_empty() {
-        let new_tree_id = store.write_tree(dir, new_tree)?;
+        let new_tree_id = store.write_tree(dir, new_tree).block_on()?;
         Ok(Merge::resolved(new_tree_id))
     } else {
         // For each side of the conflict, overwrite the entries in `new_tree` with the
@@ -475,7 +476,7 @@ fn merge_trees(merge: &Merge<Tree>) -> BackendResult<Merge<Tree>> {
             for (basename, path_conflict) in &mut conflicts {
                 new_tree.set_or_remove(basename, path_conflict.next().unwrap());
             }
-            let tree = store.write_tree(dir, new_tree.clone())?;
+            let tree = store.write_tree(dir, new_tree.clone()).block_on()?;
             new_trees.push(tree);
         }
         Ok(Merge::from_vec(new_trees))
@@ -799,13 +800,13 @@ impl Iterator for TreeDiffIterator<'_> {
                         return Some(TreeDiffEntry {
                             path,
                             values: Err(before_err),
-                        })
+                        });
                     }
                     (_, Err(after_err)) => {
                         return Some(TreeDiffEntry {
                             path,
                             values: Err(after_err),
-                        })
+                        });
                     }
                 };
                 let subdir =
diff --git a/lib/src/store.rs b/lib/src/store.rs
index 47feb23f7..8f62b01ad 100644
--- a/lib/src/store.rs
+++ b/lib/src/store.rs
@@ -215,12 +215,12 @@ impl Store {
         }
     }
 
-    pub fn write_tree(
+    pub async fn write_tree(
         self: &Arc<Self>,
         path: &RepoPath,
         tree: backend::Tree,
     ) -> BackendResult<Tree> {
-        let tree_id = self.backend.write_tree(path, &tree).block_on()?;
+        let tree_id = self.backend.write_tree(path, &tree).await?;
         let data = Arc::new(tree);
         {
             let mut locked_cache = self.tree_cache.lock().unwrap();
diff --git a/lib/src/tree.rs b/lib/src/tree.rs
index 68d7b7424..0c59e18a5 100644
--- a/lib/src/tree.rs
+++ b/lib/src/tree.rs
@@ -23,6 +23,7 @@ use std::io::Read;
 use std::sync::Arc;
 
 use itertools::Itertools;
+use pollster::FutureExt;
 use tracing::instrument;
 
 use crate::backend;
@@ -330,7 +331,7 @@ pub fn merge_trees(side1_tree: &Tree, base_tree: &Tree, side2_tree: &Tree) -> Ba
             new_tree.set_or_remove(basename, new_value);
         }
     }
-    store.write_tree(dir, new_tree)
+    store.write_tree(dir, new_tree).block_on()
 }
 
 /// Returns `Some(TreeId)` if this is a directory or missing. If it's missing,
diff --git a/lib/src/tree_builder.rs b/lib/src/tree_builder.rs
index a7dc75743..4a32f6dde 100644
--- a/lib/src/tree_builder.rs
+++ b/lib/src/tree_builder.rs
@@ -17,6 +17,8 @@
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
+use pollster::FutureExt;
+
 use crate::backend;
 use crate::backend::BackendResult;
 use crate::backend::TreeId;
@@ -95,6 +97,7 @@ impl TreeBuilder {
 
         // Write trees in reverse lexicographical order, starting with trees without
         // children.
+        // TODO: Writing trees concurrently should help on high-latency backends
         let store = &self.store;
         while let Some((dir, tree)) = trees_to_write.pop_last() {
             if let Some((parent, basename)) = dir.split() {
@@ -106,13 +109,13 @@ impl TreeBuilder {
                         // Entry would have been replaced with file (see above)
                     }
                 } else {
-                    let tree = store.write_tree(&dir, tree)?;
+                    let tree = store.write_tree(&dir, tree).block_on()?;
                     parent_tree.set(basename.to_owned(), TreeValue::Tree(tree.id().clone()));
                 }
             } else {
                 // We're writing the root tree. Write it even if empty. Return its id.
                 assert!(trees_to_write.is_empty());
-                let written_tree = store.write_tree(&dir, tree)?;
+                let written_tree = store.write_tree(&dir, tree).block_on()?;
                 return Ok(written_tree.id().clone());
             }
         }