diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 43dacaf1ea0..232952a3a33 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -1,17 +1,23 @@ name: CI # Controls when the workflow will run. +# +# The only live job here (integration-tests) builds the trace_machina/nativelink +# OCI image from the docker-compose Dockerfile and runs run_integration_tests.sh. +# That is a per-PR OCI build — the same class of heavy work #image-off-pr moved +# off PRs for image.yaml and #2877 moved off PRs for the nix integration job. +# This lane fell in the crack between those two, so defer it the same way: +# nightly cross-check + on every merge to main, not on every PR. Per-PR +# build/test signal is carried by the Nix "Bazel Dev" lane. on: push: branches: [ main ] paths-ignore: - '.github/styles/**' - 'web/**' - pull_request: - branches: [ main ] - paths-ignore: - - '.github/styles/**' - - 'web/**' + schedule: + - cron: '0 7 * * *' # 07:00 UTC daily (off-peak) + workflow_dispatch: permissions: read-all diff --git a/nativelink-scheduler/src/cache_lookup_scheduler.rs b/nativelink-scheduler/src/cache_lookup_scheduler.rs index 3d6ced82639..c8600b77d2b 100644 --- a/nativelink-scheduler/src/cache_lookup_scheduler.rs +++ b/nativelink-scheduler/src/cache_lookup_scheduler.rs @@ -231,8 +231,15 @@ impl CacheLookupScheduler { let inflight_cache_checks = self.inflight_cache_checks.clone(); // We need this spawn because we are returning a stream and this spawn will populate the stream's data. background_spawn!("cache_lookup_scheduler_add_action", async move { - // If our spawn ever dies, we will remove the action from the inflight_cache_checks map. - let _scope_guard = scope_guard; + // If our spawn ever dies before we take ownership of the map entry, remove the + // action from the inflight_cache_checks map (drop-safety). We defuse this guard + // the moment the cache check returns (below): from that point we remove the entry + // ourselves on every path with no intervening `.await`, so the guard's stale + // remove must NOT fire afterwards — a concurrent same-key request may have + // re-inserted a fresh entry in the window between our remove and this spawn's + // exit, and letting the guard run would orphan it (its oneshot tx is dropped, + // surfacing a spurious "tx hung up" to that client and forcing a retry). + let scope_guard = scope_guard; let unique_key = match &action_info.unique_qualifier { ActionUniqueQualifier::Cacheable(unique_key) => unique_key, @@ -254,6 +261,11 @@ impl CacheLookupScheduler { action_info.unique_qualifier.digest_function(), ) .await; + // The cache check has returned; from here every path removes our own map entry + // with no intervening `.await` (no drop point), so defusing the drop-guard now is + // leak-free and prevents its stale remove from clobbering a concurrent same-key + // re-insertion after we remove our entry. + scopeguard::ScopeGuard::into_inner(scope_guard); match maybe_action_result { Ok(action_result) => { let maybe_pending_txs = { diff --git a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs index cbff20b9867..de24db6d328 100644 --- a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs +++ b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs @@ -32,12 +32,14 @@ use nativelink_scheduler::cache_lookup_scheduler::CacheLookupScheduler; use nativelink_scheduler::mock_scheduler::MockActionScheduler; use nativelink_store::memory_store::MemoryStore; use nativelink_util::action_messages::{ - ActionResult, ActionStage, ActionState, ActionUniqueQualifier, OperationId, + ActionInfo, ActionResult, ActionStage, ActionState, ActionUniqueQualifier, OperationId, }; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; -use nativelink_util::operation_state_manager::{ClientStateManager, OperationFilter}; +use nativelink_util::operation_state_manager::{ + ActionStateResult, ClientStateManager, OperationFilter, +}; use nativelink_util::store_trait::{ RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; @@ -253,3 +255,95 @@ async fn find_by_client_operation_id_call_passed() -> Result<(), Error> { ); Ok(()) } + +fn queued_result( + client_operation_id: OperationId, + action_info: Arc, +) -> Box { + let (_tx, rx) = watch::channel(Arc::new(ActionState { + client_operation_id: client_operation_id.clone(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + last_transition_timestamp: SystemTime::now(), + })); + Box::new(TokioWatchActionStateResult::new( + client_operation_id, + action_info, + rx, + )) +} + +// A leader's cache-check task removes its in-flight entry before forwarding a +// miss to the inner scheduler, and that forward is an await point. A second +// request for the same key arriving in that window inserts a fresh entry. The +// leader's drop guard must not remove that fresh entry when the leader's task +// exits, or the second request's sender is dropped and it fails with +// "tx hung up". The default test runtime is single-threaded, so the +// interleaving below is deterministic. +#[nativelink_test] +async fn leader_drop_guard_does_not_orphan_a_later_same_key_request() -> Result<(), Error> { + let context = make_cache_scheduler()?; + let cache_scheduler = Arc::new(context.cache_scheduler); + let mock_scheduler = context.mock_scheduler; + // Nothing in the AC, so both requests miss and are forwarded. + let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 77)); + + // Request A becomes the leader, misses, removes its entry and parks in the + // forward to the mock scheduler, which has not answered yet. + let op_a = OperationId::default(); + let request_a = tokio::spawn({ + let cache_scheduler = cache_scheduler.clone(); + let (op_a, action_info) = (op_a.clone(), action_info.clone()); + async move { + cache_scheduler + .add_action(op_a, action_info) + .await + .map(|_| ()) + } + }); + for _ in 0..50 { + tokio::task::yield_now().await; + } + + // Request B for the same key arrives while A is parked: A's entry is gone, + // so B inserts its own. + let op_b = OperationId::default(); + let request_b = tokio::spawn({ + let cache_scheduler = cache_scheduler.clone(); + let (op_b, action_info) = (op_b.clone(), action_info.clone()); + async move { + cache_scheduler + .add_action(op_b, action_info) + .await + .map(|_| ()) + } + }); + + // Answer A's forward. A's task then finishes while B's entry is in the map. + let (forwarded_a, _) = mock_scheduler + .expect_add_action(Ok(queued_result(op_a.clone(), action_info.clone()))) + .await; + assert_eq!(forwarded_a, op_a); + + // B must still be forwarded and answered, not orphaned. + let serve_b = tokio::spawn({ + let mock_scheduler = mock_scheduler.clone(); + let action_info = action_info.clone(); + async move { + mock_scheduler + .expect_add_action(Ok(queued_result(op_b, action_info))) + .await + } + }); + let result_b = tokio::time::timeout(std::time::Duration::from_secs(5), request_b) + .await + .expect("request B never completed") + .expect("request B panicked"); + assert!( + result_b.is_ok(), + "a same-key request that arrived while the leader was forwarding was orphaned: {result_b:?}" + ); + serve_b.await.expect("serving B panicked"); + request_a.await.expect("request A panicked")?; + Ok(()) +}