Merge scheduler lease renewal for slow contract setup from the scheduler base

Author: kvlonge

crates/sqlbuild-orchestration/tests/scheduler/helpers.rs

@@ -4,15 +4,36 @@ pub(crate) use crate::scheduler::main::_acquire_lease::acquire_lease;
 pub(crate) use crate::scheduler::main::_evaluate_trigger::evaluate_trigger;
 pub(crate) use crate::scheduler::main::_scheduler_cycle::scheduler_cycle;
 use crate::scheduler::main::renew_lease::renew_lease;
+use crate::scheduler::models::SchedulerLease;
 use crate::store::errors::ControlError;
 use crate::store::main::migrate_store::migrate_store;
+use crate::store::models::ControlStore;
+use futures_util::{StreamExt, stream};
 use uuid::Uuid;
 
 use crate::store::tests::lease_helpers::expire_lease;
 use crate::store::tests::replica_helpers::second_store;
 use crate::store::tests::store_helpers::{close_store, close_test_store, open_test_store};
 use crate::store::tests::store_test_types::Backend;
 
+/// Renew `lease` every five seconds, as a running scheduler does, until a renewal fails.
+pub(crate) fn hold_lease(
+    store: &ControlStore,
+    lease: &SchedulerLease,
+) -> tokio::task::JoinHandle<()> {
+    let store = store.clone();
+    let renewals = stream::try_unfold(lease.clone(), move |lease| {
+        let store = store.clone();
+        async move {
+            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+            renew_lease(&store, &lease)
+                .await
+                .map(|renewed| Some(((), renewed)))
+        }
+    });
+    tokio::spawn(renewals.for_each(|_| async {}))
+}
+
 pub(crate) async fn leadership_contract(backend: Backend) -> Option<i64> {
     let test_store = open_test_store(backend).await?;
     let store = &test_store.store;

crates/sqlbuild-orchestration/tests/scheduler/ticks/test_retention.rs

@@ -14,6 +14,7 @@ use crate::scheduler::main::{
     _acquire_lease::acquire_lease, _evaluate_trigger::evaluate_trigger, list_ticks::list_ticks,
 };
 use crate::scheduler::models::TickOutcome;
+use crate::scheduler::tests::contracts::helpers::hold_lease;
 use crate::scheduler::tests::contracts::ticks::helpers::{
     append_covered_tick, calendar_time, shift_schedule, snapshot,
 };
@@ -67,6 +68,7 @@ fn given_expired_covered_ticks_when_pruned_and_cursor_reset_then_run_linked_and_
                 .await
                 .expect("election")
                 .expect("lease");
+            let holder = hold_lease(store, &lease);
             evaluate_trigger(store, &lease, trigger_id)
                 .await
                 .expect("catch up");
@@ -176,6 +178,7 @@ fn given_expired_covered_ticks_when_pruned_and_cursor_reset_then_run_linked_and_
                 Some("2030-01-01T00:00:00Z".parse().expect("future"))
             );
             assert!(ticks.iter().all(|tick| tick.scheduled_for < Utc::now()));
+            holder.abort();
             close_test_store(fixture).await;
         });
     }

crates/sqlbuild-orchestration/tests/sensors/test_nodes.rs

@@ -17,7 +17,7 @@ use crate::runs::main::{
 use crate::runs::models::{
     JobRunRequest, NodeOutcome, RunAuthority, RunChange, RunChangeRequest, RunOutcome,
 };
-use crate::scheduler::tests::contracts::helpers::acquire_lease;
+use crate::scheduler::tests::contracts::helpers::{acquire_lease, hold_lease};
 use crate::sensors::main::{
     begin_sensor_evaluation::begin_sensor_evaluation,
     complete_builtin_sensor::complete_builtin_sensor,
@@ -89,6 +89,7 @@ fn given_any_run_rebuilds_when_consuming_then_only_genuine_scoped_successes_queu
                 .await
                 .expect("elect")
                 .expect("leader");
+            let holder = hold_lease(store, &lease);
             let trigger = list_triggers(store, location, "development")
                 .await
                 .expect("triggers")
@@ -368,6 +369,7 @@ fn given_any_run_rebuilds_when_consuming_then_only_genuine_scoped_successes_queu
                     .count(),
                 test_case.expected_runs
             );
+            holder.abort();
             close_store(replica).await;
             close_test_store(fixture).await;
         });

crates/sqlbuild-orchestration/tests/sensors/test_runs.rs

@@ -13,7 +13,7 @@ use crate::declarations::models::TriggerChange;
 use crate::runs::main::list_runs::list_runs;
 use crate::runs::main::request_rerun::request_rerun;
 use crate::runs::models::{NodeOutcome, RerunRequest, RunOutcome};
-use crate::scheduler::tests::contracts::helpers::acquire_lease;
+use crate::scheduler::tests::contracts::helpers::{acquire_lease, hold_lease};
 use crate::sensors::main::{
     begin_sensor_evaluation::begin_sensor_evaluation,
     complete_builtin_sensor::complete_builtin_sensor,
@@ -109,6 +109,7 @@ fn given_terminal_jobs_and_groups_when_selectors_retry_then_outcomes_scope_and_r
                 .await
                 .expect("elect")
                 .expect("leader");
+            let holder = hold_lease(store, &lease);
             let mut expected: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
             let outcomes = [
                 (
@@ -504,6 +505,7 @@ fn given_terminal_jobs_and_groups_when_selectors_retry_then_outcomes_scope_and_r
                     "sensor never changes source state"
                 );
             }
+            holder.abort();
             close_store(replica).await;
             close_test_store(fixture).await;
         });

crates/sqlbuild-orchestration/tests/sensors/test_transitions.rs

@@ -10,7 +10,7 @@ use crate::declarations::main::list_triggers::list_triggers;
 use crate::declarations::models::TriggerChange;
 use crate::runs::main::list_runs::list_runs;
 use crate::runs::models::{NodeOutcome, RunOutcome};
-use crate::scheduler::tests::contracts::helpers::acquire_lease;
+use crate::scheduler::tests::contracts::helpers::{acquire_lease, hold_lease};
 use crate::sensors::main::{
     begin_sensor_evaluation::begin_sensor_evaluation,
     complete_builtin_sensor::complete_builtin_sensor,
@@ -82,6 +82,7 @@ fn given_independent_check_failures_when_nondecisive_results_arrive_then_only_ex
                 .await
                 .expect("elect")
                 .expect("leader");
+            let holder = hold_lease(store, &lease);
             let fixture_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
                 .join("../../benchmarks/duckshop/edge_cases/orchestration/check_transitions.json");
             let transitions: Vec<serde_json::Value> = serde_json::from_str(
@@ -178,6 +179,7 @@ fn given_independent_check_failures_when_nondecisive_results_arrive_then_only_ex
                     .count(),
                 test_case.expected_runs
             );
+            holder.abort();
             close_test_store(fixture).await;
         });
     }