From 774c1123ef7615b2c2eb8b0ce628a99c10069f9e Mon Sep 17 00:00:00 2001
From: Jun Wu <quark@meta.com>
Date: Tue, 16 Jan 2024 10:16:16 -0800
Subject: [PATCH] dag/import_pull: split segments for more efficient high level
 segments

Summary:
See the previous diffs for context. This is intended to further defrag segments
from the new server algorithms to ensure performance of the local graph.

Reviewed By: zzl0

Differential Revision: D52763359

fbshipit-source-id: 9f6067ebe7b1bca42597ab8c47709a1dac78e151
---
 eden/scm/lib/dag/src/namedag.rs           | 11 +++++--
 eden/scm/lib/dag/src/tests/test_sparse.rs | 35 +++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/eden/scm/lib/dag/src/namedag.rs b/eden/scm/lib/dag/src/namedag.rs
index d7abac42f9..278a703c95 100644
--- a/eden/scm/lib/dag/src/namedag.rs
+++ b/eden/scm/lib/dag/src/namedag.rs
@@ -813,7 +813,7 @@ where
             idmap_by_name: BTreeMap<&'a VertexName, Id>,
             idmap_by_id: &'a BTreeMap<Id, VertexName>,
         }
-        let server = ServerState {
+        let mut server = ServerState {
             seg_by_high: clone_data
                 .flat_segments
                 .segments
@@ -934,7 +934,14 @@ where
             }
 
             while let Some(server_high) = stack.pop() {
-                let server_seg = &server.seg_containing_id(server_high)?;
+                let mut server_seg = server.seg_containing_id(server_high)?;
+                if server_high < server_seg.high {
+                    // Split the segment for more efficient high level segments.
+                    let seg_high = server_seg.high;
+                    server.split_seg(seg_high, server_high);
+                    server_seg = server.seg_containing_id(server_high)?;
+                    assert_eq!(server_high, server_seg.high);
+                }
                 let high_vertex = server.name_by_id(server_high);
                 let client_high_id = new
                     .map
diff --git a/eden/scm/lib/dag/src/tests/test_sparse.rs b/eden/scm/lib/dag/src/tests/test_sparse.rs
index 522368ac1f..68d2369de8 100644
--- a/eden/scm/lib/dag/src/tests/test_sparse.rs
+++ b/eden/scm/lib/dag/src/tests/test_sparse.rs
@@ -298,6 +298,41 @@ async fn test_pull_remap() {
     );
 }
 
+#[tokio::test]
+async fn test_pull_split() {
+    // Check that the client can split server segments for defragmentation purpose.
+    // This just tests the basic feature.
+    let mut server = TestDag::new();
+    server.drawdag(
+        r#"
+        A1..A9--B1..B9--D  B5-D  C5-D
+               \      /
+                C1..C9
+    "#,
+        &["D"],
+    );
+    let data = server.export_pull_data("", "D").await.unwrap();
+
+    // Import only the B5 branch without importing the segment containing B5 (head: B9).
+    let mut client = server.client().await;
+    client.import_pull_data(data.clone(), "B5").await.unwrap();
+    assert!(!client.contains_vertex_locally("B9"));
+    // There is only one flat segment.
+    assert_eq!(
+        client.dump_state().await,
+        "<spans [A1:B5+0:13]>\nLv0: RH0-13[]\n0->A1 8->A9 13->B5"
+    );
+
+    // Import only the C5 branch without importing the segment containing C5 (head: C9).
+    let mut client = server.client().await;
+    client.import_pull_data(data, "C5").await.unwrap();
+    assert!(!client.contains_vertex_locally("C9"));
+    assert_eq!(
+        client.dump_state().await,
+        "<spans [A1:C5+0:13]>\nLv0: RH0-13[]\n0->A1 8->A9 9->C1 13->C5"
+    );
+}
+
 #[tokio::test]
 async fn test_pull_overlap() {
     let mut server = TestDag::new();