diff --git a/.gitignore b/.gitignore index 5a4edd14ce..a48b1d2615 100644 --- a/.gitignore +++ b/.gitignore @@ -15,7 +15,7 @@ release target tmp __fuzz__ - +/graphify-out #==============================================================================# # Files to ignore #==============================================================================# diff --git a/libdd-data-pipeline-ffi/src/trace_exporter.rs b/libdd-data-pipeline-ffi/src/trace_exporter.rs index 6016c32dcb..f19edf3a55 100644 --- a/libdd-data-pipeline-ffi/src/trace_exporter.rs +++ b/libdd-data-pipeline-ffi/src/trace_exporter.rs @@ -901,6 +901,44 @@ pub unsafe extern "C" fn ddog_trace_exporter_free(handle: Box) { let _ = catch_panic!(handle.shutdown(None), Ok(())); } +/// Discards runtime-identity-bound state without sending buffered data, then frees the exporter. +/// +/// Why this exists: SDKs that restore from a process snapshot, such as dd-trace-py in a MicroVM +/// `/run` hook, can inherit buffered stats, telemetry state, and cached agent `/info` responses +/// from the snapshotted runtime. The normal [`ddog_trace_exporter_free`] path flushes before +/// freeing, which is correct for ordinary shutdown but wrong when the SDK is discarding inherited +/// state and immediately constructing a fresh exporter for the restored runtime identity. +/// +/// This function gives those snapshot-aware callers an explicit no-flush teardown API. Ordinary +/// non-snapshot shutdown callers should continue using [`ddog_trace_exporter_free`]. +/// +/// Returns `None` when the inherited state was discarded. Once a non-null handle is accepted, +/// `*handle` is always set to null, including when an ordinary error is returned. +/// +/// # Arguments +/// +/// * handle - A non-null pointer to the TraceExporter handle. The pointee must be non-null. +#[no_mangle] +pub unsafe extern "C" fn ddog_trace_exporter_free_without_flush( + handle: &mut *mut TraceExporter, +) -> Option> { + catch_panic!( + { + let Some(exporter) = NonNull::new(*handle) else { + return gen_error!(ErrorCode::InvalidArgument); + }; + let exporter = Box::from_raw(exporter.as_ptr()); + *handle = std::ptr::null_mut(); + + match exporter.shutdown_without_flush() { + Ok(()) => None, + Err(err) => Some(Box::new(ExporterError::from(err))), + } + }, + gen_error!(ErrorCode::Panic) + ) +} + /// Send traces to the Datadog Agent. /// /// # Arguments @@ -1351,6 +1389,28 @@ mod tests { } } + #[cfg_attr(miri, ignore)] + #[test] + fn exporter_free_without_flush_test() { + unsafe { + let mut config: MaybeUninit> = MaybeUninit::uninit(); + ddog_trace_exporter_config_new(NonNull::new_unchecked(&mut config).cast()); + let cfg = config.assume_init(); + + let mut exporter: MaybeUninit> = MaybeUninit::uninit(); + let error = ddog_trace_exporter_new( + NonNull::new_unchecked(&mut exporter).cast(), + Some(cfg.borrow()), + ); + assert!(error.is_none()); + + let mut exporter = Box::into_raw(exporter.assume_init()); + assert!(ddog_trace_exporter_free_without_flush(&mut exporter).is_none()); + assert!(exporter.is_null()); + ddog_trace_exporter_config_free(cfg); + } + } + #[test] fn exporter_send_test_arguments_test() { unsafe { diff --git a/libdd-data-pipeline-ffi/src/tracer.rs b/libdd-data-pipeline-ffi/src/tracer.rs index eabfd11e74..28afc9e23c 100644 --- a/libdd-data-pipeline-ffi/src/tracer.rs +++ b/libdd-data-pipeline-ffi/src/tracer.rs @@ -1554,8 +1554,8 @@ mod tests { assert!(event.0.attributes.is_empty()); // Misaligned pointer - let buf = [0u8; 32]; - let misaligned = buf.as_ptr().add(1).cast::(); + let data = [0i64; 3]; + let misaligned = data.as_ptr().cast::().add(1).cast::(); assert!(!misaligned.is_aligned()); let err = ddog_tracer_span_event_set_int_array( Some(&mut event), diff --git a/libdd-data-pipeline/src/agent_info/mod.rs b/libdd-data-pipeline/src/agent_info/mod.rs index 5675f2c00a..ba7bed6718 100644 --- a/libdd-data-pipeline/src/agent_info/mod.rs +++ b/libdd-data-pipeline/src/agent_info/mod.rs @@ -29,6 +29,11 @@ pub fn get_agent_info() -> Option> { AGENT_INFO_CACHE.load_full() } +/// Clear the global agent info cache after restoring from an inherited runtime identity. +pub(crate) fn clear_cache_for_runtime_identity_refresh() { + AGENT_INFO_CACHE.store(None); +} + pub use fetcher::{ fetch_info, fetch_info_with_state, AgentInfoFetcher, FetchInfoStatus, ResponseObserver, }; @@ -36,5 +41,5 @@ pub use fetcher::{ #[cfg(test)] /// Clear the global agent info cache for test isolation in envs where nextest isn't used. pub fn clear_cache_for_test() { - AGENT_INFO_CACHE.store(None); + clear_cache_for_runtime_identity_refresh(); } diff --git a/libdd-data-pipeline/src/trace_buffer/mod.rs b/libdd-data-pipeline/src/trace_buffer/mod.rs index 715c087857..bc05d4e7c5 100644 --- a/libdd-data-pipeline/src/trace_buffer/mod.rs +++ b/libdd-data-pipeline/src/trace_buffer/mod.rs @@ -634,6 +634,20 @@ impl Receiver { Ok(()) } + /// Permanently close the channel, dropping any buffered chunks without exporting them. + /// + /// Unlike [`Self::reset`], which reopens the channel for a worker that keeps running (e.g. + /// after a fork), this is for a worker that is never run again. Leaving the channel + /// `Running` would let a surviving [`Sender`] keep queuing chunks that no receiver will ever + /// take, and would leave callers blocked forever in [`Sender::wait_close_done`] or + /// [`Sender::wait_flush_done`]. + fn discard(&self) -> Result<(), MutexPoisonedError> { + let mut state = self.lock_state()?; + state.batch.reset(); + self.waiter.mark_stopped(state); + Ok(()) + } + async fn receive(&self, timeout: Duration) -> Result>, MutexPoisonedError> { loop { // Enable the notify future BEFORE acquiring the lock to avoid lost wakeups: @@ -950,6 +964,13 @@ impl Worker for TraceExporterWorker { fn reset(&mut self) { let _ = self.rx.reset(); } + + // Override the default (`Worker::discard`'s `reset()` forwarding): `reset()` reopens the + // channel for continued use, which is wrong once this worker is permanently discarded (see + // `Receiver::discard`). + fn discard(&mut self) { + let _ = self.rx.discard(); + } } #[cfg(test)] @@ -958,7 +979,7 @@ mod tests { use std::sync::Arc; use std::time::{Duration, Instant}; - use libdd_shared_runtime::{BlockingRuntime, ForkSafeRuntime, SharedRuntime}; + use libdd_shared_runtime::{BlockingRuntime, ForkSafeRuntime, SharedRuntime, Worker}; use crate::trace_buffer::{BufferSize, Export, TraceBuffer, TraceBufferConfig}; use crate::trace_exporter::agent_response::AgentResponse; @@ -1550,4 +1571,46 @@ mod tests { assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 2); rt.shutdown(None).unwrap(); } + + #[test] + fn test_worker_reset_drops_buffered_chunk() { + let (sender, mut worker) = TraceBuffer::new( + TraceBufferConfig::default().flush_threshold_bytes(2), + Box::new(|_| {}), + Box::new(AssertExporter( + Box::new(|_| panic!("discard must not export buffered chunks")), + Arc::new(tokio::sync::Semaphore::new(0)), + )), + ); + + sender.send_chunk(vec![()]).unwrap(); + assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 1); + + worker.reset(); + assert_eq!(sender.queue_metrics().get_metrics().spans_queued, 0); + } + + #[test] + fn test_worker_discard_closes_channel_without_exporting() { + let (sender, mut worker) = TraceBuffer::new( + TraceBufferConfig::default().flush_threshold_bytes(2), + Box::new(|_| {}), + Box::new(AssertExporter( + Box::new(|_| panic!("discard must not export buffered chunks")), + Arc::new(tokio::sync::Semaphore::new(0)), + )), + ); + + sender.send_chunk(vec![()]).unwrap(); + + worker.discard(); + + // Unlike `reset`, which reopens the channel for a worker that keeps running (e.g. after + // a fork), `discard` is for a worker that is never run again. A channel left `Running` + // would let this sender keep queuing chunks that no receiver will ever take. + assert!(matches!( + sender.send_chunk(vec![()]), + Err(TraceBufferError::AlreadyClosed) + )); + } } diff --git a/libdd-data-pipeline/src/trace_exporter/mod.rs b/libdd-data-pipeline/src/trace_exporter/mod.rs index 7e522a4d3a..c91dbcdf78 100644 --- a/libdd-data-pipeline/src/trace_exporter/mod.rs +++ b/libdd-data-pipeline/src/trace_exporter/mod.rs @@ -46,6 +46,7 @@ use libdd_capabilities::{HttpClientCapability, LogWriterCapability, MaybeSend, S use libdd_common::tag::Tag; use libdd_common::Endpoint; use libdd_dogstatsd_client::DogStatsDClient; +use libdd_shared_runtime::shared_runtime::runtime_identity_refresh; #[cfg(not(target_arch = "wasm32"))] use libdd_shared_runtime::BlockingRuntime; use libdd_shared_runtime::{SharedRuntime, WorkerHandle}; @@ -321,6 +322,24 @@ impl< runtime.block_on(self.shutdown_async(timeout))? } + /// Discard inherited runtime-identity state without sending buffered worker data. + /// + /// This exists for SDKs that restore from a process snapshot, such as a MicroVM `/run` + /// hook. The restored process can inherit buffered stats, telemetry state, and cached + /// `/info` responses from the snapshotted runtime. Normal [`Self::shutdown`] intentionally + /// flushes those buffers, so this separate API gives snapshot-aware callers an explicit + /// opt-in path that discards inherited state before they create a fresh exporter. + /// + /// Non-snapshot shutdown callers should continue using [`Self::shutdown`]. + #[cfg(not(target_arch = "wasm32"))] + pub fn shutdown_without_flush(self) -> Result<(), TraceExporterError> + where + R: BlockingRuntime, + { + let runtime = self.shared_runtime.clone(); + runtime.block_on(self.shutdown_without_flush_async())? + } + /// Async version of [`Self::shutdown`]. /// /// # Errors @@ -344,6 +363,16 @@ impl< } } + /// Async version of [`Self::shutdown_without_flush`]. + /// + /// # Cancel safety + /// This function is *NOT* cancel safe. If cancelled, workers can be left in an invalid state. + pub async fn shutdown_without_flush_async(self) -> Result<(), TraceExporterError> { + let discard_result = self.discard_workers_without_flush().await; + agent_info::clear_cache_for_runtime_identity_refresh(); + discard_result + } + async fn shutdown_workers(self) { let mut handles: Vec = Vec::new(); @@ -375,6 +404,50 @@ impl< } } + async fn discard_workers_without_flush(&self) -> Result<(), TraceExporterError> { + let mut handles: Vec = Vec::new(); + + if let StatsComputationStatus::Enabled { worker_handle, .. } = + &**self.client_side_stats.status.load() + { + handles.push(worker_handle.clone()); + } + + if let Some(info_fetcher) = &self.workers.info_fetcher { + handles.push(info_fetcher.clone()); + } + + if let Some(dogstatsd) = &self.workers.dogstatsd { + handles.push(dogstatsd.clone()) + } + + #[cfg(feature = "telemetry")] + if let Some(telemetry) = &self.workers.telemetry { + handles.push(telemetry.clone()); + } + + let mut futures: FuturesUnordered<_> = handles + .into_iter() + .map(runtime_identity_refresh::discard_worker_without_flush) + .collect(); + + let mut first_error = None; + while let Some(result) = futures.next().await { + if let Err(e) = result { + error!("Worker failed to discard: {:?}", e); + if first_error.is_none() { + first_error = Some(e); + } + } + } + + first_error + .map(|err| { + TraceExporterError::Internal(InternalErrorKind::InvalidWorkerState(err.to_string())) + }) + .map_or(Ok(()), Err) + } + /// Send msgpack serialized traces to the agent. /// /// Sync facade over [`Self::send_async`]; panics inside an existing tokio context. @@ -2729,6 +2802,110 @@ mod single_threaded_tests { mock_stats.assert(); } + #[cfg_attr(miri, ignore)] + #[test] + fn test_shutdown_without_flush_discards_buffered_stats_and_agent_info() { + // Clear the agent info cache to ensure test isolation + agent_info::clear_cache_for_test(); + + let server = MockServer::start(); + + let mock_traces = server.mock(|when, then| { + when.method(POST) + .header("Content-type", "application/msgpack") + .path(V04_TRACES_ENDPOINT); + then.status(200).body(""); + }); + + let mock_stats = server.mock(|when, then| { + when.method(POST) + .header("Content-type", "application/msgpack") + .path(STATS_ENDPOINT); + then.status(200).body(""); + }); + + let _mock_info = server.mock(|when, then| { + when.method(GET).path(INFO_ENDPOINT); + then.status(200) + .header("content-type", "application/json") + .header("datadog-agent-state", "1") + .body(format!( + r#"{{"version":"1","client_drop_p0s":true,"endpoints":["{V04_TRACES_ENDPOINT}","{STATS_ENDPOINT}"]}}"# + )); + }); + + let runtime = Arc::new(ForkSafeRuntime::new().unwrap()); + + let mut builder = TraceExporter::::builder(); + builder + .set_url(&server.url("/")) + .set_service("test") + .set_env("staging") + .set_tracer_version("v0.1") + .set_language("nodejs") + .set_language_version("1.0") + .set_language_interpreter("v8") + .set_input_format(TraceExporterInputFormat::V04) + .set_output_format(TraceExporterOutputFormat::V04) + .set_shared_runtime(runtime.clone()) + .enable_stats(Duration::from_secs(10)); + let exporter = builder.build::().unwrap(); + + let trace_chunk = vec![SpanBytes { + service: "test".into(), + name: "test".into(), + resource: "test".into(), + r#type: "test".into(), + duration: 10, + ..Default::default() + }]; + let data = msgpack_encoder::v04::to_vec_from_v04(&[trace_chunk]); + + // Wait for the info fetcher so sending the trace starts stats computation. + while agent_info::get_agent_info().is_none() { + std::thread::sleep(Duration::from_millis(100)); + } + // The mock returns an empty response body, which is reported as an agent response error + // after the trace has still been processed for client-side stats. + assert!(exporter.send(data.as_ref()).is_err()); + + let start_time = std::time::Instant::now(); + while !exporter.is_stats_worker_active() { + if start_time.elapsed() > Duration::from_secs(10) { + panic!("Timeout waiting for stats worker to become active"); + } + std::thread::sleep(Duration::from_millis(10)); + } + + exporter.shutdown_without_flush().unwrap(); + assert!(agent_info::get_agent_info().is_none()); + runtime.shutdown(None).unwrap(); + + mock_traces.assert(); + assert_eq!( + mock_stats.calls(), + 0, + "discard must not flush buffered stats" + ); + } + + #[test] + fn test_shutdown_without_flush_reports_prior_runtime_shutdown() { + let runtime = Arc::new(ForkSafeRuntime::new().unwrap()); + let mut builder = TraceExporter::::builder(); + builder.set_shared_runtime(runtime.clone()); + let exporter = builder.build::().unwrap(); + + runtime.shutdown(None).unwrap(); + + assert!(matches!( + exporter.shutdown_without_flush(), + Err(TraceExporterError::Internal( + InternalErrorKind::InvalidWorkerState(_) + )) + )); + } + #[cfg_attr(miri, ignore)] #[test] fn test_shutdown_with_timeout() { diff --git a/libdd-shared-runtime/src/shared_runtime/fork_safe.rs b/libdd-shared-runtime/src/shared_runtime/fork_safe.rs index 06bd56f6ff..413cdaee0a 100644 --- a/libdd-shared-runtime/src/shared_runtime/fork_safe.rs +++ b/libdd-shared-runtime/src/shared_runtime/fork_safe.rs @@ -242,6 +242,7 @@ impl BlockingRuntime for ForkSafeRuntime { #[cfg(test)] mod tests { use super::*; + use crate::shared_runtime::runtime_identity_refresh::discard_worker_without_flush; use async_trait::async_trait; use std::sync::mpsc::{channel, Receiver, Sender}; use std::time::Duration; @@ -330,6 +331,52 @@ mod tests { assert_eq!(last, -1); } + #[test] + fn test_runtime_identity_refresh_discard_does_not_shutdown() { + #[derive(Debug)] + struct DiscardWorker(Sender); + + #[async_trait] + impl Worker for DiscardWorker { + async fn run(&mut self) {} + + async fn trigger(&mut self) { + std::future::pending::<()>().await; + } + + fn reset(&mut self) { + let _ = self.0.send(-2); + } + + async fn shutdown(&mut self) { + let _ = self.0.send(-1); + } + } + + let rt = tokio::runtime::Runtime::new().unwrap(); + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let (sender, receiver) = channel(); + + let handle = shared_runtime + .spawn_worker(DiscardWorker(sender), true) + .unwrap(); + + rt.block_on(async { + assert!(discard_worker_without_flush(handle).await.is_ok()); + }); + + assert_eq!(shared_runtime.workers.lock_or_panic().len(), 0); + assert_eq!( + receiver + .recv_timeout(Duration::from_secs(1)) + .expect("discard did not run"), + -2 + ); + assert!( + receiver.recv_timeout(Duration::from_millis(200)).is_err(), + "discard must not run shutdown" + ); + } #[test] fn test_before_and_after_fork_parent() { let shared_runtime = ForkSafeRuntime::new().unwrap(); diff --git a/libdd-shared-runtime/src/shared_runtime/mod.rs b/libdd-shared-runtime/src/shared_runtime/mod.rs index 5868b2bbef..124ca7a5e8 100644 --- a/libdd-shared-runtime/src/shared_runtime/mod.rs +++ b/libdd-shared-runtime/src/shared_runtime/mod.rs @@ -259,6 +259,45 @@ impl WorkerHandle { } } +/// Operations used when a restored runtime must discard inherited identity-bound state. +pub mod runtime_identity_refresh { + use super::{WorkerEntry, WorkerHandle, WorkerHandleError}; + use libdd_common::MutexExt; + + /// Stop a worker and clear restartable state without executing shutdown flushing. + /// + /// This is intended for runtime identity refreshes, such as a MicroVM `/run` hook replacing + /// an exporter after snapshot restore. Normal shutdown and fork paths should keep using + /// [`WorkerHandle::stop`] and the fork APIs so they preserve their existing behavior. + /// Keep this out of `WorkerHandle`'s general API so regular callers do not accidentally + /// choose discard semantics when they meant ordinary shutdown. + /// + /// # Errors + /// Returns an error if the worker has already been stopped. + /// + /// # Cancel safety + /// This function is *NOT* cancel safe and should not be called from + /// [`Worker::trigger`](crate::worker::Worker::trigger). If cancelled, the discarded worker can + /// end up in an invalid state. + pub async fn discard_worker_without_flush( + handle: WorkerHandle, + ) -> Result<(), WorkerHandleError> { + let mut worker = { + let mut workers_lock = handle.workers.lock_or_panic(); + let Some(position) = workers_lock + .iter() + .position(|entry| entry.id == handle.worker_id) + else { + return Err(WorkerHandleError::AlreadyStopped); + }; + let WorkerEntry { worker, .. } = workers_lock.swap_remove(position); + worker + }; + worker.discard().await?; + Ok(()) + } +} + /// Errors that can occur when using a `SharedRuntime` implementation. #[derive(Debug)] pub enum SharedRuntimeError { diff --git a/libdd-shared-runtime/src/shared_runtime/pausable_worker.rs b/libdd-shared-runtime/src/shared_runtime/pausable_worker.rs index e39291b48f..50af8fb267 100644 --- a/libdd-shared-runtime/src/shared_runtime/pausable_worker.rs +++ b/libdd-shared-runtime/src/shared_runtime/pausable_worker.rs @@ -46,7 +46,15 @@ pub(super) fn tokio_spawn_fn( pub enum PausableWorker { Running { handle: WorkerJoinHandle, + /// Cancelled by [`Self::pause`]. Only guards `trigger`/`initial_trigger`, never + /// `Worker::run`, so a normal pause (before a fork, or on ordinary shutdown) always lets + /// an in-flight `run` finish. stop_token: CancellationToken, + /// Cancelled by [`Self::discard`]. Unlike `stop_token`, this one also guards `Worker::run` + /// itself, so a runtime identity refresh can cut off an in-flight run instead of waiting + /// for it. Kept as a separate token so `discard`'s stronger cancellation never affects + /// `pause`. + discard_token: CancellationToken, }, Paused { worker: T, @@ -115,7 +123,13 @@ impl PausableWorker { }; let stop_token = CancellationToken::new(); + // Runtime identity refresh needs a stronger cancellation path than normal pause: + // it must discard in-flight work inherited from a process snapshot. Keep this + // separate from `stop_token` so normal pause/fork behavior still waits for + // `Worker::run` to finish. + let discard_token = CancellationToken::new(); let cloned_token = stop_token.clone(); + let cloned_discard_token = discard_token.clone(); let future = Box::pin(async move { // First iteration using initial_trigger. // @@ -125,32 +139,54 @@ impl PausableWorker { // not keep the runtime scheduler alive after this task is dropped. select! { biased; + _ = cloned_discard_token.cancelled() => { + return worker; + } _ = cloned_token.cancelled() => { return worker; } _ = WeakWakerFuture::new(worker.initial_trigger()) => { - worker.run().await; } } + select! { + biased; + _ = cloned_discard_token.cancelled() => { + return worker; + } + _ = worker.run() => {} + } // Regular iterations loop { select! { biased; + _ = cloned_discard_token.cancelled() => { + break; + } _ = cloned_token.cancelled() => { break; } _ = WeakWakerFuture::new(worker.trigger()) => { - worker.run().await; } } + select! { + biased; + _ = cloned_discard_token.cancelled() => { + break; + } + _ = worker.run() => {} + } } worker }); let handle = spawn_fn(future); - *self = PausableWorker::Running { handle, stop_token }; + *self = PausableWorker::Running { + handle, + stop_token, + discard_token, + }; Ok(()) } PausableWorker::InvalidState => Err(PausableWorkerError::InvalidState), @@ -165,8 +201,9 @@ impl PausableWorker { match self { PausableWorker::Running { .. } => { debug!("Waiting for worker to pause"); - let PausableWorker::Running { handle, stop_token } = - std::mem::replace(self, PausableWorker::InvalidState) + let PausableWorker::Running { + handle, stop_token, .. + } = std::mem::replace(self, PausableWorker::InvalidState) else { // Unreachable return Ok(()); @@ -190,6 +227,49 @@ impl PausableWorker { } } + /// Stop the worker for a runtime identity refresh and discard its inherited state, without + /// running the ordinary shutdown flush path. + /// + /// Calls [`Worker::discard`] rather than [`Worker::reset`]: this worker instance is never run + /// again, whereas `reset` is written for a worker that keeps running afterward (e.g. a forked + /// child). Using `reset` here would incorrectly reopen resources (e.g. a channel) that a + /// permanently discarded worker should instead close. + pub(super) async fn discard(&mut self) -> Result<(), PausableWorkerError> { + match self { + PausableWorker::Running { .. } => { + debug!("Waiting for worker to discard"); + let PausableWorker::Running { + handle, + discard_token, + .. + } = std::mem::replace(self, PausableWorker::InvalidState) + else { + // Unreachable + return Ok(()); + }; + + if !discard_token.is_cancelled() { + discard_token.cancel(); + } + + if let Ok(mut worker) = handle.await { + worker.discard(); + debug!(?worker, "Worker discarded successfully"); + *self = PausableWorker::Paused { worker }; + Ok(()) + } else { + *self = PausableWorker::InvalidState; + Err(PausableWorkerError::TaskAborted) + } + } + PausableWorker::Paused { worker } => { + worker.discard(); + Ok(()) + } + PausableWorker::InvalidState => Err(PausableWorkerError::InvalidState), + } + } + /// Reset the worker state (e.g. in a fork child). pub fn reset(&mut self) { if let PausableWorker::Paused { worker } = self { @@ -235,6 +315,42 @@ mod tests { } } + /// A worker whose `run` can be held "in flight" for a controlled duration, used to test + /// `pause`/`discard` racing against an active `run` call. + /// + /// `TestWorker` above can't exercise that race: its `run` returns immediately, so there is + /// no window in which `pause`/`discard` can observe it as in-flight. Each lifecycle method + /// reports itself on `sender` so a test can assert both whether `run` was allowed to finish + /// and the order lifecycle methods ran in. + #[derive(Debug)] + struct TestInFlightWorker { + sender: Sender<&'static str>, + run_duration: Duration, + } + + #[async_trait] + impl Worker for TestInFlightWorker { + async fn run(&mut self) { + let _ = self.sender.send("run-started"); + sleep(self.run_duration).await; + let _ = self.sender.send("run-finished"); + } + + async fn trigger(&mut self) { + std::future::pending::<()>().await; + } + + async fn initial_trigger(&mut self) {} + + fn reset(&mut self) { + let _ = self.sender.send("reset"); + } + + async fn shutdown(&mut self) { + let _ = self.sender.send("shutdown"); + } + } + #[test] fn test_restart() { let (sender, receiver) = channel::(); @@ -256,4 +372,64 @@ mod tests { pausable_worker.start(tokio_spawn_fn(&handle)).unwrap(); assert_eq!(receiver.recv().unwrap(), next_message); } + + #[test] + fn test_pause_waits_for_in_flight_run() { + let (sender, receiver) = channel::<&'static str>(); + let worker = TestInFlightWorker { + sender, + run_duration: Duration::from_millis(100), + }; + let runtime = Builder::new_multi_thread().enable_time().build().unwrap(); + let handle = runtime.handle().clone(); + let mut pausable_worker: PausableWorker> = + PausableWorker::new(Box::new(worker)); + + pausable_worker.start(tokio_spawn_fn(&handle)).unwrap(); + assert_eq!( + receiver.recv_timeout(Duration::from_secs(1)).unwrap(), + "run-started" + ); + + runtime.block_on(async { pausable_worker.pause().await.unwrap() }); + + assert_eq!( + receiver.recv_timeout(Duration::from_secs(1)).unwrap(), + "run-finished" + ); + assert!( + receiver.recv_timeout(Duration::from_millis(200)).is_err(), + "pause must not reset or shutdown the worker" + ); + } + + #[test] + fn test_discard_cancels_in_flight_run_and_resets() { + let (sender, receiver) = channel::<&'static str>(); + let worker = TestInFlightWorker { + sender, + run_duration: Duration::from_secs(60), + }; + let runtime = Builder::new_multi_thread().enable_time().build().unwrap(); + let handle = runtime.handle().clone(); + let mut pausable_worker: PausableWorker> = + PausableWorker::new(Box::new(worker)); + + pausable_worker.start(tokio_spawn_fn(&handle)).unwrap(); + assert_eq!( + receiver.recv_timeout(Duration::from_secs(1)).unwrap(), + "run-started" + ); + + runtime.block_on(async { pausable_worker.discard().await.unwrap() }); + + assert_eq!( + receiver.recv_timeout(Duration::from_secs(1)).unwrap(), + "reset" + ); + assert!( + receiver.recv_timeout(Duration::from_millis(200)).is_err(), + "discard must cancel the in-flight run and must not shutdown the worker" + ); + } } diff --git a/libdd-shared-runtime/src/worker.rs b/libdd-shared-runtime/src/worker.rs index 4666ddbf2d..f20057f8bc 100644 --- a/libdd-shared-runtime/src/worker.rs +++ b/libdd-shared-runtime/src/worker.rs @@ -38,6 +38,18 @@ pub trait Worker: std::fmt::Debug + MaybeSend { /// Reset the worker state. Called in the child after a fork to cleanup parent state. fn reset(&mut self) {} + /// Reset the worker state after it has been permanently removed from its runtime, e.g. for a + /// runtime identity refresh (see + /// [`runtime_identity_refresh`](crate::shared_runtime::runtime_identity_refresh)). + /// + /// Unlike [`reset`](Self::reset), which prepares the worker to keep running (e.g. in a forked + /// child), this worker instance is never run again. Override this when a worker holds + /// resources whose "keep running" reset behavior would be wrong once discarded permanently + /// (e.g. a channel that must be closed rather than reopened). + fn discard(&mut self) { + self.reset(); + } + /// Hook called when the app is shutting down. Can be used to flush remaining data. async fn shutdown(&mut self) {} } @@ -62,6 +74,10 @@ impl Worker for Box { (**self).reset() } + fn discard(&mut self) { + (**self).discard() + } + async fn shutdown(&mut self) { (**self).shutdown().await }