feat(topology): add opt-in reload sink drain limit

changelog.d/reload_sink_drain_limit.feature.md

@@ -0,0 +1,5 @@
+Add the `--reload-sink-drain-limit-secs` option (environment variable `VECTOR_RELOAD_SINK_DRAIN_LIMIT_SECS`). It sets how long removed and changed sinks can take to finish during a configuration reload, so that a sink that does not stop cannot block the reload. All of these sinks share one deadline. After the deadline, Vector cancels their in-flight work. A replacement sink that reuses the buffer receives the unread events. Other unread and in-flight events are lost.
+
+The option is not set by default, and Vector then waits until the sinks finish, as before. `--graceful-shutdown-limit-secs` does not apply to reloads.
+
+authors: thomasqueirozb

src/app.rs

@@ -75,6 +75,9 @@ impl ApplicationConfig {
 
         let graceful_shutdown_duration = (!opts.no_graceful_shutdown_limit)
             .then(|| Duration::from_secs(u64::from(opts.graceful_shutdown_limit_secs)));
+        let reload_sink_drain_duration = opts
+            .reload_sink_drain_limit_secs
+            .map(|secs| Duration::from_secs(u64::from(secs)));
 
         let watcher_conf = if opts.watch_config {
             Some(watcher_config(
@@ -91,6 +94,7 @@ impl ApplicationConfig {
             opts.require_healthy,
             opts.allow_empty_config,
             graceful_shutdown_duration,
+            reload_sink_drain_duration,
             signal_handler,
         )
         .await?;
@@ -650,6 +654,7 @@ pub async fn load_configs(
     require_healthy: Option<bool>,
     allow_empty_config: bool,
     graceful_shutdown_duration: Option<Duration>,
+    reload_sink_drain_duration: Option<Duration>,
     signal_handler: &mut SignalHandler,
 ) -> Result<Config, ExitCode> {
     let config_paths = config::process_paths(config_paths).ok_or(exitcode::CONFIG)?;
@@ -732,6 +737,7 @@ pub async fn load_configs(
     }
     config.healthchecks.set_require_healthy(require_healthy);
     config.graceful_shutdown_duration = graceful_shutdown_duration;
+    config.reload_sink_drain_duration = reload_sink_drain_duration;
 
     Ok(config)
 }

src/cli.rs

@@ -259,6 +259,15 @@ pub struct RootOpts {
     )]
     pub no_graceful_shutdown_limit: bool,
 
+    /// Set the duration in seconds that removed and changed sinks can take to finish during a
+    /// configuration reload. All of these sinks share one deadline, including sinks that finish
+    /// in the background. After the duration has passed, Vector cancels their in-flight work and
+    /// continues the reload. A replacement sink that reuses the buffer receives the unread
+    /// events. Other unread and in-flight events are lost. If not set, Vector waits until the
+    /// sinks finish.
+    #[arg(long, env = "VECTOR_RELOAD_SINK_DRAIN_LIMIT_SECS")]
+    pub reload_sink_drain_limit_secs: Option<NonZeroU64>,
+
     /// Set runtime allocation tracing
     #[cfg(all(unix, feature = "tikv-jemallocator"))]
     #[arg(long, env = "ALLOCATION_TRACING", default_value = "false")]

src/config/builder.rs

@@ -72,6 +72,13 @@ pub struct ConfigBuilder {
     #[doc(hidden)]
     pub graceful_shutdown_duration: Option<Duration>,
 
+    /// The duration to wait for removed and changed sinks to finish during a configuration reload.
+    /// No limit applies when unset. This value can be set using a
+    /// [cli arg](crate::cli::RootOpts::reload_sink_drain_limit_secs).
+    #[serde(default, skip)]
+    #[doc(hidden)]
+    pub reload_sink_drain_duration: Option<Duration>,
+
     /// Allow the configuration to be empty, resulting in a topology with no components.
     #[serde(default, skip)]
     #[doc(hidden)]
@@ -93,6 +100,7 @@ impl From<Config> for ConfigBuilder {
             tests,
             secret,
             graceful_shutdown_duration,
+            reload_sink_drain_duration,
         } = config;
 
         let transforms = transforms
@@ -126,6 +134,7 @@ impl From<Config> for ConfigBuilder {
             tests,
             secret,
             graceful_shutdown_duration,
+            reload_sink_drain_duration,
             allow_empty: false,
         }
     }

src/config/compiler.rs

@@ -60,6 +60,7 @@ pub fn compile(mut builder: ConfigBuilder) -> Result<(Config, Vec<String>), Vec<
         provider: _,
         secret,
         graceful_shutdown_duration,
+        reload_sink_drain_duration,
         allow_empty: _,
     } = builder;
     let graph = match Graph::new(nodes, global.wildcard_matching.unwrap_or_default()) {
@@ -126,6 +127,7 @@ pub fn compile(mut builder: ConfigBuilder) -> Result<(Config, Vec<String>), Vec<
             tests,
             secret,
             graceful_shutdown_duration,
+            reload_sink_drain_duration,
         };
 
         // Validate sinks in place

src/config/mod.rs

@@ -165,6 +165,7 @@ pub struct Config {
     tests: Vec<TestDefinition>,
     secret: IndexMap<ComponentKey, SecretBackends>,
     pub graceful_shutdown_duration: Option<Duration>,
+    pub reload_sink_drain_duration: Option<Duration>,
 }
 
 impl Config {

src/topology/builder.rs

@@ -137,6 +137,7 @@ struct Builder<'a> {
     inputs: HashMap<ComponentKey, (BufferSender<EventArray>, Inputs<OutputId>)>,
     healthchecks: HashMap<ComponentKey, Task>,
     detach_triggers: HashMap<ComponentKey, Trigger>,
+    reload_shutdown_triggers: HashMap<ComponentKey, oneshot::Sender<tokio::time::Instant>>,
     extra_context: ExtraContext,
     utilization_emitter: Option<UtilizationEmitter>,
     utilization_registry: UtilizationRegistry,
@@ -169,6 +170,7 @@ impl<'a> Builder<'a> {
             inputs: HashMap::new(),
             healthchecks: HashMap::new(),
             detach_triggers: HashMap::new(),
+            reload_shutdown_triggers: HashMap::new(),
             extra_context,
             utilization_emitter: emitter,
             utilization_registry: registry,
@@ -196,6 +198,7 @@ impl<'a> Builder<'a> {
                 healthchecks: self.healthchecks,
                 shutdown_coordinator: self.shutdown_coordinator,
                 detach_triggers: self.detach_triggers,
+                reload_shutdown_triggers: self.reload_shutdown_triggers,
                 metrics_storage: METRICS_STORAGE.clone(),
                 utilization: self
                     .utilization_emitter
@@ -719,6 +722,7 @@ impl<'a> Builder<'a> {
         };
 
         let (trigger, tripwire) = Tripwire::new();
+        let (reload_shutdown_tx, reload_shutdown_rx) = oneshot::channel();
 
         let utilization_sender = self
             .utilization_registry
@@ -742,24 +746,40 @@ impl<'a> Builder<'a> {
             let mut rx = Utilization::new(utilization_sender, component_key.clone(), rx);
 
             let events_received = register!(EventsReceived);
-            sink.run(
-                rx.by_ref()
-                    .filter(|events: &EventArray| ready(filter_events_type(events, input_type)))
-                    .inspect(|events| {
-                        events_received.emit(CountByteSize(
-                            events.len(),
-                            events.estimated_json_encoded_size_of(),
-                        ));
-                    })
-                    .take_until_if(tripwire),
-            )
-            .await
-            .map(|()| {
-                debug!("Sink finished normally.");
-                TaskOutput::Sink(rx)
-            })
-            .inspect_err(|error| debug!(%error, "Sink finished with an error."))
-            .map_err(Into::into)
+            let reload_deadline = async move {
+                if let Ok(deadline) = reload_shutdown_rx.await {
+                    tokio::time::sleep_until(deadline).await;
+                } else {
+                    // Dropping the topology during process shutdown must not cancel the sink here.
+                    futures::future::pending::<()>().await;
+                }
+            };
+            let result = select! {
+                biased;
+                result = sink.run(
+                    rx.by_ref()
+                        .filter(|events: &EventArray| ready(filter_events_type(events, input_type)))
+                        .inspect(|events| {
+                            events_received.emit(CountByteSize(
+                                events.len(),
+                                events.estimated_json_encoded_size_of(),
+                            ));
+                        })
+                        .take_until_if(tripwire),
+                ) => result,
+                () = reload_deadline => {
+                    warn!("Sink did not finish before the reload deadline. Cancelling in-flight work.");
+                    // Drop only the sink future, retaining unread buffered events for its replacement.
+                    Ok(())
+                }
+            };
+            result
+                .map(|()| {
+                    debug!("Sink stopped.");
+                    TaskOutput::Sink(rx)
+                })
+                .inspect_err(|error| debug!(%error, "Sink finished with an error."))
+                .map_err(Into::into)
         };
 
         let task = Task::new(key.clone(), typetag, sink);
@@ -806,6 +826,8 @@ impl<'a> Builder<'a> {
         self.healthchecks.insert(key.clone(), healthcheck_task);
         self.tasks.insert(key.clone(), task);
         self.detach_triggers.insert(key.clone(), trigger);
+        self.reload_shutdown_triggers
+            .insert(key.clone(), reload_shutdown_tx);
     }
 
     fn build_transform(
@@ -1104,6 +1126,8 @@ pub struct TopologyPieces {
     pub(super) healthchecks: HashMap<ComponentKey, Task>,
     pub(crate) shutdown_coordinator: SourceShutdownCoordinator,
     pub(crate) detach_triggers: HashMap<ComponentKey, Trigger>,
+    pub(crate) reload_shutdown_triggers:
+        HashMap<ComponentKey, oneshot::Sender<tokio::time::Instant>>,
     pub(crate) metrics_storage: MetricsStorage,
     pub(crate) utilization: Option<(UtilizationEmitter, UtilizationRegistry)>,
 }