Don't hand out an already-cancelled transaction from saved state (#4670) Summary: ### Overview The `test_shows_stdlib_errors_for_multiple_versions_and_paths_with_force_on` test showed some flakiness, where a `textDocument/diagnostic` response comes back `"items": []` for a just-opened file. This pointed to a real gremlin. ### Solution Resetting the cancellation on restore upholds the invariant that every transaction handed out can perform work, while keeping the cached work that makes saving worthwhile (this is a general form of the guard found in the `ProvideType` handler). Pull Request resolved: https://github.com/facebook/pyrefly/pull/4670 Test Plan: Existing tests pass; created a new/deterministic unit test (with the help of Opus 5) that covers the issue. Reviewed By: rchen152 Differential Revision: D117483183 Pulled By: NathanTempest fbshipit-source-id: 56007d2a13b8a6b3831902557e2b1508f1f10cdb
diff --git a/pyrefly/lib/lsp/non_wasm/transaction_manager.rs b/pyrefly/lib/lsp/non_wasm/transaction_manager.rs index 132eb99..2305618 100644 --- a/pyrefly/lib/lsp/non_wasm/transaction_manager.rs +++ b/pyrefly/lib/lsp/non_wasm/transaction_manager.rs
@@ -78,3 +78,88 @@ self.saved_state = Some(transaction.save(telemetry)) } } + +#[cfg(test)] +mod tests { + use std::time::Instant; + + use dupe::Dupe; + use pyrefly_build::handle::Handle; + use pyrefly_python::module_name::ModuleName; + use pyrefly_util::telemetry::QueueName; + use pyrefly_util::telemetry::TelemetryEventKind; + use pyrefly_util::telemetry::TelemetryServerState; + use pyrefly_util::thread_pool::TEST_THREAD_COUNT; + use uuid::Uuid; + + use super::*; + use crate::module::finder::DirEntryCache; + use crate::module::finder::find_import; + use crate::test::util::TestEnv; + + /// A recheck cancels the in-flight reads that block its commit. That cancellation belongs + /// to the request being aborted, so it should *not* survive into the restored transaction. + /// Clearing it must also keep the work the saved transaction already did, so the restored + /// transaction reuses the cached stdlib instead of recomputing it. + #[test] + fn test_restored_transaction_is_not_still_cancelled() { + let mut test_env = TestEnv::new(); + test_env.add("first", "x: int = 1\n"); + test_env.add("second", "y: int = \"not an int\"\n"); + let config_file = test_env.config(); + let sys_info = test_env.sys_info(); + let state = State::new(test_env.config_finder(), TEST_THREAD_COUNT); + let handle = |name: &str| { + let name = ModuleName::from_str(name); + let path = find_import(&config_file, name, None, None, &DirEntryCache::new(), None) + .finding() + .unwrap(); + Handle::new(name, path, sys_info.dupe()) + }; + + let mut manager = TransactionManager::default(); + let mut transaction = manager.non_committable_transaction(&state); + transaction.set_memory(test_env.get_memory()); + transaction.run(&[handle("first")], Require::Everything, None); + let old_cancellation = transaction.get_cancellation_handle(); + old_cancellation.cancel(); + let mut telemetry = TelemetryEvent::new_task( + TelemetryEventKind::InvalidateConfig, + TelemetryServerState { + has_sourcedb: false, + id: Uuid::new_v4(), + surface: None, + server_start_time: Instant::now(), + agent_session_id: None, + agent_invocation_id: None, + active_experiments: Vec::new(), + }, + QueueName::RecheckQueue, + 0, + Instant::now(), + ); + manager.save(transaction, &mut telemetry); + + let second = handle("second"); + let mut transaction = manager.saved_state.take().unwrap().restore().unwrap(); + assert!(old_cancellation.is_cancelled()); + assert!( + !transaction.get_cancellation_handle().is_cancelled(), + "restore should replace the saved cancellation handle" + ); + transaction.set_memory(test_env.get_memory()); + transaction.run(&[second.dupe()], Require::Everything, None); + assert!( + transaction.compute_stdlib_cached(), + "restored transaction should reuse the stdlib computed before the save" + ); + assert_eq!( + transaction + .get_errors([&second]) + .collect_errors() + .ordinary + .len(), + 1 + ); + } +}
diff --git a/pyrefly/lib/state/state.rs b/pyrefly/lib/state/state.rs index 23fd9e4..8558037 100644 --- a/pyrefly/lib/state/state.rs +++ b/pyrefly/lib/state/state.rs
@@ -702,14 +702,17 @@ impl<'a> TransactionData<'a> { /// Convert saved transaction data back into a full transaction. We can only restore if the /// underlying state is unchanged, otherwise the transaction data might make inconsistent - /// assumptions, in particular about deps/rdeps. + /// assumptions, in particular about deps/rdeps. A restored transaction always receives a + /// fresh cancellation handle (cancellation applies only to the consumer that saved it). pub(crate) fn restore(self) -> Result<Transaction<'a>, Duration> { let start = Timer::start(); let readable = self.state.state.read(); let state_lock_blocked = start.elapsed(); if self.base == readable.now { + let mut data = self; + data.todo.reset_cancellation(); Ok(Transaction { - data: self, + data, stats: Mutex::new(TelemetryTransactionStats { state_lock_blocked, ..Default::default()
diff --git a/pyrefly/lib/test/lsp/lsp_interaction/diagnostic.rs b/pyrefly/lib/test/lsp/lsp_interaction/diagnostic.rs index d873fec..bfe8070 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/diagnostic.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/diagnostic.rs
@@ -1214,7 +1214,16 @@ #[test] fn test_shows_stdlib_errors_for_multiple_versions_and_paths_with_force_on() { let test_files_root = get_test_files_root(); - let mut interaction = LspInteraction::new(); + let mut interaction = LspInteraction::new_with_args(LspInteractionArgs { + // Keep the production background-indexing path here: LazyBlocking + // closes the cancellation window instead of exercising recovery when a + // background recheck cancels an IDE request. + args: LspArgs { + indexing_mode: IndexingMode::LazyNonBlockingBackground, + ..LspInteractionArgs::default().args + }, + ..Default::default() + }); interaction.set_root(test_files_root.path().to_path_buf()); interaction .initialize(InitializeSettings {