own tree shape

3 files changed, 187 insertions, 10 deletions

diff --git a/src/api/stats.rs b/src/api/stats.rs
index 5061ff7..3851ef7 100644
--- a/src/api/stats.rs
+++ b/src/api/stats.rs
@@ -187,6 +187,18 @@ struct TopologyView {
     stable: usize,
     neighbours: usize,
     stable_neighbours: usize,
+    /// Our own writer's tree (`swarm::tree::own_shape`); absent while we aren't in the
+    /// member list's graph.
+    own_tree: Option<OwnTree>,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+struct OwnTree {
+    members: usize,
+    depth: u32,
+    mean_depth: f64,
+    off_tree: usize,
 }
 
 #[derive(Debug, Serialize)]
@@ -264,6 +276,18 @@ pub(crate) async fn get_stats(State(state): State<AppState>, headers: HeaderMap)
 /// totals over the peers from their list (`peer_list`), read once for both.
 pub(crate) async fn collect(state: &AppState, list: &PeerListResponse) -> StatsResponse {
     let node = state.stats.snapshot();
+    let own_tree = crate::swarm::tree::own_shape(state)
+        .await
+        .map(|shape| OwnTree {
+            members: shape.members,
+            depth: shape.depth,
+            mean_depth: if shape.members == 0 {
+                0.0
+            } else {
+                shape.hops as f64 / shape.members as f64
+            },
+            off_tree: shape.off,
+        });
 
     let (heads, tier, remakes) = {
         let replica = state.doc.read().await;
@@ -358,6 +382,7 @@ pub(crate) async fn collect(state: &AppState, list: &PeerListResponse) -> StatsR
                 stable: s.stable,
                 neighbours: s.neighbours,
                 stable_neighbours: s.stable_neighbours,
+                own_tree,
             }),
         },
         stability: StabilityStats {
diff --git a/src/api/stats/prometheus.rs b/src/api/stats/prometheus.rs
index 595466a..db3cef3 100644
--- a/src/api/stats/prometheus.rs
+++ b/src/api/stats/prometheus.rs
@@ -713,6 +713,32 @@ pub(super) fn render(stats: &StatsResponse, list: &PeerListResponse) -> String {
             "Current neighbours with the stable tier.",
             view.stable_neighbours as f64,
         );
+        if let Some(tree) = &view.own_tree {
+            out.single(
+                "own_tree_members",
+                "gauge",
+                "Other members in this node's own writer's tree.",
+                tree.members as f64,
+            );
+            out.single(
+                "own_tree_depth",
+                "gauge",
+                "Most hops from this node to a member along its own writer's tree.",
+                f64::from(tree.depth),
+            );
+            out.single(
+                "own_tree_mean_depth",
+                "gauge",
+                "Mean hops from this node to a member along its own writer's tree.",
+                tree.mean_depth,
+            );
+            out.single(
+                "own_tree_off_members",
+                "gauge",
+                "Members this node's own writer's tree doesn't reach (they get its writes by announcement).",
+                tree.off_tree as f64,
+            );
+        }
     }
 
     let stability = &stats.stability;
diff --git a/src/swarm/tree.rs b/src/swarm/tree.rs
index c6baff9..015d62c 100644
--- a/src/swarm/tree.rs
+++ b/src/swarm/tree.rs
@@ -174,6 +174,105 @@ pub(crate) fn whole(
     )
 }
 
