dag: move client Dag construction to TestDag

Summary: This will make it easier to test client / server dags in upcoming changes.

Reviewed By: sfilipco

Differential Revision: D27629318

fbshipit-source-id: e3137654613aa3208a8f2e4b9f4ddfe73871f2c5

eden/scm/lib/dag/src/namedag.rs

@@ -457,7 +457,7 @@ where
     }
 
     /// Attempt to get a snapshot of this graph.
-    fn try_snapshot(&self) -> Result<Arc<Self>> {
+    pub(crate) fn try_snapshot(&self) -> Result<Arc<Self>> {
         if let Some(s) = self.snapshot.read().deref() {
             if s.dag.version() == self.dag.version() {
                 return Ok(Arc::clone(s));

eden/scm/lib/dag/src/tests/test_dag.rs

@@ -8,20 +8,28 @@
 use crate::nameset::SyncNameSetQuery;
 use crate::ops::DagAddHeads;
 use crate::ops::DagAlgorithm;
+use crate::ops::DagExportCloneData;
+use crate::ops::DagImportCloneData;
 use crate::ops::DagPersistent;
 use crate::ops::IdConvert;
+use crate::protocol;
+use crate::protocol::RemoteIdConvertProtocol;
 use crate::render::render_namedag;
 use crate::NameDag;
+use crate::Result;
 use crate::Vertex;
 use nonblocking::non_blocking_result;
+use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::sync::Arc;
 
 /// Dag structure for testing purpose.
 pub struct TestDag {
     pub dag: NameDag,
     pub seg_size: usize,
     pub dir: tempfile::TempDir,
+    pub output: Arc<Mutex<Vec<String>>>,
 }
 
 impl TestDag {
@@ -35,7 +43,12 @@ impl TestDag {
     pub fn new_with_segment_size(seg_size: usize) -> Self {
         let dir = tempfile::tempdir().unwrap();
         let dag = NameDag::open(dir.path().join("n")).unwrap();
-        Self { dir, dag, seg_size }
+        Self {
+            dir,
+            dag,
+            seg_size,
+            output: Default::default(),
+        }
     }
 
     /// Add vertexes to the graph.
@@ -97,6 +110,39 @@ impl TestDag {
         .unwrap()
     }
 
+    /// Use this DAG as the "server", return the "client" Dag that has lazy Vertexes.
+    pub async fn client(&self) -> TestDag {
+        let data = self.dag.export_clone_data().await.unwrap();
+        let mut client = TestDag::new();
+        client.dag.import_clone_data(data).await.unwrap();
+        let remote = self.remote_protocol(client.output.clone());
+        client.dag.set_remote_protocol(remote);
+        client
+    }
+
+    /// Remote protocol used to resolve Id <-> Vertex remotely using the test dag
+    /// as the "server".
+    ///
+    /// Logs of the remote access will be written to `output`.
+    pub fn remote_protocol(
+        &self,
+        output: Arc<Mutex<Vec<String>>>,
+    ) -> Arc<dyn RemoteIdConvertProtocol> {
+        let remote = ProtocolMonitor {
+            inner: Box::new(self.dag.try_snapshot().unwrap()),
+            output,
+        };
+        Arc::new(remote)
+    }
+
+    /// Output of remote protocols since the last call.
+    pub fn output(&self) -> Vec<String> {
+        let mut result = Vec::new();
+        let mut output = self.output.lock();
+        std::mem::swap(&mut result, &mut *output);
+        result
+    }
+
     fn validate(&self) {
         // All vertexes should be accessible, and round-trip through IdMap.
         for v in non_blocking_result(self.dag.all()).unwrap().iter().unwrap() {
@@ -108,6 +154,35 @@ impl TestDag {
     }
 }
 
+struct ProtocolMonitor {
+    inner: Box<dyn RemoteIdConvertProtocol>,
+    output: Arc<Mutex<Vec<String>>>,
+}
+
+#[async_trait::async_trait]
+impl RemoteIdConvertProtocol for ProtocolMonitor {
+    async fn resolve_names_to_relative_paths(
+        &self,
+        heads: Vec<Vertex>,
+        names: Vec<Vertex>,
+    ) -> Result<Vec<(protocol::AncestorPath, Vec<Vertex>)>> {
+        let msg = format!("resolve names: {:?}, heads: {:?}", &names, &heads);
+        self.output.lock().push(msg);
+        self.inner
+            .resolve_names_to_relative_paths(heads, names)
+            .await
+    }
+
+    async fn resolve_relative_paths_to_names(
+        &self,
+        paths: Vec<protocol::AncestorPath>,
+    ) -> Result<Vec<(protocol::AncestorPath, Vec<Vertex>)>> {
+        let msg = format!("resolve paths: {:?}", &paths);
+        self.output.lock().push(msg);
+        self.inner.resolve_relative_paths_to_names(paths).await
+    }
+}
+
 fn get_heads_and_parents_func_from_ascii(
     text: &str,
 ) -> (Vec<Vertex>, HashMap<Vertex, Vec<Vertex>>) {

eden/scm/lib/dag/src/tests/test_sparse.rs

@@ -7,46 +7,10 @@
 
 use super::TestDag;
 use crate::ops::DagAlgorithm;
-use crate::ops::DagExportCloneData;
-use crate::ops::DagImportCloneData;
 use crate::ops::IdConvert;
-use crate::protocol;
-use crate::protocol::RemoteIdConvertProtocol;
 use crate::Id;
-use crate::Result;
 use crate::VertexName;
 use futures::TryStreamExt;
-use parking_lot::Mutex;
-use std::sync::Arc;
-
-struct ProtocolMonitor {
-    inner: Box<dyn RemoteIdConvertProtocol>,
-    output: Arc<Mutex<Vec<String>>>,
-}
-
-#[async_trait::async_trait]
-impl RemoteIdConvertProtocol for ProtocolMonitor {
-    async fn resolve_names_to_relative_paths(
-        &self,
-        heads: Vec<VertexName>,
-        names: Vec<VertexName>,
-    ) -> Result<Vec<(protocol::AncestorPath, Vec<VertexName>)>> {
-        let msg = format!("resolve names: {:?}, heads: {:?}", &names, &heads);
-        self.output.lock().push(msg);
-        self.inner
-            .resolve_names_to_relative_paths(heads, names)
-            .await
-    }
-
-    async fn resolve_relative_paths_to_names(
-        &self,
-        paths: Vec<protocol::AncestorPath>,
-    ) -> Result<Vec<(protocol::AncestorPath, Vec<VertexName>)>> {
-        let msg = format!("resolve paths: {:?}", &paths);
-        self.output.lock().push(msg);
-        self.inner.resolve_relative_paths_to_names(paths).await
-    }
-}
 
 #[tokio::test]
 async fn test_sparse_dag() {
@@ -73,9 +37,7 @@ async fn test_sparse_dag() {
     server2.drawdag("A-B-C G-C", &["C"]);
     server2.drawdag("C-D-E-M-X J-E", &["M"]);
 
-    let data = server2.dag.export_clone_data().await.unwrap();
-    let mut client = TestDag::new();
-    client.dag.import_clone_data(data).await.unwrap();
+    let client = server2.client().await;
 
     // Note: some ids (ex. 11) does not have matching name in its IdMap.
     // The server-side non-master (X) is not cloned.
@@ -98,17 +60,7 @@ async fn test_sparse_dag() {
 "#
     );
 
-    // Without remote protocol. Cannot solve id <-> names.
-    assert!(client.dag.vertex_name(Id(9)).await.is_err());
-    assert!(client.dag.vertex_id("A".into()).await.is_err());
-
     // With remote protocol. Be able to resolve id <-> names.
-    let output = Arc::new(Default::default());
-    let protocol = ProtocolMonitor {
-        inner: Box::new(server1.dag),
-        output: Arc::clone(&output),
-    };
-    client.dag.set_remote_protocol(Arc::new(protocol));
     assert_eq!(client.dag.vertex_name(Id(9)).await.unwrap(), "C".into());
     assert_eq!(client.dag.vertex_id("E".into()).await.unwrap(), Id(11));
 
@@ -124,7 +76,7 @@ async fn test_sparse_dag() {
     );
 
     assert_eq!(
-        output.lock().clone(),
+        client.output(),
         [
             "resolve paths: [D~1]",
             "resolve names: [E], heads: [M]",