From 670f88b530fb6ed98b053a3ecd72a102cb40ee4c Mon Sep 17 00:00:00 2001
From: Niko Matsakis <nikomat@amazon.com>
Date: Fri, 15 Nov 2024 16:58:58 +0000
Subject: [PATCH] another parallel test scenario

---
 tests/parallel/cycle_a_t1_b_t2.rs | 79 +++++++++++++++++++++++++++++++
 tests/parallel/main.rs            |  1 +
 tests/parallel/setup.rs           |  2 +
 3 files changed, 82 insertions(+)
 create mode 100644 tests/parallel/cycle_a_t1_b_t2.rs

diff --git a/tests/parallel/cycle_a_t1_b_t2.rs b/tests/parallel/cycle_a_t1_b_t2.rs
new file mode 100644
index 000000000..5ea5c447d
--- /dev/null
+++ b/tests/parallel/cycle_a_t1_b_t2.rs
@@ -0,0 +1,79 @@
+//! Test a specific cycle scenario:
+//!
+//! ```text
+//! Thread T1          Thread T2
+//! ---------          ---------
+//!    |                  |
+//!    v                  |
+//! query_a()             |
+//!  ^  |                 v
+//!  |  +------------> query_b()
+//!  |                    |
+//!  +--------------------+
+//! ``````
+
+use salsa::CycleRecoveryAction;
+
+use crate::setup::{Knobs, KnobsDatabase};
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
+struct CycleValue(u32);
+
+const MIN: CycleValue = CycleValue(0);
+const MAX: CycleValue = CycleValue(22);
+
+// Signal 1: T1 has entered `query_a`
+// Signal 2: T2 has entered `query_b`
+
+#[salsa::tracked(cycle_fn=query_a_cycle_fn, cycle_initial=query_a_initial)]
+fn query_a(db: &dyn KnobsDatabase) -> CycleValue {
+    db.signal(1);
+
+    // Wait for Thread T2 to enter `query_b` before we continue.
+    db.wait_for(2);
+
+    query_b(db)
+}
+
+fn query_a_cycle_fn(
+    _db: &dyn KnobsDatabase,
+    _value: &CycleValue,
+    _count: u32,
+) -> CycleRecoveryAction<CycleValue> {
+    CycleRecoveryAction::Iterate
+}
+
+fn query_a_initial(_db: &dyn KnobsDatabase) -> CycleValue {
+    MIN
+}
+
+#[salsa::tracked]
+fn query_b(db: &dyn KnobsDatabase) -> CycleValue {
+    // Wait for Thread T1 to enter `query_a` before we continue.
+    db.wait_for(1);
+
+    db.signal(2);
+
+    let a_value = query_a(db);
+    CycleValue(a_value.0 + 1).min(MAX)
+}
+
+#[test]
+fn the_test() {
+    std::thread::scope(|scope| {
+        let db_t1 = Knobs::default();
+        let db_t2 = db_t1.clone();
+
+        // Thread 1:
+        scope.spawn(move || {
+            let r = query_a(&db_t1);
+            assert_eq!(r, MAX);
+        });
+
+        // Thread 2:
+        scope.spawn(move || {
+            let r = query_b(&db_t2);
+            assert_eq!(r, MAX);
+        });
+    });
+}
diff --git a/tests/parallel/main.rs b/tests/parallel/main.rs
index 6b34dc06f..e4e423937 100644
--- a/tests/parallel/main.rs
+++ b/tests/parallel/main.rs
@@ -1,5 +1,6 @@
 mod setup;
 
+mod cycle_a_t1_b_t2;
 mod cycle_ab_peeping_c;
 mod parallel_cancellation;
 mod parallel_map;
diff --git a/tests/parallel/setup.rs b/tests/parallel/setup.rs
index c266731a0..b67d4f667 100644
--- a/tests/parallel/setup.rs
+++ b/tests/parallel/setup.rs
@@ -9,8 +9,10 @@ use crate::signal::Signal;
 /// a certain behavior.
 #[salsa::db]
 pub(crate) trait KnobsDatabase: Database {
+    /// Signal that we are entering stage 1.
     fn signal(&self, stage: usize);
 
+    /// Wait until we reach stage `stage` (no-op if we have already reached that stage).
     fn wait_for(&self, stage: usize);
 }