Concurrent pantsd: Swap uncacheability for run-scoped results

src/rust/graph/src/entry.rs

@@ -58,7 +58,7 @@ impl Generation {
 #[derive(Debug)]
 pub(crate) enum NodeInterrupt<N: Node> {
     Dirtied,
-    Aborted(NodeResult<N>),
+    Aborted(N::Error),
 }
 
 ///
@@ -105,10 +105,12 @@ impl<N: Node> EntryResult<N> {
         }
     }
 
-    fn has_uncacheable_deps(&self) -> bool {
+    /// The Run which produced this result, if it may only be consumed within that Run.
+    fn run_scope(&self) -> Option<RunId> {
         match self {
-            EntryResult::Uncacheable(_, _) | EntryResult::UncacheableDependencies(_, _) => true,
-            EntryResult::Clean(..) | EntryResult::Dirty(..) => false,
+            EntryResult::Uncacheable(_, run_id)
+            | EntryResult::UncacheableDependencies(_, run_id) => Some(*run_id),
+            EntryResult::Clean(..) | EntryResult::Dirty(..) => None,
         }
     }
 
@@ -165,11 +167,31 @@ impl<N: Node> AsRef<N::Item> for EntryResult<N> {
     }
 }
 
-pub type NodeResult<N> = (
-    Result<<N as Node>::Item, <N as Node>::Error>,
-    Generation,
-    bool,
-);
+#[derive(Clone, Debug)]
+pub(crate) struct NodeResult<N: Node> {
+    pub(crate) result: Result<N::Item, N::Error>,
+    pub(crate) generation: Generation,
+    /// The Run which produced this result, if it may only be consumed within that Run.
+    pub(crate) run_scope: Option<RunId>,
+}
+
+impl<N: Node> NodeResult<N> {
+    fn completed(result: &EntryResult<N>, generation: Generation) -> Self {
+        NodeResult {
+            result: Ok(result.as_ref().clone()),
+            generation,
+            run_scope: result.run_scope(),
+        }
+    }
+
+    fn invalidated(generation: Generation) -> Self {
+        NodeResult {
+            result: Err(N::Error::invalidated()),
+            generation: generation.next(),
+            run_scope: None,
+        }
+    }
+}
 
 #[derive(Debug)]
 pub(crate) enum EntryState<N: Node> {
@@ -466,9 +488,9 @@ impl<N: Node> Entry<N> {
         tokio::select! {
           interrupt_item = sender.interrupted() => {
             match interrupt_item {
-              Some(NodeInterrupt::Aborted(res)) => {
-                  // We were aborted via terminate: complete with the given res.
-                  break (Some(res.0), context2.complete(entry2.node()))
+              Some(NodeInterrupt::Aborted(err)) => {
+                  // We were aborted via terminate: complete with the given error.
+                  break (Some(Err(err)), context2.complete(entry2.node()))
               }
               Some(NodeInterrupt::Dirtied) => {
                   // Attempt to clean the Node, and cancel it if we fail.
@@ -550,10 +572,22 @@ impl<N: Node> Entry<N> {
                 ..
             } => {
                 if let Some(receiver) = pending_value.receiver() {
+                    let entry = self.clone();
+                    let context = context.clone();
                     return async move {
-                        receiver.recv().await.unwrap_or_else(|| {
-                            (Err(N::Error::invalidated()), generation.next(), true)
-                        })
+                        let node_result = receiver
+                            .recv()
+                            .await
+                            .unwrap_or_else(|| NodeResult::invalidated(generation));
+                        match node_result.run_scope {
+                            Some(run_id)
+                                if node_result.result.is_ok() && run_id != context.run_id() =>
+                            {
+                                // We attached to another Run's attempt: re-enter for our own value.
+                                entry.get_node_result(&context, entry_id).await
+                            }
+                            _ => node_result,
+                        }
                     }
                     .boxed();
                 }
@@ -564,12 +598,7 @@ impl<N: Node> Entry<N> {
                 generation,
                 ..
             } if result.is_clean(context) => {
-                return future::ready((
-                    Ok(result.as_ref().clone()),
-                    generation,
-                    result.has_uncacheable_deps(),
-                ))
-                .boxed();
+                return future::ready(NodeResult::completed(result, generation)).boxed();
             }
             _ => (),
         };
@@ -644,7 +673,7 @@ impl<N: Node> Entry<N> {
             receiver
                 .recv()
                 .await
-                .unwrap_or_else(|| (Err(N::Error::invalidated()), generation.next(), true))
+                .unwrap_or_else(|| NodeResult::invalidated(generation))
         }
         .boxed()
     }
@@ -732,7 +761,11 @@ impl<N: Node> Entry<N> {
                     previous_result.dirty();
                 }
                 generation = generation.next();
-                sender.send((Err(e), generation, true));
+                sender.send(NodeResult {
+                    result: Err(e),
+                    generation,
+                    run_scope: Some(context.run_id()),
+                });
                 EntryState::NotStarted {
                     run_token,
                     generation,
@@ -748,11 +781,7 @@ impl<N: Node> Entry<N> {
                     // Node was re-executed (ie not cleaned) and had a different result value.
                     generation = generation.next()
                 };
-                sender.send((
-                    Ok(next_result.as_ref().clone()),
-                    generation,
-                    next_result.has_uncacheable_deps(),
-                ));
+                sender.send(NodeResult::completed(&next_result, generation));
                 EntryState::Completed {
                     result: next_result,
                     pollers: Vec::new(),
@@ -771,11 +800,7 @@ impl<N: Node> Entry<N> {
                     self.cacheable_with_output(Some(result.as_ref())),
                     has_uncacheable_deps,
                 );
-                sender.send((
-                    Ok(result.as_ref().clone()),
-                    generation,
-                    result.has_uncacheable_deps(),
-                ));
+                sender.send(NodeResult::completed(&result, generation));
                 EntryState::Completed {
                     result,
                     pollers: Vec::new(),
@@ -928,17 +953,8 @@ impl<N: Node> Entry<N> {
         let inner = &mut *self.state.lock();
         test_trace_log!("Terminating node {:?} with {:?}", self.node, err);
         for state in inner.states_mut() {
-            if let EntryState::Running {
-                pending_value,
-                generation,
-                ..
-            } = state
-            {
-                let _ = pending_value.try_interrupt(NodeInterrupt::Aborted((
-                    Err(err.clone()),
-                    generation.next(),
-                    true,
-                )));
+            if let EntryState::Running { pending_value, .. } = state {
+                let _ = pending_value.try_interrupt(NodeInterrupt::Aborted(err.clone()));
             };
         }
     }

src/rust/graph/src/lib.rs

@@ -5,7 +5,7 @@ mod context;
 mod entry;
 mod node;
 
-use crate::entry::{Entry, Generation, RunToken};
+use crate::entry::{Entry, Generation, NodeResult, RunToken};
 
 use std::collections::VecDeque;
 use std::fs::File;
@@ -528,9 +528,11 @@ impl<N: Node> Graph<N> {
 
         // Return the state of the destination, retrying the dst to handle Node invalidation.
         let context = context.clone();
-        let (result, generation, uncacheable) = loop {
+        let node_result = loop {
             match entry.get_node_result(&context, entry_id).await {
-                (Err(err), _, _) if err == N::Error::invalidated() => {
+                NodeResult {
+                    result: Err(err), ..
+                } if err == N::Error::invalidated() => {
                     let node = {
                         let inner = self.inner.lock();
                         inner.unsafe_entry_for_id(entry_id).node().clone()
@@ -547,12 +549,16 @@ impl<N: Node> Graph<N> {
         };
 
         if src_id.is_some()
-            && let Err(e) = context.dep_record(entry_id, generation, uncacheable)
+            && let Err(e) = context.dep_record(
+                entry_id,
+                node_result.generation,
+                node_result.run_scope.is_some(),
+            )
         {
-            return (Err(e), generation);
+            return (Err(e), node_result.generation);
         }
 
-        (result, generation)
+        (node_result.result, node_result.generation)
     }
 
     ///
@@ -625,11 +631,10 @@ impl<N: Node> Graph<N> {
                         .clone();
 
                     async move {
-                        let (_, generation, uncacheable) =
-                            dep_entry.get_node_result(context, dep_id).await;
-                        if generation == previous_dep_generation {
+                        let node_result = dep_entry.get_node_result(context, dep_id).await;
+                        if node_result.generation == previous_dep_generation {
                             // Matched.
-                            Ok(uncacheable)
+                            Ok(node_result.run_scope.is_some())
                         } else {
                             // Did not match. We error here to trigger fail-fast in `try_join_all`.
                             log::debug!(

src/rust/graph/src/tests.rs

@@ -12,6 +12,7 @@ use futures::future;
 use parking_lot::Mutex;
 use rand::RngExt;
 use task_executor::Executor;
+use tokio::sync::oneshot;
 use tokio::time::{error::Elapsed, sleep, timeout};
 
 use crate::context::Context;
@@ -464,6 +465,54 @@ async fn uncacheable_deps_is_cleaned_for_the_session() {
     assert_no_change_within_session(&context);
 }
 
+#[tokio::test]
+async fn uncacheable_node_reruns_for_concurrent_run() {
+    let _logger = env_logger::try_init();
+    let graph = empty_graph();
+
+    let mut uncacheable = HashSet::new();
+    uncacheable.insert(TNode::new(1));
+
+    // Run A blocks in the uncacheable node so run B observes its attempt in flight.
+    let (gate, entered, release) = TGate::new();
+    let context_a = {
+        let mut gates = HashMap::new();
+        gates.insert(TNode::new(1), gate);
+        graph.context(
+            TContext::new()
+                .with_uncacheable(uncacheable.clone())
+                .with_gates(gates),
+        )
+    };
+    let context_b = graph.context(TContext::new().with_uncacheable(uncacheable));
+
+    let graph_a = graph.clone();
+    let context_a2 = context_a.clone();
+    let fut_a = tokio::spawn(async move { graph_a.create(TNode::new(2), &context_a2).await });
+    entered.await.unwrap();
+
+    let graph_b = graph.clone();
+    let context_b2 = context_b.clone();
+    let fut_b = tokio::spawn(async move { graph_b.create(TNode::new(2), &context_b2).await });
+    for _ in 0..10 {
+        tokio::task::yield_now().await;
+    }
+    // Yields are not synchronization: assert B is still pending, i.e. attached to A.
+    assert!(!fut_b.is_finished());
+    release.send(()).unwrap();
+
+    let expected = vec![T(0, 0), T(1, 0), T(2, 0)];
+    assert_eq!(fut_a.await.unwrap(), Ok(expected.clone()));
+    assert_eq!(fut_b.await.unwrap(), Ok(expected));
+
+    // B attached to A's attempt, but re-ran the uncacheable node rather than consuming it.
+    assert_eq!(
+        context_a.runs(),
+        vec![TNode::new(2), TNode::new(1), TNode::new(0)]
+    );
+    assert_eq!(context_b.runs(), vec![TNode::new(1)]);
+}
+
 #[tokio::test]
 async fn live_runs_released_when_last_holder_drops() {
     let _logger = env_logger::try_init();
@@ -944,6 +993,7 @@ struct TContext {
     edges: Arc<HashMap<TNode, Vec<TNode>>>,
     delays_pre: Arc<HashMap<TNode, Duration>>,
     delays_post: Arc<HashMap<TNode, Duration>>,
+    gates: Arc<HashMap<TNode, TGate>>,
     // Nodes which should error when they run.
     errors: Arc<HashSet<TNode>>,
     non_restartable: Arc<HashSet<TNode>>,
@@ -952,13 +1002,49 @@ struct TContext {
     runs: Arc<Mutex<Vec<TNode>>>,
 }
 
+///
+/// A gate for a TNode: the first time the node runs, it signals `entered` and then blocks until
+/// `release` fires.
+///
+struct TGate {
+    entered: Mutex<Option<oneshot::Sender<()>>>,
+    release: Mutex<Option<oneshot::Receiver<()>>>,
+}
+
+impl TGate {
+    /// Returns a TGate, a handle to await its entry, and a handle to release it.
+    fn new() -> (TGate, oneshot::Receiver<()>, oneshot::Sender<()>) {
+        let (entered_tx, entered_rx) = oneshot::channel();
+        let (release_tx, release_rx) = oneshot::channel();
+        (
+            TGate {
+                entered: Mutex::new(Some(entered_tx)),
+                release: Mutex::new(Some(release_rx)),
+            },
+            entered_rx,
+            release_tx,
+        )
+    }
+
+    async fn pass(&self) {
+        if let Some(entered) = self.entered.lock().take() {
+            let _ = entered.send(());
+        }
+        let release = self.release.lock().take();
+        if let Some(release) = release {
+            release.await.unwrap();
+        }
+    }
+}
+
 impl TContext {
     fn new() -> TContext {
         TContext {
             salt: Arc::new(AtomicUsize::new(0)),
             edges: Arc::default(),
             delays_pre: Arc::default(),
             delays_post: Arc::default(),
+            gates: Arc::default(),
             errors: Arc::default(),
             non_restartable: Arc::default(),
             uncacheable: Arc::default(),
@@ -984,6 +1070,12 @@ impl TContext {
         self
     }
 
+    /// Gates awaited before a node requests its dependencies.
+    fn with_gates(mut self, gates: HashMap<TNode, TGate>) -> TContext {
+        self.gates = Arc::new(gates);
+        self
+    }
+
     fn with_errors(mut self, errors: HashSet<TNode>) -> TContext {
         self.errors = Arc::new(errors);
         self
@@ -1030,6 +1122,9 @@ impl TContext {
     }
 
     async fn maybe_delay_pre(&self, node: &TNode) {
+        if let Some(gate) = self.gates.get(node) {
+            gate.pass().await;
+        }
         if let Some(delay) = self.delays_pre.get(node) {
             sleep(*delay).await;
         }