autumn_plugin_aws_lambda/
runtime.rs

1//! The Lambda runtime loop and the readiness wait.
2
3use std::time::Duration;
4
5use crate::PluginError;
6use crate::proxy::LambdaProxy;
7use crate::timing::{as_ms, next_delay_ms};
8
9/// First delay of the readiness wait.
10const READY_INITIAL: Duration = Duration::from_millis(5);
11/// Longest delay of the readiness wait.
12const READY_MAX: Duration = Duration::from_millis(100);
13/// Name of the runtime loop thread.
14const THREAD_NAME: &str = "aws-lambda-runtime";
15/// Time between two Runtime API checks in concurrent mode.
16const WATCH_INTERVAL: Duration = Duration::from_secs(1);
17/// Failed Runtime API checks in sequence before the process exits.
18const WATCH_LIMIT: u32 = 5;
19
20/// How the plugin sends responses to Lambda.
21#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
22#[non_exhaustive]
23pub enum ResponseMode {
24    /// Read the full body (up to 6 MiB), then send it. It works with all
25    /// HTTP event sources. This is the default.
26    #[default]
27    Buffered,
28    /// Stream the body. Use it only with a Function URL or with API
29    /// Gateway response streaming.
30    Streaming,
31}
32
33/// The state of Autumn startup.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub(crate) enum Startup {
36    /// Startup hooks still run.
37    Pending,
38    /// Autumn is ready.
39    Complete,
40    /// Autumn shuts down.
41    Stopping,
42}
43
44/// Settings for the runtime loop thread.
45#[derive(Debug)]
46pub(crate) struct LoopSettings {
47    pub proxy: LambdaProxy,
48    pub mode: ResponseMode,
49    /// `None`: wait until startup is complete, with no time limit.
50    pub readiness_timeout: Option<Duration>,
51    /// Events to handle at the same time.
52    pub concurrency: u32,
53    /// Address of the Runtime API (`host:port`).
54    pub runtime_api: String,
55}
56
57/// Calls `ready` until it returns `true` or until `timeout` ends.
58///
59/// The delay starts at `initial` (minimum 1 ms) and doubles up to `max`.
60/// The last sleep stops at `timeout`. Returns the last value of `ready`.
61/// Verus twin: `verification/core.rs::readiness_probes`.
62pub async fn wait_until(
63    ready: impl Fn() -> bool,
64    timeout: Duration,
65    initial: Duration,
66    max: Duration,
67) -> bool {
68    let timeout_ms = as_ms(timeout);
69    let mut delay_ms = as_ms(initial).max(1);
70    let max_ms = as_ms(max).max(delay_ms);
71    let start = tokio::time::Instant::now();
72    loop {
73        if ready() {
74            return true;
75        }
76        let slept_ms = as_ms(start.elapsed());
77        if slept_ms >= timeout_ms {
78            return false;
79        }
80        let step = delay_ms.min(timeout_ms - slept_ms);
81        tokio::time::sleep(Duration::from_millis(step)).await;
82        delay_ms = next_delay_ms(delay_ms, max_ms);
83    }
84}
85
86/// Starts the runtime loop on a new thread.
87///
88/// The thread waits for Autumn startup. If Autumn shuts down first, the
89/// loop does not start. The `lambda_http` loop future is not `Send`, so
90/// `tokio::spawn` cannot accept it. The loop gets its own thread and
91/// runtime.
92///
93/// The loop runs until the process ends. If it stops or panics, the
94/// process exits with code 1. Lambda then starts a new execution
95/// environment. This exit does not run Autumn shutdown hooks.
96///
97/// # Errors
98///
99/// Returns an error if the thread cannot start.
100pub(crate) fn spawn_loop(
101    settings: LoopSettings,
102    startup: impl Fn() -> Startup + Send + 'static,
103) -> Result<(), PluginError> {
104    std::thread::Builder::new()
105        .name(THREAD_NAME.to_owned())
106        .spawn(move || {
107            let runtime = match build_runtime(settings.concurrency) {
108                Ok(runtime) => runtime,
109                Err(error) => {
110                    tracing::error!(%error, "cannot build the Lambda runtime; exiting");
111                    std::process::exit(1);
112                }
113            };
114            let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
115                runtime.block_on(run_when_ready(settings, startup))
116            }));
117            if matches!(outcome, Ok(None)) {
118                return;
119            }
120            tracing::error!(?outcome, "AWS Lambda runtime loop stopped; exiting");
121            std::process::exit(1);
122        })
123        .map_err(PluginError::Runtime)?;
124    Ok(())
125}
126
127/// Waits for startup, then runs the loop. Returns `None` if the loop does
128/// not start.
129#[allow(
130    clippy::future_not_send,
131    reason = "the lambda_http loop is not Send; it runs on its own thread"
132)]
133async fn run_when_ready(
134    settings: LoopSettings,
135    startup: impl Fn() -> Startup,
136) -> Option<Result<(), lambda_http::Error>> {
137    let limit = settings.readiness_timeout.unwrap_or(Duration::MAX);
138    wait_until(
139        || startup() != Startup::Pending,
140        limit,
141        READY_INITIAL,
142        READY_MAX,
143    )
144    .await;
145    match startup() {
146        Startup::Stopping => {
147            tracing::info!("Autumn shuts down; the Lambda loop does not start");
148            return None;
149        }
150        Startup::Pending => tracing::warn!(
151            ?limit,
152            "Autumn startup is not complete; the Lambda loop starts now"
153        ),
154        Startup::Complete => {}
155    }
156    if settings.concurrency > 1 {
157        // In concurrent mode, lambda_runtime retries a failed `/next` with
158        // no delay and does not stop. This check stops the process instead.
159        tokio::spawn(watch_runtime_api(settings.runtime_api));
160    }
161    tracing::info!(mode = ?settings.mode, concurrency = settings.concurrency, "AWS Lambda runtime loop runs");
162    Some(run(settings.proxy, settings.mode).await)
163}
164
165/// Builds the loop runtime. One thread is enough for one event at a time.
166fn build_runtime(concurrency: u32) -> std::io::Result<tokio::runtime::Runtime> {
167    if concurrency <= 1 {
168        return tokio::runtime::Builder::new_current_thread()
169            .enable_all()
170            .build();
171    }
172    let cores = std::thread::available_parallelism().map_or(1, std::num::NonZero::get);
173    let wanted = usize::try_from(concurrency).unwrap_or(usize::MAX);
174    tokio::runtime::Builder::new_multi_thread()
175        .worker_threads(cores.min(wanted))
176        .enable_all()
177        .build()
178}
179
180/// Exits the process when the Runtime API refuses connections
181/// [`WATCH_LIMIT`] times in sequence.
182async fn watch_runtime_api(address: String) {
183    let mut failures = 0;
184    loop {
185        tokio::time::sleep(WATCH_INTERVAL).await;
186        let connect = tokio::net::TcpStream::connect(address.as_str());
187        match tokio::time::timeout(WATCH_INTERVAL, connect).await {
188            Ok(Ok(_)) => failures = 0,
189            _ => failures += 1,
190        }
191        if failures >= WATCH_LIMIT {
192            tracing::error!(%address, "AWS Lambda Runtime API is not reachable; exiting");
193            std::process::exit(1);
194        }
195    }
196}
197
198/// Runs the Lambda runtime loop with `proxy`. It returns only on a fatal error.
199///
200/// `lambda_http` uses concurrent mode when `AWS_LAMBDA_MAX_CONCURRENCY` is
201/// more than 1. Otherwise, it handles one event at a time.
202#[allow(
203    clippy::future_not_send,
204    reason = "the lambda_http loop is not Send; it runs on its own thread"
205)]
206async fn run(proxy: LambdaProxy, mode: ResponseMode) -> Result<(), lambda_http::Error> {
207    let proxy = proxy.response_mode(mode);
208    match mode {
209        ResponseMode::Streaming => lambda_http::run_with_streaming_response_concurrent(proxy).await,
210        _ => lambda_http::run_concurrent(proxy).await,
211    }
212}
autumn_plugin_aws_lambda/runtime.rs

autumn_plugin_aws_lambda/
runtime.rs
