migration: Send IsReady to all upstream full mat nodes

On restart, we have a race condition that causes a panic when
"recovering" the `DfState`. This occurs on joins with fully
materialized nodes, when replay begins. Before replay, all nodes will
be sent a `Ready` message, but base nodes (which must open files from
disk), "ready" themselves asynchronously - meaning, the Leader does
not block for them to become ready.

The problem we are encountering is when a replay that contains a join
node begins, the `DomainMigrationPlan` ensures only the left-side is
ready (by sending an `IsReady` message). The right-side has been sent
a `Ready` message, but we have not guaranteed that it is, indeed,
"ready".

This CL checks each ReplayPath, in the `DomainMigrationPlan`'s `Plan`
phase, and if the path contains fully-materialized nodes, it checks to
see if there's a join node. If there is a join, it walks up the DAG to
look for both base nodes and fully materialized nodes. The indices of
those nodes are added to a set, and a further part of the
`DomainMigrationPlan::commit()` will send an `IsReady` message to
those additional nodes.

This will not cause any additional traffic at startup as there's
already a `HashSet` in `DomainMigrationPlan::commit()` which tracks if
a `IsReady` has already been sent to a target node.

Fixes: REA-5126
Change-Id: Icbc815ce0a4aa4ed1b03aa734c785122c0839cbb
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/8563
Tested-by: Buildkite CI
Reviewed-by: Michael Zink <[email protected]>

Author: jasobrown-rs

readyset-server/src/controller/migrate/materialization/mod.rs

@@ -1325,6 +1325,11 @@ impl Materializations {
                         },
                     )?;
                 }
+                for (domain, node) in pending.additional_ancestors {
+                    if non_ready_nodes.remove(&(domain, node)) {
+                        dmp.add_message(domain, DomainRequest::IsReady { node })?;
+                    }
+                }
 
                 dmp.add_message(
                     pending.source_domain,

readyset-server/src/controller/migrate/materialization/plan.rs

@@ -68,6 +68,12 @@ pub(super) struct PendingReplay {
     pub(super) source: LocalNodeIndex,
     pub(super) source_domain: DomainIndex,
     pub(super) target_domain: DomainIndex,
+
+    // If the segment contained a join, we need to ensure any fully materialized
+    // nodes in all ancestors are `Ready` before we start the replay. Else on restart, we have a
+    // race condition between the replay hitting the join node and doing a lookup on the other
+    // parent, and the other parent actually having it's base tables opened.
+    pub(super) additional_ancestors: HashSet<(DomainIndex, LocalNodeIndex)>,
 }
 
 impl<'a> Plan<'a> {
@@ -286,6 +292,45 @@ impl<'a> Plan<'a> {
             .collect()
     }
 
+    /// Find the ancestors of any joins for the given path. This is to ensure that all parents
+    /// are `Ready` (and will have been sent an `IsReady` message) before we start a replay.
+    ///
+    /// We assume that since the node traversal is topologically sorted (see
+    /// `Materializations::commit()`), any parent node would have been sent a `Ready` message.
+    ///
+    /// To do this, we look for all join nodes in the provided `path`, and walk the tree of
+    /// the each parent. We then find all fully materialized nodes in all ancestors.
+    fn find_additional_ancestors(
+        &self,
+        path: &RawReplayPath,
+    ) -> HashSet<(DomainIndex, LocalNodeIndex)> {
+        let mut additional_ancestors = HashSet::new();
+
+        let mut seen = HashSet::new();
+        for segment in path.segments() {
+            let n = &self.graph[segment.node];
+            if n.is_join().is_ok_and(|b| b) {
+                // as the `path.segments()` iterator traverses downward through the graph,
+                // we can skip ancestors of the current segment we've already looked at.
+                let mut queue = n.ancestors().unwrap_or_default();
+                queue.retain(|a| !seen.contains(a));
+
+                while let Some(parent) = queue.pop() {
+                    let node = &self.graph[parent];
+                    if node.is_base() || node.requires_full_materialization() {
+                        additional_ancestors.insert((node.domain(), node.local_addr()));
+                    } else {
+                        queue.extend(node.ancestors().unwrap_or_default());
+                    }
+
+                    seen.insert(segment.node);
+                }
+            }
+        }
+
+        additional_ancestors
+    }
+
     /// Finds the appropriate replay paths for the given index, and inform all domains on those
     /// paths about them. It also notes if any data backfills will need to be run, which is
     /// eventually reported back by `finalize`.
@@ -718,6 +763,10 @@ impl<'a> Plan<'a> {
                         }
                     }
                 } else {
+                    // for full materializations, we need to identifiy the parents that
+                    // are also fully materialized, and make sure to wait for them to be `Ready`.
+                    let additional_ancestors = self.find_additional_ancestors(&path);
+
                     // for full materializations, the last domain should report when it's done
                     if i == segments.len() - 1 {
                         if let DomainRequest::SetupReplayPath {
@@ -742,6 +791,7 @@ impl<'a> Plan<'a> {
                                         .last()
                                         .ok_or_else(|| internal_err!())?
                                         .0,
+                                    additional_ancestors,
                                 });
                             }
                         }