+/// Whom we and our neighbours in `graph` reach, by pair key, as far as we know (our
+/// own, and what our neighbours said).
+async fn known_reaches<'g>(state: &AppState, graph: &'g Graph) -> HashMap<&'g str, Vec<u64>> {
+    let me = state.local_id();
+    let mut reaches: HashMap<&str, Vec<u64>> = HashMap::new();
+    if let Some((me, _)) = graph.get_key_value(me) {
+        reaches.insert(me.as_str(), reach::published(state));
+    }
+    for member in graph.get(me).into_iter().flatten() {
+        if let Some(peer_state) = state.get_peer_state(member).await
+            && let Some(keys) = reach::of(peer_state.lock().await.reaches.as_ref())
+        {
+            reaches.insert(member.as_str(), keys);
+        }
+    }
+    reaches
+}
+
+/// Our own tree, as our writes go along it (`Plan`, as every member computes its part).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct Shape {
+    /// The other members in it.
+    pub(crate) members: usize,
+    /// The most hops from us to one of them.
+    pub(crate) depth: u32,
+    /// The hops to each, summed (for the mean).
+    pub(crate) hops: u64,
+    /// Members it doesn't reach: without a parent there (no member they reach comes
+    /// before them), below one without, or not connected to us in the member list's
+    /// graph. They get our writes by announcement.
+    pub(crate) off: usize,
+}
+
+/// The shape of our own tree; None while the member list isn't known, or we aren't in
+/// it. What members other than us and our neighbours reach is taken from the graph, as
+/// in `plan`.
+pub(crate) async fn own_shape(state: &AppState) -> Option<Shape> {
+    let me = state.local_id();
+    let (graph, distance) = distances(state, topology::pair_key(me))?;
+    let reaches = known_reaches(state, &graph).await;
+    let reached = |by: &str, member: &str| {
+        reaches
+            .get(by)
+            .is_none_or(|keys| keys.contains(&topology::pair_key(member)))
+    };
+    Some(shape(&graph, &distance, &reached, me))
+}
+
+/// `own_shape` over `graph`, from `me` at distance 0.
+fn shape(
+    graph: &Graph,
+    distance: &HashMap<String, u32>,
+    reached: &dyn Fn(&str, &str) -> bool,
+    me: &str,
+) -> Shape {
+    let parents: HashMap<&str, Option<&str>> = graph
+        .keys()
+        .filter(|member| distance.get(member.as_str()).is_some_and(|d| *d > 0))
+        .map(|member| (member.as_str(), parent(graph, distance, reached, member)))
+        .collect();
+    // Hops from us along parents; None where the chain ends without us. Parents come
+    // before their children, so chains end.
+    let mut hops: HashMap<&str, Option<u32>> = HashMap::from([(me, Some(0))]);
+    fn along<'g>(
+        member: &'g str,
+        parents: &HashMap<&'g str, Option<&'g str>>,
+        hops: &mut HashMap<&'g str, Option<u32>>,
+    ) -> Option<u32> {
+        if let Some(known) = hops.get(member) {
+            return *known;
+        }
+        let found = parents
+            .get(member)
+            .copied()
+            .flatten()
+            .and_then(|parent| along(parent, parents, hops))
+            .map(|hops| hops + 1);
+        hops.insert(member, found);
+        found
+    }
+    let mut shape = Shape {
+        members: 0,
+        depth: 0,
+        hops: 0,
+        off: graph.len() - 1 - parents.len(),
+    };
+    for member in parents.keys() {
+        match along(member, &parents, &mut hops) {
+            Some(h) => {
+                shape.members += 1;
+                shape.depth = shape.depth.max(h);
+                shape.hops += u64::from(h);
+            }
+            None => shape.off += 1,
+        }
+    }
+    shape
+}
+
 /// What `writer`'s tree means for us: our children in it, and our neighbours it leaves
 /// out (`uncovered`).
 pub(crate) struct Plan {
@@ -195,16 +294,7 @@ pub(crate) async fn plan(state: &AppState, writer: u64) -> Option<Plan> {
     if !distance.contains_key(me) {
         return None;
     }
-    // Whom we and our neighbours reach, by pair key.
-    let mut reaches: HashMap<&str, Vec<u64>> = HashMap::new();
-    reaches.insert(me, reach::published(state));
-    for member in graph.get(me).into_iter().flatten() {
-        if let Some(peer_state) = state.get_peer_state(member).await
-            && let Some(keys) = reach::of(peer_state.lock().await.reaches.as_ref())
-        {
-            reaches.insert(member.as_str(), keys);
-        }
-    }
+    let reaches = known_reaches(state, &graph).await;
     let reached = |by: &str, member: &str| {
         reaches
             .get(by)
@@ -277,3 +367,39 @@ pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Option
     }
     Some(plan)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Hops along parents from us; members without a parent, below one without, or
+    /// not connected count as off the tree.
+    #[test]
+    fn our_tree_in_numbers() {
+        let link = |graph: &mut Graph, a: &str, b: &str| {
+            graph.entry(a.into()).or_default().insert(b.into());
+            graph.entry(b.into()).or_default().insert(a.into());
+        };
+        let mut graph = Graph::new();
+        link(&mut graph, "me", "a");
+        link(&mut graph, "a", "b");
+        link(&mut graph, "me", "c");
+        link(&mut graph, "c", "e");
+        graph.insert("d".into(), BTreeSet::new());
+        let distance: HashMap<String, u32> = [("me", 0), ("a", 1), ("c", 1), ("b", 2), ("e", 2)]
+            .into_iter()
+            .map(|(member, d)| (member.to_owned(), d))
+            .collect();
+        // c doesn't reach us: no parent, nor e below it.
+        let reached = |by: &str, member: &str| !(by == "c" && member == "me");
+        assert_eq!(
+            shape(&graph, &distance, &reached, "me"),
+            Shape {
+                members: 2,
+                depth: 2,
+                hops: 3,
+                off: 3,
+            }
+        );
+    }
+}