|
| 1 | +use lambda_extension::Extension; |
| 2 | +use lambda_runtime::{service_fn, tracing, Error, LambdaEvent}; |
| 3 | +use serde::{Deserialize, Serialize}; |
| 4 | +use tokio::signal::unix::{signal, SignalKind}; |
| 5 | + |
| 6 | +/// This is also a made-up example. Requests come into the runtime as unicode |
| 7 | +/// strings in json format, which can map to any structure that implements `serde::Deserialize` |
| 8 | +/// The runtime pays no attention to the contents of the request payload. |
| 9 | +#[derive(Deserialize)] |
| 10 | +struct Request { |
| 11 | + command: String, |
| 12 | +} |
| 13 | + |
| 14 | +/// This is a made-up example of what a response structure may look like. |
| 15 | +/// There is no restriction on what it can be. The runtime requires responses |
| 16 | +/// to be serialized into json. The runtime pays no attention |
| 17 | +/// to the contents of the response payload. |
| 18 | +#[derive(Serialize)] |
| 19 | +struct Response { |
| 20 | + req_id: String, |
| 21 | + msg: String, |
| 22 | +} |
| 23 | + |
| 24 | +pub(crate) async fn handler(event: LambdaEvent<Request>) -> Result<Response, Error> { |
| 25 | + // extract some useful info from the request |
| 26 | + let command = event.payload.command; |
| 27 | + |
| 28 | + tracing::info!("executing {command:#}"); |
| 29 | + |
| 30 | + // prepare the response |
| 31 | + let resp = Response { |
| 32 | + req_id: event.context.request_id, |
| 33 | + msg: format!("Command {} executed.", command), |
| 34 | + }; |
| 35 | + |
| 36 | + // return `Response` (it will be serialized to JSON automatically by the runtime) |
| 37 | + Ok(resp) |
| 38 | +} |
| 39 | + |
| 40 | +#[tokio::main] |
| 41 | +async fn main() -> Result<(), Error> { |
| 42 | + // Let's pretend this Lambda function is: |
| 43 | + // - sensitive to the tokio runtime being blocked due to heavy async concurrency in its handler |
| 44 | + // - stays fairly warm |
| 45 | + // - is somewhat long-running |
| 46 | + // In that case, we might prefer to avoid brief runtime blocks from flushing to stdout in the local context. |
| 47 | + // ref: https://github.com/tokio-rs/tracing/issues/2653 |
| 48 | + // |
| 49 | + // So, we can use a dedicated logging thread with tracing_appender. The downside is that we need |
| 50 | + // to flush any outstanding entries before shutdown - hence the need for graceful shutdown. |
| 51 | + let (writer, log_guard) = tracing_appender::non_blocking(std::io::stdout()); |
| 52 | + tracing::init_default_subscriber_with_writer(writer); |
| 53 | + |
| 54 | + // You need an extension registered with the Lambda orchestrator in order for your process |
| 55 | + // to receive a SIGTERM for graceful shutdown. |
| 56 | + // |
| 57 | + // We accomplish this here by registering a no-op internal extension, which immediately returns. |
| 58 | + // (Unfortunately it does need to subscribe to at least one event type, in order to get a SIGTERM signal). |
| 59 | + // This adds overhead, but very little. |
| 60 | + // |
| 61 | + // You could also run a useful internal extension, such as in `../extension-internal-flush`. |
| 62 | + // |
| 63 | + // If an external extension (for example CloudWatch Lambda Insights) is running, it would register itself from outside |
| 64 | + // your Rust process and also get you SIGTERM. In that case, you don't need an internal extension at all. |
| 65 | + let extension = Extension::new() |
| 66 | + // Internal extensions only support INVOKE events. |
| 67 | + .with_events(&["INVOKE"]) |
| 68 | + // Internal extension names MUST be unique within a given Lambda function. |
| 69 | + .with_extension_name("no-op") |
| 70 | + // Extensions MUST be registered before calling lambda_runtime::run(), which ends the Init |
| 71 | + // phase and begins the Invoke phase. |
| 72 | + .register() |
| 73 | + .await |
| 74 | + .expect("could not register extension"); |
| 75 | + |
| 76 | + // Handle SIGTERM signal. |
| 77 | + // |
| 78 | + // In this graceful shutdown, we flush our background logging thread by dropping its guard. |
| 79 | + // |
| 80 | + // More on signal handling: |
| 81 | + // https://tokio.rs/tokio/topics/shutdown |
| 82 | + // https://rust-cli.github.io/book/in-depth/signals.html |
| 83 | + tokio::spawn(async move { |
| 84 | + let mut sigint = signal(SignalKind::interrupt()).unwrap(); |
| 85 | + let mut sigterm = signal(SignalKind::terminate()).unwrap(); |
| 86 | + tokio::select! { |
| 87 | + _sigint = sigint.recv() => { |
| 88 | + tracing::info!("[runtime] SIGINT received"); |
| 89 | + tracing::info!("[runtime] Graceful shutdown in progress ..."); |
| 90 | + std::mem::drop(log_guard); |
| 91 | + |
| 92 | + eprintln!("[runtime] Graceful shutdown completed"); |
| 93 | + std::process::exit(0); |
| 94 | + }, |
| 95 | + _sigterm = sigterm.recv()=> { |
| 96 | + tracing::info!("[runtime] SIGTERM received"); |
| 97 | + tracing::info!("[runtime] Graceful shutdown in progress ..."); |
| 98 | + std::mem::drop(log_guard); |
| 99 | + |
| 100 | + eprintln!("[runtime] Graceful shutdown completed"); |
| 101 | + std::process::exit(0); |
| 102 | + }, |
| 103 | + } |
| 104 | + }); |
| 105 | + |
| 106 | + // TODO: add biased! to always poll the handler future first, once supported: |
| 107 | + // https://github.com/tokio-rs/tokio/issues/7304 |
| 108 | + tokio::try_join!(lambda_runtime::run(service_fn(handler)), extension.run(),)?; |
| 109 | + |
| 110 | + Ok(()) |
| 111 | +} |
0 commit comments