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