From 60b8fea5235b1c359d8bd0c654ddd587ca461509 Mon Sep 17 00:00:00 2001 From: Petar Korudzhiev Date: Wed, 7 Oct 2026 10:45:15 +0300 Subject: [PATCH 1/2] Add ROTEL_LAMBDA_FLUSH_MODE and flush buffered telemetry on shutdown In adaptive mode, once invocations arrive faster than one a minute the extension flushes at the start of an invocation, at most every 20 s. An environment that then goes idle holds that telemetry until Lambda reclaims it. On SHUTDOWN the pipelines were cancelled and the exporters given 500 ms to drain, which was often not enough to reconnect and export, so the batch was dropped ("Exporters did not exit on channel close, cancelling."). - ROTEL_LAMBDA_FLUSH_MODE=adaptive|after-call (default adaptive, unchanged). after-call always exports after the response, so nothing is buffered while the environment is frozen. - On SHUTDOWN, force-flush logs, pipelines and exporters (bounded to 1.3 s of the 2 s shutdown budget) before cancelling. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CBAmNpXHPvXfhFq4havtNT --- src/lifecycle/flush_control.rs | 56 ++++++++++++++++++++++++++++++++- src/main.rs | 57 ++++++++++++++++++++++++++++++++-- 2 files changed, 109 insertions(+), 4 deletions(-) diff --git a/src/lifecycle/flush_control.rs b/src/lifecycle/flush_control.rs index 8da71f4..2e06fad 100644 --- a/src/lifecycle/flush_control.rs +++ b/src/lifecycle/flush_control.rs @@ -20,10 +20,23 @@ pub trait Clock { fn now(&self) -> u64; } +/// How the extension decides when to export. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] +pub enum FlushPolicy { + /// Flush after every invocation until a regular invocation rate is detected, + /// then flush periodically at the start of invocations. + #[default] + Adaptive, + /// Always flush after the function's response is sent. Nothing is held while the + /// execution environment is frozen, so nothing is lost if it is reclaimed while idle. + AfterCall, +} + pub struct FlushControl { rate: InvocationRate, inner: Arc>, clock: C, + policy: FlushPolicy, } struct Inner { @@ -56,12 +69,17 @@ impl PeriodicFlushControl { impl FlushControl { pub fn new(clock: C) -> Self { + Self::with_policy(clock, FlushPolicy::default()) + } + + pub fn with_policy(clock: C, policy: FlushPolicy) -> Self { Self { clock: clock.clone(), rate: InvocationRate::default(), inner: Arc::new(Mutex::new(Inner { last_flush: clock.now(), })), + policy, } } @@ -69,7 +87,12 @@ impl FlushControl { let now_millis = self.clock.now(); self.rate.add(now_millis); - let mode = match self.rate.is_faster_than(ACTIVE_INVOCATION_RATE_MILLIS) { + let adaptive = match self.policy { + FlushPolicy::AfterCall => None, + FlushPolicy::Adaptive => self.rate.is_faster_than(ACTIVE_INVOCATION_RATE_MILLIS), + }; + + let mode = match adaptive { // Not initialized, stick to flush per call None => AfterCall, @@ -182,6 +205,37 @@ mod tests { } } + #[test] + fn test_after_call_policy_never_goes_periodic() { + let clock = TestClock::new(1000); + let mut flush_control = FlushControl::with_policy(clock.clone(), FlushPolicy::AfterCall); + + // The same fast rate that switches the adaptive policy to Periodic + for _ in 1..=50 { + clock.advance(ACTIVE_INVOCATION_RATE_MILLIS / 10); + match flush_control.pick() { + FlushMode::AfterCall => {} + _ => panic!("Expected AfterCall mode for every invocation"), + } + } + } + + #[test] + fn test_default_policy_is_adaptive() { + let clock = TestClock::new(1000); + let mut flush_control = FlushControl::new(clock.clone()); + + for _ in 1..=20 { + clock.advance(ACTIVE_INVOCATION_RATE_MILLIS / 2); + let _ = flush_control.pick(); + } + + match flush_control.pick() { + FlushMode::Periodic(_) => {} + _ => panic!("Expected the default policy to switch to Periodic"), + } + } + #[test] fn test_transition_from_periodic_to_after_call() { let clock = TestClock::new(1000); diff --git a/src/main.rs b/src/main.rs index 19b6857..a307d1a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,7 @@ use rotel_extension::env::{EnvArnParser, resolve_secrets}; use rotel_extension::lambda; use rotel_extension::lambda::telemetry_api::TelemetryAPI; use rotel_extension::lifecycle::flush_control::{ - Clock, DEFAULT_FLUSH_INTERVAL_MILLIS, FlushControl, FlushMode, + Clock, DEFAULT_FLUSH_INTERVAL_MILLIS, FlushControl, FlushMode, FlushPolicy, }; use rustls::crypto::CryptoProvider; use std::collections::HashMap; @@ -51,6 +51,11 @@ pub const FLUSH_LOGS_TIMEOUT_MILLIS: u64 = 100; // can be short, simply forces b pub const FLUSH_PIPELINE_TIMEOUT_MILLIS: u64 = 500; pub const FLUSH_EXPORTERS_TIMEOUT_MILLIS: u64 = 3_000; +// Lambda gives external extensions two seconds in total to shut down. Spend most of it +// exporting what is still buffered, and keep the rest for tearing down the pipelines. +pub const SHUTDOWN_TOTAL_MILLIS: u64 = 2_000; +pub const SHUTDOWN_FLUSH_TIMEOUT_MILLIS: u64 = 1_300; + #[derive(Debug, Parser)] #[command(name = "rotel-lambda-extension")] #[command(bin_name = "rotel-lambda-extension")] @@ -72,6 +77,17 @@ struct Arguments { /// Environment environment: String, + #[arg( + value_enum, + long, + env = "ROTEL_LAMBDA_FLUSH_MODE", + default_value = "adaptive" + )] + /// When to export: `adaptive` switches to periodic flushing at the start of + /// invocations once the invocation rate is regular; `after-call` always flushes + /// after the response, so nothing is buffered while the environment is frozen. + lambda_flush_mode: LambdaFlushModeArg, + // This is ignored in these options, but we keep it here to avoid an error on unknown // options #[arg(long)] @@ -95,6 +111,21 @@ pub enum LogFormatArg { Json, } +#[derive(Copy, Clone, PartialEq, Eq, Debug, ValueEnum)] +pub enum LambdaFlushModeArg { + Adaptive, + AfterCall, +} + +impl From for FlushPolicy { + fn from(arg: LambdaFlushModeArg) -> Self { + match arg { + LambdaFlushModeArg::Adaptive => FlushPolicy::Adaptive, + LambdaFlushModeArg::AfterCall => FlushPolicy::AfterCall, + } + } +} + fn main() -> ExitCode { let start_time = Instant::now(); @@ -139,6 +170,7 @@ fn main() -> ExitCode { port_map, telemetry_listener, &opt.environment, + opt.lambda_flush_mode.into(), ) { Ok(_) => {} Err(e) => { @@ -197,6 +229,7 @@ async fn run_extension( port_map: HashMap, telemetry_listener: Listener, env: &String, + flush_policy: FlushPolicy, ) -> Result<(), BoxError> { let mut tapi_join_set = JoinSet::new(); let mut agent_join_set = JoinSet::new(); @@ -301,7 +334,8 @@ async fn run_extension( }; handle_next_response(next_evt); - let mut flush_control = FlushControl::new(SystemClock {}); + info!(?flush_policy, "Using flush policy"); + let mut flush_control = FlushControl::with_policy(SystemClock {}, flush_policy); 'outer: loop { let mode = flush_control.pick(); @@ -428,7 +462,24 @@ async fn run_extension( } // We have two seconds to completely shutdown - let final_stop = Instant::now().add(Duration::from_secs(2)); + let final_stop = Instant::now().add(Duration::from_millis(SHUTDOWN_TOTAL_MILLIS)); + + // Export what is still buffered (in periodic mode, up to one flush interval of + // telemetry) before the pipelines are cancelled; otherwise it is dropped. + if timeout( + Duration::from_millis(SHUTDOWN_FLUSH_TIMEOUT_MILLIS), + force_flush( + &mut flush_logs_tx, + &mut flush_pipeline_tx, + &mut flush_exporters_tx, + &mut default_flush_interval, + ), + ) + .await + .is_err() + { + warn!("Timed out flushing telemetry on shutdown"); + } // Wait up to 500ms for the TelemetryAPI to shutdown, this will stop the logs pipeline telemetry_cancel.cancel(); From 77e8220fc20a70102add97599d2c88223f5ca511 Mon Sep 17 00:00:00 2001 From: Petar Korudzhiev Date: Wed, 7 Oct 2026 11:34:35 +0300 Subject: [PATCH 2/2] Add ROTEL_LAMBDA_FORWARD_FUNCTION_LOGS to skip function stdout/stderr Functions that already export their logs over OTLP get every line twice: once from the SDK and once from the Telemetry API function stream. When the setting is false the extension does not subscribe to function records; platform records (which drive flushing) and the extension's own logs are unchanged. Defaults to true. Also documents ROTEL_LAMBDA_FLUSH_MODE and the shutdown flush. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CBAmNpXHPvXfhFq4havtNT --- README.md | 8 ++++++++ src/lambda/api.rs | 38 +++++++++++++++++++++++++++++++++----- src/main.rs | 16 +++++++++++++++- 3 files changed, 56 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 7db741c..7e95c37 100644 --- a/README.md +++ b/README.md @@ -171,6 +171,14 @@ If the extension detects a regular invocation pattern—such as invocations occu For long-running invocations, a **global backup timer** is used to flush telemetry periodically. This timer is reset whenever a regular flush occurs, ensuring that telemetry is still sent even if invocation patterns become irregular. +Set `ROTEL_LAMBDA_FLUSH_MODE=after-call` to always flush at the end of each invocation. Periodic flushing holds telemetry while the execution environment is frozen between invocations; if Lambda reclaims an idle environment, whatever is still buffered must be exported within the shutdown window. Flushing after every call keeps nothing buffered across invocations, at the cost of the billed duration described above. The default is `adaptive`. + +On shutdown the extension flushes any buffered telemetry before stopping its exporters. + +## Function logs + +The extension subscribes to the Lambda Telemetry API and exports the function's stdout and stderr as logs. If your function already exports its logs over OTLP (for example through the OpenTelemetry SDK's log provider), every line would arrive twice. Set `ROTEL_LAMBDA_FORWARD_FUNCTION_LOGS=false` to skip function logs; the extension's own logs are still exported. The default is `true`. + ## Examples These are example repos demonstrating how to use the Rotel Lambda Extension. diff --git a/src/lambda/api.rs b/src/lambda/api.rs index 7fca197..f96892d 100644 --- a/src/lambda/api.rs +++ b/src/lambda/api.rs @@ -113,14 +113,11 @@ pub async fn telemetry_subscribe( client: Client>, ext_id: &str, addr: &SocketAddr, + forward_function_logs: bool, ) -> Result<(), BoxError> { let sub = serde_json::json!(TelemetryAPISubscribe { schema_version: TELEMETRY_API_SCHEMA.to_string(), - types: vec![ - "platform".to_string(), - "function".to_string(), - "extension".to_string() - ], + types: subscription_types(forward_function_logs), buffering: TelemetryAPISubscribeBuffering { // todo: these are the defaults from API ref, consider adjusting max_items: 1000, @@ -154,6 +151,19 @@ pub async fn telemetry_subscribe( Ok(()) } +/// The Telemetry API record types to subscribe to. Platform records drive the +/// flush after each invocation and extension records carry this extension's own +/// logs, so both are always on. Function records are the function's stdout and +/// stderr; skip them when the function already exports its logs over OTLP. +fn subscription_types(forward_function_logs: bool) -> Vec { + let mut types = vec!["platform".to_string()]; + if forward_function_logs { + types.push("function".to_string()); + } + types.push("extension".to_string()); + types +} + fn lambda_api_url(path: &str) -> Result { let base_api = std::env::var("AWS_LAMBDA_RUNTIME_API") .map_err(|e| format!("Unable to read AWS_LAMBDA_RUNTIME_API: {:?}", e))?; @@ -164,3 +174,21 @@ fn lambda_api_url(path: &str) -> Result { Ok(format!("http://{}{}", base_api, path)) } } + +#[cfg(test)] +mod tests { + use super::subscription_types; + + #[test] + fn test_subscribes_to_function_logs_by_default() { + assert_eq!( + subscription_types(true), + vec!["platform", "function", "extension"] + ); + } + + #[test] + fn test_skips_function_logs_when_disabled() { + assert_eq!(subscription_types(false), vec!["platform", "extension"]); + } +} diff --git a/src/main.rs b/src/main.rs index a307d1a..5ee7d99 100644 --- a/src/main.rs +++ b/src/main.rs @@ -88,6 +88,17 @@ struct Arguments { /// after the response, so nothing is buffered while the environment is frozen. lambda_flush_mode: LambdaFlushModeArg, + #[arg( + long, + env = "ROTEL_LAMBDA_FORWARD_FUNCTION_LOGS", + default_value_t = true, + action = clap::ArgAction::Set, + value_parser = clap::builder::BoolishValueParser::new() + )] + /// Export the function's stdout and stderr as logs. Turn this off when the + /// function already exports its logs over OTLP, or every line arrives twice. + lambda_forward_function_logs: bool, + // This is ignored in these options, but we keep it here to avoid an error on unknown // options #[arg(long)] @@ -171,6 +182,7 @@ fn main() -> ExitCode { telemetry_listener, &opt.environment, opt.lambda_flush_mode.into(), + opt.lambda_forward_function_logs, ) { Ok(_) => {} Err(e) => { @@ -230,6 +242,7 @@ async fn run_extension( telemetry_listener: Listener, env: &String, flush_policy: FlushPolicy, + forward_function_logs: bool, ) -> Result<(), BoxError> { let mut tapi_join_set = JoinSet::new(); let mut agent_join_set = JoinSet::new(); @@ -303,6 +316,7 @@ async fn run_extension( client.clone(), &r.extension_id, &telemetry_listener.bound_address()?, + forward_function_logs, ) .await { @@ -334,7 +348,7 @@ async fn run_extension( }; handle_next_response(next_evt); - info!(?flush_policy, "Using flush policy"); + info!(?flush_policy, forward_function_logs, "Using flush policy"); let mut flush_control = FlushControl::with_policy(SystemClock {}, flush_policy); 'outer: loop {