Skip to main content

autumn_plugin_aws_lambda/
plugin.rs

1//! The Autumn plugin.
2
3use std::net::SocketAddr;
4use std::time::Duration;
5
6use autumn_web::AppState;
7use autumn_web::app::AppBuilder;
8use autumn_web::config::ServerConfig;
9use autumn_web::plugin::Plugin;
10
11use crate::activation::{RUNTIME_API_ENV, lambda_concurrency, missing_lambda_env};
12use crate::proxy::LambdaProxy;
13use crate::runtime::{LoopSettings, Startup, spawn_loop};
14use crate::upstream::Upstream;
15use crate::{Activation, PluginError, ResponseMode};
16
17/// Autumn sets this variable to run one task instead of the server.
18const TASK_ENV: &str = "AUTUMN_RUN_TASK";
19
20/// Runs an Autumn app on AWS Lambda.
21///
22/// Register it with `.plugin(AwsLambdaPlugin::new())`.
23#[derive(Debug, Clone)]
24pub struct AwsLambdaPlugin {
25    activation: Activation,
26    response_mode: ResponseMode,
27    upstream: Option<SocketAddr>,
28    readiness_timeout: Option<Duration>,
29    timeout_margin: Duration,
30}
31
32impl Default for AwsLambdaPlugin {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38impl AwsLambdaPlugin {
39    /// Default time between the end of the upstream call and the
40    /// invocation deadline.
41    pub const DEFAULT_TIMEOUT_MARGIN: Duration = Duration::from_millis(100);
42
43    /// Makes the plugin with default settings.
44    #[must_use]
45    pub const fn new() -> Self {
46        Self {
47            activation: Activation::Auto,
48            response_mode: ResponseMode::Buffered,
49            upstream: None,
50            readiness_timeout: None,
51            timeout_margin: Self::DEFAULT_TIMEOUT_MARGIN,
52        }
53    }
54
55    /// Sets when the runtime loop starts. Default: [`Activation::Auto`].
56    #[must_use]
57    pub const fn activation(mut self, activation: Activation) -> Self {
58        self.activation = activation;
59        self
60    }
61
62    /// Sets how the plugin sends responses to Lambda.
63    /// Default: [`ResponseMode::Buffered`].
64    #[must_use]
65    pub const fn response_mode(mut self, mode: ResponseMode) -> Self {
66        self.response_mode = mode;
67        self
68    }
69
70    /// Sets the upstream address. Default: from Autumn `[server]` config.
71    ///
72    /// With this setting, the plugin does not check `[server]`. The address
73    /// must reach the Autumn listener.
74    #[must_use]
75    pub const fn upstream(mut self, addr: SocketAddr) -> Self {
76        self.upstream = Some(addr);
77        self
78    }
79
80    /// Sets the maximum time to wait for Autumn startup before the first
81    /// event. When the time ends, the loop starts. Default: no limit.
82    #[must_use]
83    pub const fn readiness_timeout(mut self, timeout: Duration) -> Self {
84        self.readiness_timeout = Some(timeout);
85        self
86    }
87
88    /// Sets the time between the end of the upstream call and the
89    /// invocation deadline. Default: [`Self::DEFAULT_TIMEOUT_MARGIN`].
90    #[must_use]
91    pub const fn timeout_margin(mut self, margin: Duration) -> Self {
92        self.timeout_margin = margin;
93        self
94    }
95
96    /// Decides what to do at startup. `env` returns the value of one
97    /// variable, or `None`.
98    ///
99    /// Returns `Ok(None)` if the plugin stays idle. Returns
100    /// `Ok(Some(upstream))` if the loop starts.
101    pub(crate) fn plan(
102        &self,
103        server: &ServerConfig,
104        env: impl Fn(&str) -> Option<String>,
105    ) -> Result<Option<Upstream>, PluginError> {
106        if !self.activation.is_active(env(RUNTIME_API_ENV).is_some()) {
107            return Ok(None);
108        }
109        // In task mode, Autumn does not start the listener.
110        if env(TASK_ENV).is_some_and(|task| !task.trim().is_empty()) {
111            return Ok(None);
112        }
113        let names = missing_lambda_env(env);
114        if !names.is_empty() {
115            return Err(PluginError::MissingEnv { names });
116        }
117        let upstream = match self.upstream {
118            Some(addr) => Upstream::new(addr),
119            None => Upstream::from_server_config(server)?,
120        };
121        Ok(Some(upstream))
122    }
123
124    /// Starts the runtime loop if [`plan`](Self::plan) returns an upstream.
125    fn start(&self, state: &AppState) -> Result<(), PluginError> {
126        let env = |name: &str| std::env::var(name).ok();
127        let server = &state.config_arc().server;
128        let Some(upstream) = self.plan(server, env)? else {
129            tracing::debug!("AWS Lambda is not active; the plugin is idle");
130            return Ok(());
131        };
132        let slow = slow_shutdown_settings(server);
133        if !slow.is_empty() {
134            tracing::warn!(
135                settings = ?slow,
136                "Lambda gives little time at shutdown; set these to 0 or 1"
137            );
138        }
139        tracing::info!(%upstream, "AWS Lambda detected; the plugin sends events to Autumn");
140        let settings = LoopSettings {
141            proxy: LambdaProxy::new(upstream).timeout_margin(self.timeout_margin),
142            mode: self.response_mode,
143            readiness_timeout: self.readiness_timeout,
144            concurrency: lambda_concurrency(env),
145            runtime_api: env(RUNTIME_API_ENV).unwrap_or_default(),
146        };
147        let state = state.clone();
148        spawn_loop(settings, move || {
149            let probes = state.probes();
150            if probes.is_shutting_down() {
151                Startup::Stopping
152            } else if probes.is_startup_complete() {
153                Startup::Complete
154            } else {
155                Startup::Pending
156            }
157        })
158    }
159}
160
161/// Returns the `[server]` settings that are too slow for Lambda shutdown.
162///
163/// Lambda stops the process soon after `SIGTERM`.
164#[must_use]
165pub fn slow_shutdown_settings(server: &ServerConfig) -> Vec<&'static str> {
166    let mut slow = Vec::new();
167    if server.prestop_grace_secs > 0 {
168        slow.push("server.prestop_grace_secs");
169    }
170    if server.shutdown_timeout_secs > 1 {
171        slow.push("server.shutdown_timeout_secs");
172    }
173    slow
174}
175
176impl Plugin for AwsLambdaPlugin {
177    fn build(self, app: AppBuilder) -> AppBuilder {
178        app.on_startup(move |state| {
179            let result = self.start(&state).map_err(Into::into);
180            async move { result }
181        })
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use std::collections::HashMap;
188
189    use super::*;
190    use crate::PluginError;
191
192    fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
193        let map: HashMap<String, String> = pairs
194            .iter()
195            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
196            .collect();
197        move |name: &str| map.get(name).cloned()
198    }
199
200    const LAMBDA_ENV: [(&str, &str); 4] = [
201        ("AWS_LAMBDA_RUNTIME_API", "127.0.0.1:9001"),
202        ("AWS_LAMBDA_FUNCTION_NAME", "f"),
203        ("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128"),
204        ("AWS_LAMBDA_FUNCTION_VERSION", "$LATEST"),
205    ];
206
207    #[test]
208    fn plugin_registers_once() {
209        let app = autumn_web::app()
210            .plugin(AwsLambdaPlugin::new())
211            .plugin(AwsLambdaPlugin::new());
212        assert!(app.has_plugin(std::any::type_name::<AwsLambdaPlugin>()));
213    }
214
215    #[test]
216    fn auto_is_idle_outside_lambda() {
217        let plan = AwsLambdaPlugin::new().plan(&ServerConfig::default(), env(&[]));
218        assert!(matches!(plan, Ok(None)));
219    }
220
221    #[test]
222    fn never_is_idle_on_lambda() {
223        let plan = AwsLambdaPlugin::new()
224            .activation(Activation::Never)
225            .plan(&ServerConfig::default(), env(&LAMBDA_ENV));
226        assert!(matches!(plan, Ok(None)));
227    }
228
229    #[test]
230    fn always_fails_outside_lambda() {
231        let plan = AwsLambdaPlugin::new()
232            .activation(Activation::Always)
233            .plan(&ServerConfig::default(), env(&[]));
234        assert!(matches!(plan, Err(PluginError::MissingEnv { names, .. }) if names.len() == 4));
235    }
236
237    #[test]
238    fn auto_fails_on_partial_lambda_env() {
239        let plan = AwsLambdaPlugin::new().plan(&ServerConfig::default(), env(&LAMBDA_ENV[..1]));
240        assert!(matches!(plan, Err(PluginError::MissingEnv { names, .. }) if names.len() == 3));
241    }
242
243    #[test]
244    fn auto_on_lambda_uses_server_config() {
245        let plan = AwsLambdaPlugin::new()
246            .plan(&ServerConfig::default(), env(&LAMBDA_ENV))
247            .expect("valid plan")
248            .expect("active");
249        assert_eq!(plan.to_string(), "http://127.0.0.1:3000");
250    }
251
252    #[test]
253    fn upstream_override_wins() {
254        let addr: SocketAddr = "0.0.0.0:8080".parse().expect("addr");
255        let plan = AwsLambdaPlugin::new()
256            .upstream(addr)
257            .plan(&ServerConfig::default(), env(&LAMBDA_ENV))
258            .expect("valid plan")
259            .expect("active");
260        assert_eq!(plan.to_string(), "http://127.0.0.1:8080");
261    }
262
263    #[test]
264    fn task_mode_is_idle_on_lambda() {
265        let mut pairs = LAMBDA_ENV.to_vec();
266        pairs.push(("AUTUMN_RUN_TASK", "backfill"));
267        let plan = AwsLambdaPlugin::new().plan(&ServerConfig::default(), env(&pairs));
268        assert!(matches!(plan, Ok(None)));
269    }
270
271    #[test]
272    fn blank_task_name_is_not_task_mode() {
273        let mut pairs = LAMBDA_ENV.to_vec();
274        pairs.push(("AUTUMN_RUN_TASK", "  "));
275        let plan = AwsLambdaPlugin::new().plan(&ServerConfig::default(), env(&pairs));
276        assert!(matches!(plan, Ok(Some(_))));
277    }
278
279    #[test]
280    fn always_starts_on_lambda() {
281        let plan = AwsLambdaPlugin::new()
282            .activation(Activation::Always)
283            .plan(&ServerConfig::default(), env(&LAMBDA_ENV));
284        assert!(matches!(plan, Ok(Some(_))));
285    }
286
287    #[test]
288    fn upstream_override_skips_server_checks() {
289        let server = ServerConfig {
290            port: 0,
291            unix_socket: Some("/tmp/app.sock".to_owned()),
292            ..ServerConfig::default()
293        };
294        let addr: SocketAddr = "127.0.0.1:9000".parse().expect("addr");
295        let plan = AwsLambdaPlugin::new()
296            .upstream(addr)
297            .plan(&server, env(&LAMBDA_ENV));
298        assert!(matches!(plan, Ok(Some(up)) if up.addr() == addr));
299    }
300
301    #[test]
302    fn server_errors_stop_the_plan() {
303        let server = ServerConfig {
304            port: 0,
305            ..ServerConfig::default()
306        };
307        let plan = AwsLambdaPlugin::new().plan(&server, env(&LAMBDA_ENV));
308        assert!(matches!(plan, Err(PluginError::ZeroPort)));
309    }
310
311    #[test]
312    fn default_equals_new() {
313        let a = AwsLambdaPlugin::default();
314        assert_eq!(a.readiness_timeout, None);
315        assert_eq!(a.timeout_margin, AwsLambdaPlugin::DEFAULT_TIMEOUT_MARGIN);
316        assert_eq!(a.activation, Activation::Auto);
317        assert_eq!(a.response_mode, ResponseMode::Buffered);
318        assert!(a.upstream.is_none());
319    }
320}