From fa9b177fe60b9aed6c89ef7f3ae4ceae6032ecf1 Mon Sep 17 00:00:00 2001 From: Derek Date: Thu, 12 Nov 2026 18:04:32 +1100 Subject: [PATCH] fix(reload): start, stop and rebuild sources on a config reload Added sources start on reload. Removed sources stop. A source changed under the same id is rebuilt from its new settings. Before, an added source only logged that a restart was needed. Startup and reload build sources through Scheduler::start_sources and a SourceFactory. main's DriverFactory holds the old startup loop. Each reload runs resolve_config_specs first, so env: and vault: specs never reach a driver as literal values. A connection is rebuilt only when the hash of its own block, secrets included, changes. A fetch task can sit in a flush or cursor commit outside its cancellation selects. Cancelling alone let two generations share a checkpoint. A replaced or removed generation moves to a stopping set. Its task is joined and its driver stopped, each within 10 s, before another starts. The fetch task stops its own driver before it ends, replacing a drop guard that spawned the stop unawaited. A failed join still gets its driver stopped. Generations whose shutdown is unconfirmed stay in the stopping set. The watcher retries every 30 s. The source restarts once the old generation is down. RunningSources::shutdown waits for every fetch task, running or stopping, then stops every driver it holds. main calls it in place of its drain over running drivers. A custom shape can stop through the new RowSource::stop. New settings that fail to build are logged. The old generation keeps running. Remove PipelineState::reload_config because it had no caller. Refs #134 --- crates/core/src/lib.rs | 6 + crates/fetcher/src/driver.rs | 11 +- crates/fetcher/src/main.rs | 419 +++++------ crates/fetcher/src/pipeline/mod.rs | 42 -- crates/fetcher/src/scheduler/mod.rs | 1044 +++++++++++++++++++++++---- docs/internals.md | 4 +- 6 files changed, 1119 insertions(+), 407 deletions(-) diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 7fbe751..8146724 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -357,6 +357,12 @@ pub trait RowSource: Send + Sync { /// Health check: credentials resolve and the provider answers. fn probe(&self) -> BoxFuture<'_, Result<()>>; + + /// Stop what the source runs between ticks and wait for it to finish. + /// Its driver calls this when it stops, possibly more than once. + fn stop(&self) -> BoxFuture<'_, ()> { + Box::pin(std::future::ready(())) + } } #[cfg(test)] diff --git a/crates/fetcher/src/driver.rs b/crates/fetcher/src/driver.rs index a4b0dbc..be69660 100644 --- a/crates/fetcher/src/driver.rs +++ b/crates/fetcher/src/driver.rs @@ -410,11 +410,14 @@ impl Driver { /// Stop what the shape runs in the background and wait for it to finish. /// /// A file tailer writes its final checkpoint, including the last - /// acknowledged tick, which only the driver's checkpoint store holds. - /// Every other shape runs nothing between ticks. Safe to call again. + /// acknowledged tick, which only the driver's checkpoint store holds. A + /// custom shape runs its own [`RowSource::stop`], and a REST or database + /// shape runs nothing between ticks. Safe to call again. pub async fn stop(&self) { - let Shape::File(shape) = &self.shape else { - return; + let shape = match &self.shape { + Shape::File(shape) => shape, + Shape::Custom(source) => return source.stop().await, + Shape::Rest(_) | Shape::Db(_) => return, }; let mut committed = HashMap::new(); for unit in shape.units() { diff --git a/crates/fetcher/src/main.rs b/crates/fetcher/src/main.rs index cca896b..9464f2e 100644 --- a/crates/fetcher/src/main.rs +++ b/crates/fetcher/src/main.rs @@ -16,7 +16,6 @@ #[global_allocator] static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; -use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -39,11 +38,11 @@ use dfe_fetcher::extractor::container::ContainerExtractor; use dfe_fetcher::extractor::vector::VectorManager; use dfe_fetcher::ingest; use dfe_fetcher::metrics::Metrics; -use dfe_fetcher::pipeline::{IntakeAcks, Orchestrator}; -use dfe_fetcher::scheduler::Scheduler; -use dfe_fetcher_core::SourceMaturity; -use dfe_fetcher_db::DbShape; -use dfe_fetcher_file::FileShape; +use dfe_fetcher::pipeline::{IntakeAcks, Orchestrator, PipelineState}; +use dfe_fetcher::scheduler::{Planned, Scheduler, SourceFactory, fingerprint}; +use dfe_fetcher_db::{DbInstance, DbShape}; +use dfe_fetcher_file::{FileInstance, FileShape}; +use dfe_fetcher_rest::{ExchangeClient, HttpClient, RestInstance}; use scalo::governor::{MemoryPressureSource, PressureSource, UnifiedPressure}; use scalo::memory::MemoryGuard; @@ -270,20 +269,146 @@ async fn main() { } } -/// How long shutdown waits for the source drivers to stop what they run, so a -/// file tailer's final checkpoint is written before the process exits. -const DRIVER_STOP_TIMEOUT: Duration = Duration::from_secs(10); +/// What one connection's driver is built from: its own block, and the +/// deployment-wide settings the driver bakes in. +#[derive(serde::Serialize)] +struct SourceSpec { + kind: SourceKind, + in_flight: usize, +} + +/// The block one connection is configured by. +#[derive(serde::Serialize)] +enum SourceKind { + Rest(Box), + Db(Box), + File(Box), +} + +/// Every connection `config` schedules, one per connection id: the typed +/// blocks that are shipped profiles underneath (`sources.github`, +/// `sources.okta`, ...), then `sources.rest`, `sources.db` and `sources.file`. +fn plan_sources(config: &Config) -> anyhow::Result>> { + let in_flight = config.accumulate.in_flight; + let planned = |id: String, interval_secs: Option, kind: SourceKind| { + let spec = SourceSpec { kind, in_flight }; + anyhow::Ok(Planned { + fingerprint: fingerprint(&spec)?, + id, + interval_secs, + spec, + }) + }; + let rest = config + .sources + .builtin_instances()? + .into_iter() + .map(|b| (b.connection_id, b.instance)) + .chain( + config + .sources + .rest + .iter() + .filter(|(_, instance)| instance.enabled) + .map(|(id, instance)| (id.clone(), instance.clone())), + ); + let mut out = Vec::new(); + for (id, instance) in rest { + let interval = instance.interval_secs; + out.push(planned(id, interval, SourceKind::Rest(Box::new(instance)))?); + } + for (id, instance) in config.sources.db.iter().filter(|(_, i)| i.enabled) { + let kind = SourceKind::Db(Box::new(instance.clone())); + out.push(planned(id.clone(), instance.interval_secs, kind)?); + } + for (id, instance) in config.sources.file.iter().filter(|(_, i)| i.enabled) { + let kind = SourceKind::File(Box::new(instance.clone())); + out.push(planned(id.clone(), instance.interval_secs, kind)?); + } + Ok(out) +} + +/// Builds every connection's driver over the process's one pipeline, memory +/// guard, checkpoint store and HTTP clients. +struct DriverFactory { + instance_id: String, + pipeline_state: Arc, + metrics: Arc, + pressure: Option>, + checkpoints: Option>, + http_client: HttpClient, + exchange_client: Arc, +} + +impl SourceFactory for DriverFactory { + type Spec = SourceSpec; -/// One connection's fetch task, ready to schedule: a source type expands -/// into one entry per connection, each keyed on its connection id (cursor -/// key + metric/log label). -struct SpawnEntry { - driver: Arc, - interval_secs: Option, + fn plan(&self, config: &Config) -> anyhow::Result>> { + plan_sources(config) + } - /// A child of the global shutdown token, held by both the driver and the - /// fetch task, so a reload that drops this source stops just this one. - cancel: CancellationToken, + /// The driver takes its source's OWN cancel token, so dropping the source + /// from the config aborts its in-flight tick through the driver's + /// non-alerting shutdown path. + fn build( + &self, + id: &str, + spec: SourceSpec, + cancel: CancellationToken, + ) -> anyhow::Result { + // Every block a cursor fetches and every chunk a reader holds is + // leased on the memory guard until its rows are out. + let lease: Arc = Arc::new( + dfe_fetcher::driver::GuardLease(Arc::clone(self.pipeline_state.memory_guard())), + ); + let shape = match spec.kind { + SourceKind::Rest(instance) => { + let profile = dfe_fetcher_rest::profile::bound::resolve_profile( + &instance, + dfe_fetcher::profiles::shipped(), + )?; + Shape::for_rest_instance( + &profile, + &instance, + id, + self.http_client.clone(), + &self.exchange_client, + )? + } + SourceKind::Db(instance) => Shape::Db(Box::new( + DbShape::from_instance(&instance, id, &lease) + .map_err(|e| anyhow::anyhow!("sources.db.{id}: {e}"))?, + )), + SourceKind::File(instance) => Shape::File(Box::new( + FileShape::from_instance(&instance, id, &lease) + .map_err(|e| anyhow::anyhow!("sources.file.{id}: {e}"))?, + )), + }; + // A scheduled source's cursor advances only once its records are delivered. + self.pipeline_state.publish_guarantee( + "scheduled", + Some(&IntakeAcks { + kind: scalo::transport::AckKind::Pull, + enabled: true, + }), + ); + Ok(Driver::new(DriverParts { + shape, + connection_id: id.to_owned(), + instance_id: self.instance_id.clone(), + shared_config: self.pipeline_state.shared_config(), + emitter: Emitter::new( + Arc::clone(&self.pipeline_state), + Arc::clone(&self.metrics), + spec.in_flight, + ), + pressure: self.pressure.clone(), + memory_guard: Arc::clone(self.pipeline_state.memory_guard()), + checkpoints: self.checkpoints.clone(), + metrics: Arc::clone(&self.metrics), + shutdown: cancel, + })) + } } /// Main service loop -- called by the ServiceApp lifecycle after logging and config. @@ -501,188 +626,38 @@ async fn run_fetcher_service( } }; - // Framework drivers commit unit checkpoints to the same store the - // scheduler keeps its window cursors in. - let framework_cursor_store = cursor_store.clone(); - // Create scheduler (reads interval/jitter/window_hours from shared_config per tick) let scheduler = Scheduler::new( &config.scheduler, orchestrator.shared_config(), - cursor_store, + cursor_store.clone(), instance_id.clone(), ); - // One spawn entry per connection, keyed on the connection id (a typed - // block's `connections[].id`, or the type name for a single connection). - let mut entries: Vec = Vec::new(); - // One Driver per connection, all sharing the memory guard's pressure // latch so a buffer that grows anywhere pauses polling everywhere. let pressure = build_pressure(&config, pipeline_state.memory_guard())?; - // Each driver takes its source's OWN cancel token, not the global one, so - // dropping the source from the config aborts its in-flight tick through - // the driver's non-alerting shutdown path. - let framework_driver = |shape: Shape, id: &str, cancel: CancellationToken| { - Driver::new(DriverParts { - shape, - connection_id: id.to_owned(), - instance_id: instance_id.clone(), - shared_config: orchestrator.shared_config(), - emitter: Emitter::new( - Arc::clone(&pipeline_state), - Arc::clone(&metrics), - config.accumulate.in_flight, - ), - pressure: pressure.clone(), - memory_guard: Arc::clone(pipeline_state.memory_guard()), - checkpoints: framework_cursor_store.clone(), - metrics: Arc::clone(&metrics), - shutdown: cancel, - }) + let factory = DriverFactory { + instance_id: instance_id.clone(), + pipeline_state: Arc::clone(&pipeline_state), + metrics: Arc::clone(&metrics), + pressure: pressure.clone(), + // Framework drivers commit unit checkpoints to the same store the + // scheduler keeps its window cursors in. + checkpoints: cursor_store, + http_client: dfe_fetcher_rest::http_client()?, + // One exchange client for the whole process: every instance's + // credential mode posts its token exchange through it. + exchange_client: dfe_fetcher_rest::exchange_client()?, }; - - // REST instances: the typed blocks that are shipped profiles underneath - // (`sources.github`, `sources.okta`, `sources.slack`, ...), then every - // `sources.rest` entry. - let http_client = dfe_fetcher_rest::http_client()?; - // One exchange client for the whole process: every instance's credential - // mode posts its token exchange through it. - let exchange_client = dfe_fetcher_rest::exchange_client()?; - let builtin = config.sources.builtin_instances()?; - let rest_instances = builtin - .iter() - .map(|b| (b.connection_id.as_str(), &b.instance)) - .chain( - config - .sources - .rest - .iter() - .filter(|(_, instance)| instance.enabled) - .map(|(id, instance)| (id.as_str(), instance)), - ); - for (id, instance) in rest_instances { - let profile = dfe_fetcher_rest::profile::bound::resolve_profile( - instance, - dfe_fetcher::profiles::shipped(), - )?; - let shape = Shape::for_rest_instance( - &profile, - instance, - id, - http_client.clone(), - &exchange_client, - )?; - let cancel = shutdown_token.child_token(); - entries.push(SpawnEntry { - driver: Arc::new(framework_driver(shape, id, cancel.clone())), - interval_secs: instance.interval_secs, - cancel, - }); - } - - // Database instances: the same driver over the DB shape; every block a - // cursor fetches is leased on the memory guard until its rows are out. - for (id, instance) in &config.sources.db { - if !instance.enabled { - continue; - } - let lease: Arc = Arc::new( - dfe_fetcher::driver::GuardLease(Arc::clone(pipeline_state.memory_guard())), - ); - let shape = DbShape::from_instance(instance, id, &lease) - .map_err(|e| anyhow::anyhow!("sources.db.{id}: {e}"))?; - let cancel = shutdown_token.child_token(); - entries.push(SpawnEntry { - driver: Arc::new(framework_driver( - Shape::Db(Box::new(shape)), - id, - cancel.clone(), - )), - interval_secs: instance.interval_secs, - cancel, - }); - } - - // File instances: the same driver over the file shape; every chunk a - // reader holds and every pass the tailer hands over is leased likewise. - for (id, instance) in &config.sources.file { - if !instance.enabled { - continue; - } - let lease: Arc = Arc::new( - dfe_fetcher::driver::GuardLease(Arc::clone(pipeline_state.memory_guard())), - ); - let shape = FileShape::from_instance(instance, id, &lease) - .map_err(|e| anyhow::anyhow!("sources.file.{id}: {e}"))?; - let cancel = shutdown_token.child_token(); - entries.push(SpawnEntry { - driver: Arc::new(framework_driver( - Shape::File(Box::new(shape)), - id, - cancel.clone(), - )), - interval_secs: instance.interval_secs, - cancel, - }); - } - - // Start one fetch task per connection. The entries are consumed so the - // fetch task holds the only `Arc`: a cancelled task then drops its - // driver and releases the connection's credentials and client. - let mut running: HashMap = HashMap::new(); - let mut drivers: Vec> = Vec::new(); - for entry in entries { - // Surface non-stable sources at startup: the profile declares the - // maturity, and anything not stable is code-complete but not - // production-validated until promoted. - let maturity = entry.driver.maturity(); - if maturity != SourceMaturity::Stable { - warn!( - source = entry.driver.name(), - maturity = %maturity, - "source is {maturity} maturity - not production-validated; \ - behaviour and config may change. See docs/cloud-setup/{}.md", - entry.driver.name(), - ); - } - - let initial_interval = scheduler.effective_interval(entry.interval_secs); - let ready_state = Arc::clone(&pipeline_state); - info!( - source = entry.driver.name(), - interval_secs = initial_interval.as_secs(), - "Starting fetch schedule (interval is hot-reloaded)" - ); - - running.insert( - entry.driver.connection_id().to_owned(), - entry.cancel.clone(), - ); - drivers.push(Arc::clone(&entry.driver)); - scheduler.spawn_source_task( - entry.driver, - entry.interval_secs, - Arc::clone(&metrics), - entry.cancel, - Arc::new(move || ready_state.is_ready()), - ); - } - - // A scheduled source's cursor advances only once its records are delivered. - if !running.is_empty() { - pipeline_state.publish_guarantee( - "scheduled", - Some(&IntakeAcks { - kind: scalo::transport::AckKind::Pull, - enabled: true, - }), - ); - } - - // Diff the running fetch tasks against every reload, so a source dropped - // from the config stops fetching without a restart. - scheduler.spawn_source_watch(running, shutdown_token.clone()); + let ready_state = Arc::clone(&pipeline_state); + let sources = scheduler.start_sources( + factory, + &config, + Arc::clone(&metrics), + Arc::new(move || ready_state.is_ready()), + shutdown_token.clone(), + )?; // Start ingest HTTP server (for container extractors using HTTP // communication). A listener that cannot start leaves the intake dead @@ -752,18 +727,10 @@ async fn run_fetcher_service( // Run pipeline orchestrator (blocks until shutdown) let run = orchestrator.run().await; - // The fetch tasks end on the same shutdown, but nothing waits for them, so - // a file tailer's final checkpoint is written here, before the runtime - // goes. A tick that will not release the tailer cannot hold the exit. - if tokio::time::timeout( - DRIVER_STOP_TIMEOUT, - futures::future::join_all(drivers.iter().map(|d| d.stop())), - ) - .await - .is_err() - { - warn!("a source driver did not stop in time; its tailer's last checkpoint may be missing"); - } + // The fetch tasks end on the same shutdown. Waiting for them, and for + // every driver still stopping, writes each file tailer's final checkpoint + // before the runtime goes. + sources.shutdown().await; // Before any exit: process::exit skips drops, leaving the managed Vector // containers running and their inline configs on disk. @@ -1002,4 +969,54 @@ mod tests { Some("no enabled sources, container extractors, vector receiver or ingest listener") ); } + + /// A source's fingerprint moves with any setting its driver is built + /// from, a secret included, and with nothing else: the reload rebuilds + /// exactly the sources a restart would build differently. + #[test] + fn a_source_fingerprint_moves_with_its_own_settings_only() { + let prints = |config: &Config| -> std::collections::BTreeMap { + plan_sources(config) + .expect("plans") + .into_iter() + .map(|p| (p.id, p.fingerprint)) + .collect() + }; + let mut config = Config::default(); + config + .sources + .rest + .insert("api".into(), RestInstance::default()); + let db = DbInstance { + enabled: true, + ..DbInstance::default() + }; + config.sources.db.insert("pg".into(), db); + let base = prints(&config); + assert_eq!(base.keys().collect::>(), ["api", "pg"]); + assert_eq!(prints(&config), base, "the same config plans the same"); + + let mut unrelated = config.clone(); + unrelated.scheduler.default_interval_secs += 1; + assert_eq!(prints(&unrelated), base, "a global edit rebuilds nothing"); + + let mut secret = config.clone(); + if let Some(pg) = secret.sources.db.get_mut("pg") { + pg.connection_string = scalo::SensitiveString::from("postgres://elsewhere"); + } + let moved = prints(&secret); + assert_eq!(moved["api"], base["api"], "a sibling is untouched"); + assert_ne!( + moved["pg"], base["pg"], + "a changed secret rebuilds its source" + ); + + let mut in_flight = config.clone(); + in_flight.accumulate.in_flight += 1; + let moved = prints(&in_flight); + assert!( + base.iter().all(|(id, print)| moved[id] != *print), + "a setting every driver bakes in rebuilds them all" + ); + } } diff --git a/crates/fetcher/src/pipeline/mod.rs b/crates/fetcher/src/pipeline/mod.rs index 5307c61..09c869c 100644 --- a/crates/fetcher/src/pipeline/mod.rs +++ b/crates/fetcher/src/pipeline/mod.rs @@ -926,18 +926,6 @@ impl PipelineState { self.memory_guard.limit_bytes(), ); } - - /// Reload configuration. - /// - /// # Errors - /// - /// Returns an error when the reload cannot be applied. - pub fn reload_config(&self, new_config: Config) -> Result<()> { - self.shared_config.update(new_config); - let version = self.shared_config.version(); - info!(version, "Configuration reloaded successfully"); - Ok(()) - } } /// Main pipeline orchestrator. @@ -1525,36 +1513,6 @@ mod tests { assert!(!state.is_ready(), "and must stall the fetch loop"); } - // -- reload_config tests -- - - #[test] - fn test_reload_config_updates_and_increments_version() { - let config = Config::default(); - let shared = SharedConfig::new(config); - let initial_version = shared.version(); - let metrics = Arc::new(Metrics::new()); - let state = PipelineState::for_tests(shared, metrics, None); - - let mut new_config = Config::default(); - new_config.scheduler.default_interval_secs = 999; - state - .reload_config(new_config) - .expect("reload_config should succeed"); - - // Version incremented - assert_eq!( - state.shared_config().version(), - initial_version + 1, - "version should increment on reload" - ); - // Config content applied - assert_eq!( - state.config().scheduler.default_interval_secs, - 999, - "reloaded config should be returned by config()" - ); - } - // -- Orchestrator tests -- /// The default config with the DLQ off: its default spool path is the diff --git a/crates/fetcher/src/scheduler/mod.rs b/crates/fetcher/src/scheduler/mod.rs index 84460b2..0e98756 100644 --- a/crates/fetcher/src/scheduler/mod.rs +++ b/crates/fetcher/src/scheduler/mod.rs @@ -23,14 +23,152 @@ use std::time::Duration; use chrono::Utc; use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; use tracing::{debug, error, info, warn}; -use crate::config::{MissingCursor, SchedulerConfig, SharedConfig}; +use crate::config::resolve::resolve_config_specs; +use crate::config::{Config, MissingCursor, SchedulerConfig, SharedConfig}; +use crate::credential::CredentialError; use crate::cursor::{CursorStore, CursorValue}; use crate::driver::Driver; use crate::error::Error; use crate::metrics::Metrics; -use dfe_fetcher_core::FetchWindow; +use dfe_fetcher_core::{FetchWindow, SourceMaturity}; + +/// How long a driver gets to stop what it runs: at shutdown, and before a +/// rebuilt source's new driver starts, so a file tailer's last checkpoint is +/// written first. +pub const DRIVER_STOP_TIMEOUT: Duration = Duration::from_secs(10); + +/// How soon a sync that did not settle is tried again: a spec that did not +/// resolve, or a generation not yet confirmed stopped. +const SYNC_RETRY: Duration = Duration::from_secs(30); + +/// The pipeline readiness a fetch task polls before each tick. +pub type ReadyFn = Arc bool + Send + Sync>; + +/// One connection a config schedules, before its driver is built. +pub struct Planned { + /// The connection id: the cursor key and the metric and log label. + pub id: String, + /// A digest of everything the driver is built from, so a reload that + /// changes any of it rebuilds the source. + pub fingerprint: u64, + /// The per-source interval override. + pub interval_secs: Option, + /// What the factory builds the driver from. + pub spec: S, +} + +/// The digest [`Planned::fingerprint`] holds: `value` serialised with its +/// secrets exposed, so a changed secret spec counts as a change, then hashed +/// so no plaintext outlives the call. +/// +/// # Errors +/// +/// Returns the serialiser's error. +pub fn fingerprint(value: &T) -> serde_json::Result { + use std::hash::{Hash, Hasher}; + let bytes = scalo::expose_during(|| serde_json::to_vec(value))?; + let mut hasher = std::hash::DefaultHasher::new(); + bytes.hash(&mut hasher); + Ok(hasher.finish()) +} + +/// Expands a config into the connections it schedules, and builds each +/// one's driver. +pub trait SourceFactory: Send + Sync + 'static { + /// What one connection's driver is built from. + type Spec: Send; + + /// Every connection `config` schedules. + /// + /// # Errors + /// + /// Returns an error when the config's sources do not expand. + fn plan(&self, config: &Config) -> anyhow::Result>>; + + /// Build the driver of connection `id`, which shuts down on `cancel`. + /// + /// # Errors + /// + /// Returns an error when the connection's settings do not build. + fn build( + &self, + id: &str, + spec: Self::Spec, + cancel: CancellationToken, + ) -> anyhow::Result; +} + +/// One built source, and its fetch task once launched. +struct Running { + driver: Arc, + fingerprint: u64, + interval_secs: Option, + /// A child of the global shutdown token, held by both the driver and the + /// fetch task, so stopping one source leaves the rest running. + cancel: CancellationToken, + task: Option>, +} + +/// A cancelled generation whose driver is not yet confirmed stopped. No +/// other generation of its source starts until it is. +struct Stopping { + driver: Arc, + /// `None` once the task has been joined, or while a settle awaits it. + task: Option>, +} + +/// Every generation of every source, keyed on connection id. +#[derive(Default)] +struct Generations { + running: HashMap, + stopping: HashMap, +} + +/// The sources [`Scheduler::start_sources`] keeps equal to the live config, +/// and the generations still stopping, which shutdown waits for too. +#[derive(Clone)] +pub struct RunningSources { + generations: Arc>, + tasks: TaskTracker, +} + +impl RunningSources { + /// The driver of every running and stopping generation. + #[must_use] + pub fn drivers(&self) -> Vec> { + let generations = self.generations.lock(); + generations + .running + .values() + .map(|r| Arc::clone(&r.driver)) + .chain(generations.stopping.values().map(|s| Arc::clone(&s.driver))) + .collect() + } + + /// Wait for every fetch task to end, then stop every driver still held, + /// both within one [`DRIVER_STOP_TIMEOUT`]. Call once the shutdown token + /// is cancelled. + pub async fn shutdown(&self) { + let deadline = tokio::time::Instant::now() + DRIVER_STOP_TIMEOUT; + self.tasks.close(); + if tokio::time::timeout_at(deadline, self.tasks.wait()) + .await + .is_err() + { + warn!("a fetch task did not end in time for shutdown"); + } + let drivers = self.drivers(); + let stops = futures::future::join_all(drivers.iter().map(|d| d.stop())); + if tokio::time::timeout_at(deadline, stops).await.is_err() { + warn!( + "a source driver did not stop in time, so its tailer's last checkpoint may be missing" + ); + } + } +} /// Fetch scheduler that coordinates timing and concurrency. /// @@ -42,11 +180,13 @@ use dfe_fetcher_core::FetchWindow; /// that fails is retried over the width that failed, so a source catching up /// after an outage drains in bounded steps instead of asking for everything /// at once. +#[derive(Clone)] pub struct Scheduler { shared_config: SharedConfig, concurrency_semaphore: Arc, cursor_store: Option>, instance_id: String, + tasks: TaskTracker, } impl Scheduler { @@ -64,6 +204,7 @@ impl Scheduler { concurrency_semaphore: semaphore, cursor_store, instance_id, + tasks: TaskTracker::new(), } } @@ -81,6 +222,9 @@ impl Scheduler { /// When a cursor store is configured, the scheduler reads the last fetch /// position before each fetch to compute a `FetchWindow`, and writes the /// cursor back once the tick has awaited every acknowledgement. + /// + /// The task stops the driver before it ends, so joining the returned + /// handle waits for the driver's last checkpoint. pub fn spawn_source_task( &self, driver: Arc, @@ -88,16 +232,14 @@ impl Scheduler { metrics: Arc, shutdown: CancellationToken, is_ready: Arc bool + Send + Sync>, - ) { + ) -> tokio::task::JoinHandle<()> { let semaphore = Arc::clone(&self.concurrency_semaphore); let cursor_store = self.cursor_store.clone(); let instance_id = self.instance_id.clone(); let shared_config = self.shared_config.clone(); - tokio::spawn(async move { - // Whichever way this task ends (shutdown, a reload that dropped - // the source, a closed semaphore), the driver's tailer is stopped. - let _stop = StopDriverOnDrop(Arc::clone(&driver)); + let stopping = Arc::clone(&driver); + let fetch = async move { // The cursor key and the per-source metric and log labels are keyed // on the CONNECTION id, not the type name, so each account or // tenant of a type checkpoints and reports independently; a @@ -124,8 +266,8 @@ impl Scheduler { } } - // The last failed attempt's window. A reload never respawns a live - // source's task, so a frozen window survives one. + // The last failed attempt's window. A reload respawns only a + // source whose settings changed, so a frozen window survives one. let mut retry: Option = None; loop { @@ -286,44 +428,61 @@ impl Scheduler { break; } } - }); + }; + self.tasks.spawn(async move { + Box::pin(fetch).await; + // However the loop ended, the tailer has stopped before the task + // does, so a caller that joins the task waits for its last checkpoint. + stopping.stop().await; + }) } - /// Watch the shared config and cancel the fetch task of any source a - /// reload has removed. + /// Start a fetch task for every connection `config` schedules, then keep + /// the running set equal to the live config's on every reload, as a + /// restart would: an added source starts, a removed one stops, and one + /// whose settings changed is rebuilt from the new settings. + /// + /// `config` is the startup config with its specs resolved. Each reload's + /// specs are resolved the same way before anything is built from it. + /// + /// # Errors /// - /// `running` maps a connection id to the child token its fetch task and - /// its driver hold, and the watcher is its only writer -- no lock. The - /// diff is on IDENTITY alone: a source whose settings changed but whose id - /// is still configured keeps running, because the driver and the scheduler - /// re-read the config every tick. - pub fn spawn_source_watch( + /// Returns the first connection that does not plan or build: at startup + /// a source that cannot run refuses the process. + pub fn start_sources( &self, - mut running: HashMap, + factory: F, + config: &Config, + metrics: Arc, + is_ready: ReadyFn, shutdown: CancellationToken, - ) { - let shared_config = self.shared_config.clone(); - - tokio::spawn(async move { - let mut versions = shared_config.subscribe(); - - // Reconcile once up front: a reload between the startup spawn and - // this subscription would otherwise go unseen. - reconcile_running_sources(&shared_config, &mut running); - - loop { - tokio::select! { - biased; - () = shutdown.cancelled() => return, - changed = versions.changed() => { - if changed.is_err() { - return; - } - } - } - reconcile_running_sources(&shared_config, &mut running); - } - }); + ) -> anyhow::Result { + let set = SourceSet { + scheduler: self.clone(), + factory, + metrics, + is_ready, + shutdown, + running: RunningSources { + generations: Arc::default(), + tasks: self.tasks.clone(), + }, + }; + // The receiver reports every update since the config was created, so + // a reload that landed before this call is synced first. + let versions = self.shared_config.subscribe(); + let built: anyhow::Result> = set + .factory + .plan(config)? + .into_iter() + .map(|planned| set.build(planned)) + .collect(); + for running in built? { + set.launch(running); + } + let running = set.running.clone(); + tokio::spawn(set.watch(versions)); + Ok(running) } /// Calculate the effective interval for a source (for logging at startup). @@ -338,55 +497,245 @@ impl Scheduler { } } -/// Stops a driver's background work when the fetch task holding it ends. -struct StopDriverOnDrop(Arc); +/// The fetch tasks of every scheduled source, and what they are built with. +struct SourceSet { + scheduler: Scheduler, + factory: F, + metrics: Arc, + is_ready: ReadyFn, + shutdown: CancellationToken, + running: RunningSources, +} -impl Drop for StopDriverOnDrop { - fn drop(&mut self) { - let driver = Arc::clone(&self.0); - if let Ok(runtime) = tokio::runtime::Handle::try_current() { - runtime.spawn(async move { driver.stop().await }); +impl SourceSet { + /// Build `planned`'s driver over a fresh child of the shutdown token. + fn build(&self, planned: Planned) -> anyhow::Result<(String, Running)> { + let cancel = self.shutdown.child_token(); + let driver = self + .factory + .build(&planned.id, planned.spec, cancel.clone())?; + let running = Running { + driver: Arc::new(driver), + fingerprint: planned.fingerprint, + interval_secs: planned.interval_secs, + cancel, + task: None, + }; + Ok((planned.id, running)) + } + + /// Start a built source's fetch task and record it as running. + fn launch(&self, (id, mut running): (String, Running)) { + let driver = Arc::clone(&running.driver); + // Not stable means code-complete but not production-validated. + let maturity = driver.maturity(); + if maturity != SourceMaturity::Stable { + warn!( + source = driver.name(), + maturity = %maturity, + "source is {maturity} maturity - not production-validated; \ + behaviour and config may change. See docs/cloud-setup/{}.md", + driver.name(), + ); } + info!( + source = driver.name(), + interval_secs = self + .scheduler + .effective_interval(running.interval_secs) + .as_secs(), + "Starting fetch schedule (interval is hot-reloaded)" + ); + running.task = Some(self.scheduler.spawn_source_task( + driver, + running.interval_secs, + Arc::clone(&self.metrics), + running.cancel.clone(), + Arc::clone(&self.is_ready), + )); + self.running.generations.lock().running.insert(id, running); } -} -/// Cancel the fetch task of every running source the live config no longer -/// schedules, and warn about a source it schedules that is not running. -/// -/// Cancelling the child token stops the task at its next select and aborts an -/// in-flight tick through the driver's non-alerting shutdown path. Adding a -/// source still needs a restart (dfe-fetcher#134), so its id is named in a -/// warning rather than silently ignored. -fn reconcile_running_sources( - shared_config: &SharedConfig, - running: &mut HashMap, -) { - let desired: BTreeSet = shared_config.with(|config| { - config - .sources - .scheduled_connection_ids() - .into_iter() - .map(str::to_owned) - .collect() - }); + /// Cancel `id`'s running generation and move it to `stopping`, where + /// shutdown still sees it, then settle it. Returns `false` while it is + /// not confirmed stopped. + async fn retire(&self, id: &str) -> bool { + { + let mut generations = self.running.generations.lock(); + let Some(running) = generations.running.remove(id) else { + return true; + }; + // Cancelling the child token stops the task at its next select and + // aborts an in-flight tick through the driver's non-alerting path. + running.cancel.cancel(); + let stopping = Stopping { + driver: running.driver, + task: running.task, + }; + generations.stopping.insert(id.to_owned(), stopping); + } + self.settle(id).await + } - running.retain(|id, cancel| { - if desired.contains(id) { + /// Confirm `id`'s stopping generation is down: its task ended, then its + /// driver stopped, each within [`DRIVER_STOP_TIMEOUT`]. Anything not + /// confirmed keeps the generation in `stopping` for a later sync. + async fn settle(&self, id: &str) -> bool { + let task = match self.running.generations.lock().stopping.get_mut(id) { + None => return true, + Some(stopping) => stopping.task.take(), + }; + if let Some(mut task) = task { + match tokio::time::timeout(DRIVER_STOP_TIMEOUT, &mut task).await { + Err(_) => { + warn!(source = %id, "the previous fetch task did not end in time, so the source stays stopped until it does"); + if let Some(stopping) = self.running.generations.lock().stopping.get_mut(id) { + stopping.task = Some(task); + } + return false; + } + // A task that panicked never stopped its driver, which the + // stop below does instead. + Ok(Err(e)) => warn!(source = %id, error = %e, "the previous fetch task failed"), + Ok(Ok(())) => {} + } + } + let driver = self + .running + .generations + .lock() + .stopping + .get(id) + .map(|stopping| Arc::clone(&stopping.driver)); + let Some(driver) = driver else { return true; + }; + if tokio::time::timeout(DRIVER_STOP_TIMEOUT, driver.stop()) + .await + .is_err() + { + warn!(source = %id, "the previous driver did not stop in time, so the source stays stopped until it does"); + return false; } - info!(source = %id, "source removed from the config; stopping its fetch schedule"); - cancel.cancel(); - false - }); + self.running.generations.lock().stopping.remove(id); + true + } - for id in desired - .iter() - .filter(|id| !running.contains_key(id.as_str())) - { - warn!( - source = %id, - "source added to the config; a restart is needed to start fetching it" - ); + /// Whether no earlier generation of `id` is still stopping. One whose + /// task has ended is settled here. + async fn settled(&self, id: &str) -> bool { + let ending = match self.running.generations.lock().stopping.get(id) { + None => return true, + Some(stopping) => stopping.task.as_ref().is_some_and(|t| !t.is_finished()), + }; + !ending && self.settle(id).await + } + + /// Bring the running set to `planned`, one generation of a source at a + /// time: a replaced or removed generation is confirmed stopped before + /// another starts. A source whose driver does not build is logged and + /// left as it was. + async fn reconcile(&self, planned: Vec>) { + let stopping: Vec = { + let generations = self.running.generations.lock(); + generations.stopping.keys().cloned().collect() + }; + for id in stopping { + self.settled(&id).await; + } + + let desired: BTreeSet<&str> = planned.iter().map(|p| p.id.as_str()).collect(); + let removed: Vec = self + .running + .generations + .lock() + .running + .keys() + .filter(|id| !desired.contains(id.as_str())) + .cloned() + .collect(); + for id in removed { + info!(source = %id, "source removed from the config, stopping its fetch schedule"); + self.retire(&id).await; + } + + for planned in planned { + let previous = self + .running + .generations + .lock() + .running + .get(&planned.id) + .map(|r| r.fingerprint); + let change = match previous { + Some(fingerprint) if fingerprint == planned.fingerprint => continue, + Some(_) => "changed", + None => "added", + }; + let id = planned.id.clone(); + if !self.settled(&id).await { + warn!(source = %id, "source {change} in the config, but its previous generation is still stopping, so it is not started yet"); + continue; + } + let built = match self.build(planned) { + Ok(built) => built, + Err(e) => { + error!( + source = %id, + error = %e, + "source {change} in the config does not build, so the running set keeps its previous state" + ); + continue; + } + }; + if !self.retire(&id).await { + continue; + } + info!(source = %id, "source {change} in the config, starting it from the new settings"); + self.launch(built); + } + } + + /// Bring the running set to the live config's sources, resolved as at + /// startup. + async fn sync(&self) -> Result<(), CredentialError> { + let mut config = self.scheduler.shared_config.get(); + resolve_config_specs(&mut config).await?; + match self.factory.plan(&config) { + Ok(planned) => self.reconcile(planned).await, + Err(e) => { + error!(error = %e, "the reloaded sources do not expand, so the running set is unchanged"); + } + } + Ok(()) + } + + /// Sync on every reload until shutdown. A spec that does not resolve (a + /// secrets backend outage), or a generation still stopping, is retried + /// without waiting for a reload. + async fn watch(self, mut versions: tokio::sync::watch::Receiver) { + let mut retry = false; + loop { + tokio::select! { + biased; + () = self.shutdown.cancelled() => return, + changed = versions.changed() => { + if changed.is_err() { + return; + } + } + () = tokio::time::sleep(SYNC_RETRY), if retry => {} + } + let synced = self.sync().await; + retry = synced.is_err() || !self.running.generations.lock().stopping.is_empty(); + if let Err(e) = synced { + error!( + error = %e, + retry_secs = SYNC_RETRY.as_secs(), + "a reloaded source spec does not resolve, so the running set is unchanged" + ); + } + } } } @@ -862,12 +1211,29 @@ mod tests { connection_id: &str, source: Box, ) -> Arc { + Arc::new(driver_with( + shared, + metrics, + connection_id, + source, + CancellationToken::new(), + )) + } + + /// [`driver_over`] shutting down on `cancel`. + fn driver_with( + shared: &SharedConfig, + metrics: &Arc, + connection_id: &str, + source: Box, + cancel: CancellationToken, + ) -> Driver { let state = Arc::new(PipelineState::for_tests( shared.clone(), Arc::clone(metrics), None, )); - Arc::new(Driver::new(DriverParts { + Driver::new(DriverParts { shape: Shape::Custom(source), connection_id: connection_id.to_owned(), instance_id: "inst".into(), @@ -877,8 +1243,8 @@ mod tests { memory_guard: Arc::clone(state.memory_guard()), checkpoints: None, metrics: Arc::clone(metrics), - shutdown: CancellationToken::new(), - })) + shutdown: cancel, + }) } /// A driver over [`Counting`] under `connection_id`. @@ -1689,95 +2055,459 @@ mod tests { shutdown.cancel(); } - /// A reload that removes the source must stop its fetch schedule: the - /// watcher cancels the source's child token, the task returns and the - /// driver it held is freed. + /// One build a [`CountingFactory`] made: the connection id and its ticks. + type Build = (String, Arc); + + /// Where every [`Generation`] logs its ticks and its completed stops. A + /// stop waits for a permit while `hold` is set. + struct StopGate { + hold: AtomicBool, + gate: tokio::sync::Semaphore, + log: Mutex>, + } + + impl Default for StopGate { + fn default() -> Self { + Self { + hold: AtomicBool::new(false), + gate: tokio::sync::Semaphore::new(0), + log: Mutex::default(), + } + } + } + + impl StopGate { + /// Let the held stop finish, holding none after it. + fn release(&self) { + self.hold.store(false, Ordering::SeqCst); + self.gate.add_permits(1); + } + + fn position(&self, entry: &str) -> Option { + self.log.lock().unwrap().iter().position(|e| e == entry) + } + } + + /// A shape that counts its ticks and logs them, and its stops, to a + /// [`StopGate`] as build number `build`. + struct Generation { + units: Vec, + ticks: Arc, + build: usize, + stops: Arc, + } + + impl RowSource for Generation { + fn name(&self) -> &'static str { + "generation" + } + fn maturity(&self) -> SourceMaturity { + SourceMaturity::Alpha + } + fn units(&self) -> &[UnitSpec] { + &self.units + } + fn rows<'a>(&'a self, _tick: TickCtx<'a>) -> RowStream<'a> { + self.ticks.fetch_add(1, Ordering::Relaxed); + let entry = format!("tick {}", self.build); + self.stops.log.lock().unwrap().push(entry); + futures::stream::empty().boxed() + } + fn probe(&self) -> BoxFuture<'_, dfe_fetcher_core::Result<()>> { + Box::pin(async { Ok(()) }) + } + fn stop(&self) -> BoxFuture<'_, ()> { + Box::pin(async move { + if self.stops.hold.load(Ordering::SeqCst) { + self.stops.gate.acquire().await.unwrap().forget(); + } + let entry = format!("stopped {}", self.build); + self.stops.log.lock().unwrap().push(entry); + }) + } + } + + /// A factory whose connections are the enabled `sources.rest` entries, + /// each a [`Generation`] driver. It records each build's id and tick + /// count, and the AWS region every plan was handed. + #[derive(Clone)] + struct CountingFactory { + shared: SharedConfig, + metrics: Arc, + builds: Arc>>, + regions: Arc>>, + stops: Arc, + } + + impl SourceFactory for CountingFactory { + type Spec = (); + + fn plan(&self, config: &Config) -> anyhow::Result>> { + let region = config.sources.aws.region.clone(); + self.regions.lock().unwrap().push(region); + config + .sources + .rest + .iter() + .filter(|(_, instance)| instance.enabled) + .map(|(id, instance)| { + Ok(Planned { + id: id.clone(), + fingerprint: fingerprint(instance)?, + interval_secs: instance.interval_secs, + spec: (), + }) + }) + .collect() + } + + fn build(&self, id: &str, (): (), cancel: CancellationToken) -> anyhow::Result { + let ticks = Arc::new(AtomicU64::new(0)); + let mut builds = self.builds.lock().unwrap(); + builds.push((id.to_owned(), Arc::clone(&ticks))); + let source = Box::new(Generation { + units: vec![UnitSpec::new("unit", UnitShape::Incremental, "t")], + ticks, + build: builds.len() - 1, + stops: Arc::clone(&self.stops), + }); + Ok(driver_with(&self.shared, &self.metrics, id, source, cancel)) + } + } + + /// Advance the paused clock `secs` seconds, a second at a time, letting + /// every task run between steps. + async fn run_for(secs: u64) { + for _ in 0..secs { + tokio::time::advance(Duration::from_secs(1)).await; + for _ in 0..4 { + tokio::task::yield_now().await; + } + } + } + + /// A reload has the effect a restart on the reloaded config would: an + /// added source starts, a removed one stops and frees its driver, a + /// changed one is rebuilt from its new settings, and an unchanged one + /// keeps its task. Every reload is planned from its specs resolved. #[tokio::test] - async fn a_source_dropped_by_a_reload_stops_ticking() { + #[allow(unsafe_code)] + async fn a_reload_starts_stops_and_rebuilds_sources_as_a_restart_would() { tokio::time::pause(); + // SAFETY: test-only; a var name no other test sets. + unsafe { std::env::set_var("DFE_FETCHER_TEST_RELOAD_REGION", "eu-west-1") }; - // `scheduled_connection_ids` reports "counting" only while the - // instance is in the config; the post-reload config keeps the same - // cadence so the tick count can only change by the cancel. - let interval_secs = 2; - let with_source = || { - let mut cfg = test_config_no_jitter(); - cfg.scheduler.default_interval_secs = interval_secs; - cfg + let source = dfe_fetcher_rest::RestInstance::default(); + let changed = dfe_fetcher_rest::RestInstance { + interval_secs: Some(3), + ..Default::default() }; - let mut cfg = with_source(); - cfg.sources - .rest - .insert("counting".into(), dfe_fetcher_rest::RestInstance::default()); - let shared = SharedConfig::new(cfg); + // (row, before, after, builds, whether the first driver survives) + let rows = [ + ("added", None, Some(&source), 1, true), + ("removed", Some(&source), None, 1, false), + ("changed", Some(&source), Some(&changed), 2, false), + ("unchanged", Some(&source), Some(&source), 1, true), + ]; + + for (row, before, after, want_builds, first_survives) in rows { + let config_with = |instance: Option<&dfe_fetcher_rest::RestInstance>| { + let mut cfg = test_config_no_jitter(); + cfg.scheduler.default_interval_secs = 2; + if let Some(instance) = instance { + cfg.sources.rest.insert("counting".into(), instance.clone()); + } + cfg + }; + let startup = config_with(before); + let shared = SharedConfig::new(startup.clone()); + let scheduler = Scheduler::new(&startup.scheduler, shared.clone(), None, "test".into()); + let metrics = Arc::new(Metrics::new()); + let factory = CountingFactory { + shared: shared.clone(), + metrics: Arc::clone(&metrics), + builds: Arc::default(), + regions: Arc::default(), + stops: Arc::default(), + }; + let shutdown = CancellationToken::new(); + let running = scheduler + .start_sources( + factory.clone(), + &startup, + Arc::clone(&metrics), + Arc::new(|| true), + shutdown.clone(), + ) + .expect("startup builds"); - let scheduler_config = SchedulerConfig { - default_interval_secs: interval_secs, - max_concurrent_fetches: 10, - jitter_percent: 0, - }; - let scheduler = Scheduler::new(&scheduler_config, shared.clone(), None, "test".into()); + run_for(5).await; + let first: Vec> = + running.drivers().iter().map(Arc::downgrade).collect(); + let first_ticks: Vec = factory + .builds + .lock() + .unwrap() + .iter() + .map(|(_, ticks)| ticks.load(Ordering::Relaxed)) + .collect(); + assert!( + first_ticks.iter().all(|&ticks| ticks >= 1), + "{row}: a configured source ticks, got {first_ticks:?}" + ); - let metrics = Arc::new(Metrics::new()); - let ticks = Arc::new(AtomicU64::new(0)); - let driver = counting_driver(&shared, &metrics, "counting", Arc::clone(&ticks)); - assert_eq!(driver.connection_id(), "counting"); - let driver_alive = Arc::downgrade(&driver); + // An unrelated edit with a spec in it, which the plan must see resolved. + let mut reloaded = config_with(after); + reloaded.sources.aws.enabled = true; + reloaded.sources.aws.region = "env:DFE_FETCHER_TEST_RELOAD_REGION".into(); + shared.update(reloaded); + run_for(10).await; + + let builds = factory.builds.lock().unwrap(); + assert_eq!(builds.len(), want_builds, "{row}: builds"); + let want_ids: BTreeSet = + after.map(|_| "counting".to_string()).into_iter().collect(); + let ids: BTreeSet = + running.generations.lock().running.keys().cloned().collect(); + assert_eq!(ids, want_ids, "{row}: running sources"); + if after.is_some() { + let (_, live) = builds.last().expect("a build"); + assert!( + live.load(Ordering::Relaxed) >= 1, + "{row}: the source the reload leaves configured ticks" + ); + } + for ((_, ticks), before_reload) in builds.iter().zip(&first_ticks) { + if !first_survives { + assert_eq!( + ticks.load(Ordering::Relaxed), + *before_reload, + "{row}: a stopped driver must not tick again" + ); + } + } + assert_eq!( + first.iter().all(|driver| driver.upgrade().is_some()), + first_survives, + "{row}: whether the startup driver is still held" + ); + assert_eq!( + factory.regions.lock().unwrap().last().map(String::as_str), + Some("eu-west-1"), + "{row}: the reload is planned from its resolved specs" + ); + assert!( + !shutdown.is_cancelled(), + "{row}: a reload never shuts the process down" + ); + shutdown.cancel(); + } - let shutdown = CancellationToken::new(); - let cancel = shutdown.child_token(); - scheduler.spawn_source_task( - driver, - None, - Arc::clone(&metrics), - cancel.clone(), - Arc::new(|| true), - ); - scheduler.spawn_source_watch( - HashMap::from([("counting".to_string(), cancel.clone())]), - shutdown.clone(), - ); + // SAFETY: test-only, as above. + unsafe { std::env::remove_var("DFE_FETCHER_TEST_RELOAD_REGION") }; + } - for _ in 0..10 { - tokio::time::advance(Duration::from_secs(1)).await; - tokio::task::yield_now().await; - } - let ticks_while_live = ticks.load(Ordering::Relaxed); - assert!( - ticks_while_live >= 1, - "the source should tick while it is configured, got {ticks_while_live}" - ); + /// A cursor store that records the end of every cursor written, and holds + /// a write at its start while `hold` is set, until a permit is added. + struct GatedStore { + hold: AtomicBool, + gate: tokio::sync::Semaphore, + held: AtomicU64, + ends: Mutex>>, + cursor: Mutex>, + } - // Drop the source: the unit-level equivalent of a rewritten - // fetcher.yaml that no longer lists it. - shared.update(with_source()); + impl GatedStore { + /// Let the held write finish, holding none after it. + fn release(&self) { + self.hold.store(false, Ordering::SeqCst); + self.gate.add_permits(1); + } - // No time passes here, so nothing but the cancel can end the cadence. - for _ in 0..10 { - tokio::task::yield_now().await; + fn held(&self) -> u64 { + self.held.load(Ordering::SeqCst) } - assert!( - cancel.is_cancelled(), - "the watcher should cancel the removed source" - ); + } - for _ in 0..20 { - tokio::time::advance(Duration::from_secs(1)).await; - tokio::task::yield_now().await; + #[async_trait::async_trait] + impl CursorStore for GatedStore { + async fn get(&self, _: &str) -> dfe_fetcher_core::error::Result> { + Ok(self.cursor.lock().unwrap().clone()) } - assert_eq!( - ticks.load(Ordering::Relaxed), - ticks_while_live, - "a removed source must not tick again" + async fn set(&self, _: &str, value: &CursorValue) -> dfe_fetcher_core::error::Result<()> { + if self.hold.load(Ordering::SeqCst) { + self.held.fetch_add(1, Ordering::SeqCst); + self.gate.acquire().await.unwrap().forget(); + self.held.fetch_sub(1, Ordering::SeqCst); + } + self.ends.lock().unwrap().push(value.last_fetch_end); + *self.cursor.lock().unwrap() = Some(value.clone()); + Ok(()) + } + async fn delete(&self, _: &str) -> dfe_fetcher_core::error::Result<()> { + Ok(()) + } + } + + /// A source runs one generation at a time. A generation caught committing + /// its cursor, outside any cancellation select, or stopping its driver, is + /// confirmed stopped before another starts: on a change, on a remove + /// followed by a re-add, and past the join timeout, where the retry timer + /// starts the source once the old one is down. The cursor never goes + /// backwards, and shutdown waits for the last driver to stop. + #[tokio::test] + async fn a_source_runs_one_generation_at_a_time_across_reloads() { + tokio::time::pause(); + let config_with = |interval_secs: Option>| { + let mut cfg = test_config_no_jitter(); + cfg.scheduler.default_interval_secs = 2; + if let Some(interval_secs) = interval_secs { + let instance = dfe_fetcher_rest::RestInstance { + interval_secs, + ..Default::default() + }; + cfg.sources.rest.insert("counting".into(), instance); + } + cfg + }; + let startup = config_with(Some(None)); + let shared = SharedConfig::new(startup.clone()); + let store = Arc::new(GatedStore { + hold: AtomicBool::new(true), + gate: tokio::sync::Semaphore::new(0), + held: AtomicU64::new(0), + ends: Mutex::default(), + cursor: Mutex::default(), + }); + let cursors: Arc = Arc::clone(&store) as Arc; + let scheduler = Scheduler::new( + &startup.scheduler, + shared.clone(), + Some(cursors), + "test".into(), ); + let metrics = Arc::new(Metrics::new()); + let factory = CountingFactory { + shared: shared.clone(), + metrics: Arc::clone(&metrics), + builds: Arc::default(), + regions: Arc::default(), + stops: Arc::default(), + }; + let shutdown = CancellationToken::new(); + let running = scheduler + .start_sources( + factory.clone(), + &startup, + metrics, + Arc::new(|| true), + shutdown.clone(), + ) + .expect("startup builds"); + let ticks = |build: usize| { + factory.builds.lock().unwrap()[build] + .1 + .load(Ordering::Relaxed) + }; + let builds = || factory.builds.lock().unwrap().len(); + let ids = || { + let ids: Vec = running.generations.lock().running.keys().cloned().collect(); + ids + }; + let hold_next_write = || async { + store.hold.store(true, Ordering::SeqCst); + for _ in 0..30 { + if store.held() == 1 { + return; + } + run_for(1).await; + } + panic!("no write was held"); + }; + + // A change while the first generation is committing its cursor. + hold_next_write().await; + shared.update(config_with(Some(Some(3)))); + run_for(5).await; + assert_eq!(builds(), 2, "the new generation is built"); + assert_eq!(ticks(1), 0, "and does not run while the old one commits"); + store.release(); + run_for(10).await; + assert!(ticks(1) >= 1, "it runs once the old generation has ended"); + + // A remove, then a re-add, while the second generation commits. + hold_next_write().await; + shared.update(config_with(None)); + run_for(3).await; + shared.update(config_with(Some(Some(3)))); + run_for(3).await; + assert_eq!(builds(), 2, "the re-add waits for the removed generation"); + store.release(); + run_for(10).await; + assert_eq!(builds(), 3, "then starts a third"); + assert!(ticks(2) >= 1, "which runs"); + + // A change whose old generation outlives the join timeout. + hold_next_write().await; + shared.update(config_with(Some(Some(4)))); + run_for(DRIVER_STOP_TIMEOUT.as_secs() + 2).await; + assert_eq!(builds(), 4, "the replacement is built"); + assert_eq!(ticks(3), 0, "but never started beside the old generation"); + assert!(ids().is_empty(), "so the source is stopped"); + store.release(); + run_for(3).await; + assert!(ids().is_empty(), "until the next retry"); + run_for(SYNC_RETRY.as_secs() + 10).await; + assert_eq!(ids(), ["counting"], "which starts it with no reload"); + assert_eq!(builds(), 5); + assert!(ticks(4) >= 1, "and it runs"); + + // A change whose old generation is held stopping its driver. + factory.stops.hold.store(true, Ordering::SeqCst); + shared.update(config_with(Some(Some(5)))); + run_for(5).await; + assert_eq!(builds(), 6, "the replacement is built"); + assert_eq!(ticks(5), 0, "and waits for the old driver to stop"); + assert_eq!(factory.stops.position("stopped 4"), None); + factory.stops.release(); + run_for(10).await; + let stopped = factory + .stops + .position("stopped 4") + .expect("the old driver stopped"); + let started = factory + .stops + .position("tick 5") + .expect("the replacement ticks"); assert!( - driver_alive.upgrade().is_none(), - "the cancelled task should have dropped its driver" + stopped < started, + "the old driver stopped before the replacement ran" ); + + let ends = store.ends.lock().unwrap().clone(); + assert!(ends.len() >= 4, "every generation committed, got {ends:?}"); assert!( - !shutdown.is_cancelled(), - "cancelling one source must not shut the process down" + ends.windows(2).all(|w| w[0] <= w[1]), + "the cursor never goes backwards: {ends:?}" ); + assert!(!shutdown.is_cancelled()); + + // Shutdown waits for a removed generation still stopping its driver. + factory.stops.hold.store(true, Ordering::SeqCst); + shared.update(config_with(None)); + run_for(DRIVER_STOP_TIMEOUT.as_secs() + 2).await; + assert!(ids().is_empty(), "the source is removed"); + shutdown.cancel(); + let drained = tokio::spawn({ + let running = running.clone(); + async move { running.shutdown().await } + }); + run_for(3).await; + assert!(!drained.is_finished(), "shutdown waits for the held stop"); + factory.stops.release(); + run_for(2).await; + assert!(drained.is_finished(), "and ends once it completes"); + assert!(factory.stops.position("stopped 5").is_some()); } /// Spawn a paced source over a cursor store seeded at `last_end` (empty diff --git a/docs/internals.md b/docs/internals.md index bc740d4..2b38a77 100644 --- a/docs/internals.md +++ b/docs/internals.md @@ -422,9 +422,7 @@ flowchart TB Earlier layers win. Validation binds every profile instance (a shipped profile by name or an inline one), every database instance and every file instance, so a bad field fails the load with its path rather than the first -tick. The per-source filter, the routes, the scheduler interval and jitter are -hot-reloaded and read once per tick; `enabled`, credentials and topics need a -restart. +tick. The per-source filter, the routes, the scheduler interval and jitter are read once per tick. A reload that adds, removes or changes a source starts, stops or rebuilds that source as a restart would, its `env:` and `vault:` specs resolved first. The ingest listener and the container and Vector extractors are built once, so a change to them needs a restart. ## Multi-Tenant Deployment