autumn_plugin_aws_lambda/
plugin.rs1use 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
17const TASK_ENV: &str = "AUTUMN_RUN_TASK";
19
20#[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 pub const DEFAULT_TIMEOUT_MARGIN: Duration = Duration::from_millis(100);
42
43 #[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 #[must_use]
57 pub const fn activation(mut self, activation: Activation) -> Self {
58 self.activation = activation;
59 self
60 }
61
62 #[must_use]
65 pub const fn response_mode(mut self, mode: ResponseMode) -> Self {
66 self.response_mode = mode;
67 self
68 }
69
70 #[must_use]
75 pub const fn upstream(mut self, addr: SocketAddr) -> Self {
76 self.upstream = Some(addr);
77 self
78 }
79
80 #[must_use]
83 pub const fn readiness_timeout(mut self, timeout: Duration) -> Self {
84 self.readiness_timeout = Some(timeout);
85 self
86 }
87
88 #[must_use]
91 pub const fn timeout_margin(mut self, margin: Duration) -> Self {
92 self.timeout_margin = margin;
93 self
94 }
95
96 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 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 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#[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}