@@ -32,6 +32,7 @@ pub mod streaming;
3232
3333/// Utilities to initialize and use `tracing` and `tracing-subscriber` in Lambda Functions.
3434#[ cfg( feature = "tracing" ) ]
35+ #[ cfg_attr( docsrs, doc( cfg( feature = "tracing" ) ) ) ]
3536pub use lambda_runtime_api_client:: tracing;
3637
3738/// Types available to a Lambda function.
@@ -123,3 +124,110 @@ where
123124 let runtime = Runtime :: new ( handler) . layer ( layers:: TracingLayer :: new ( ) ) ;
124125 runtime. run ( ) . await
125126}
127+
128+ /// Spawns a task that will be execute a provided async closure when the process
129+ /// receives unix graceful shutdown signals. If the closure takes longer than 500ms
130+ /// to execute, an unhandled `SIGKILL` signal might be received.
131+ ///
132+ /// You can use this future to execute cleanup or flush related logic prior to runtime shutdown.
133+ ///
134+ /// This function must be called prior to [lambda_runtime::run()].
135+ ///
136+ /// Note that this implicitly also registers and drives a no-op internal extension that subscribes to no events.
137+ /// This extension will be named `_lambda-rust-runtime-no-op-graceful-shutdown-helper`. This extension name
138+ /// can not be reused by other registered extensions. This is necessary in order to receive graceful shutdown signals.
139+ ///
140+ /// This extension is cheap to run because it receives no events, but is not zero cost. If you have another extension
141+ /// registered already, you might prefer to manually construct your own graceful shutdown handling without the dummy extension.
142+ ///
143+ /// For more information on general AWS Lambda graceful shutdown handling, see:
144+ /// https://github.com/aws-samples/graceful-shutdown-with-aws-lambda
145+ ///
146+ /// # Panics
147+ ///
148+ /// This function panics if:
149+ /// - this function is called after [lambda_runtime::run()]
150+ /// - this function is called outside of a context that has access to the tokio i/o
151+ /// - the no-op extension cannot be registered
152+ /// - either signal listener panics [tokio::signal::unix](https://docs.rs/tokio/latest/tokio/signal/unix/fn.signal.html#errors)
153+ ///
154+ /// # Example
155+ /// ```no_run
156+ /// use lambda_runtime::{Error, service_fn, LambdaEvent};
157+ /// use serde_json::Value;
158+ ///
159+ /// #[tokio::main]
160+ /// async fn main() -> Result<(), Error> {
161+ /// let func = service_fn(func);
162+ ///
163+ /// let (writer, log_guard) = tracing_appender::non_blocking(std::io::stdout());
164+ /// lambda_runtime::tracing::init_default_subscriber_with_writer(writer);
165+ ///
166+ /// let shutdown_hook = || async move {
167+ /// std::mem::drop(log_guard);
168+ /// };
169+ /// lambda_runtime::spawn_graceful_shutdown_handler(shutdown_hook);
170+ ///
171+ /// lambda_runtime::run(func).await?;
172+ /// Ok(())
173+ /// }
174+ ///
175+ /// async fn func(event: LambdaEvent<Value>) -> Result<Value, Error> {
176+ /// Ok(event.payload)
177+ /// }
178+ /// ```
179+ #[ cfg( all( unix, feature = "graceful-shutdown" ) ) ]
180+ #[ cfg_attr( docsrs, doc( cfg( all( unix, feature = "tokio-rt" ) ) ) ) ]
181+ pub fn spawn_graceful_shutdown_handler < Fut > ( shutdown_hook : impl FnOnce ( ) -> Fut + Send + ' static )
182+ where
183+ Fut : Future < Output = ( ) > + Send + ' static ,
184+ {
185+ tokio:: task:: spawn ( async move {
186+ // You need an extension registered with the Lambda orchestrator in order for your process
187+ // to receive a SIGTERM for graceful shutdown.
188+ //
189+ // We accomplish this here by registering a no-op internal extension, which does not subscribe to any events.
190+ //
191+ // This extension is cheap to run since after it connects to the lambda orchestration, the connection
192+ // will just wait forever for data to come, which never comes, so it won't cause wakes.
193+ let extension = lambda_extension:: Extension :: new ( )
194+ // Don't subscribe to any event types
195+ . with_events ( & [ ] )
196+ // Internal extension names MUST be unique within a given Lambda function.
197+ . with_extension_name ( "_lambda-rust-runtime-no-op-graceful-shutdown-helper" )
198+ // Extensions MUST be registered before calling lambda_runtime::run(), which ends the Init
199+ // phase and begins the Invoke phase.
200+ . register ( )
201+ . await
202+ . expect ( "could not register no-op extension for graceful shutdown" ) ;
203+
204+ let graceful_shutdown_future = async move {
205+ let mut sigint = tokio:: signal:: unix:: signal ( tokio:: signal:: unix:: SignalKind :: interrupt ( ) ) . unwrap ( ) ;
206+ let mut sigterm = tokio:: signal:: unix:: signal ( tokio:: signal:: unix:: SignalKind :: terminate ( ) ) . unwrap ( ) ;
207+ tokio:: select! {
208+ _sigint = sigint. recv( ) => {
209+ eprintln!( "[runtime] SIGINT received" ) ;
210+ eprintln!( "[runtime] Graceful shutdown in progress ..." ) ;
211+ shutdown_hook( ) . await ;
212+ eprintln!( "[runtime] Graceful shutdown completed" ) ;
213+ std:: process:: exit( 0 ) ;
214+ } ,
215+ _sigterm = sigterm. recv( ) => {
216+ eprintln!( "[runtime] SIGTERM received" ) ;
217+ eprintln!( "[runtime] Graceful shutdown in progress ..." ) ;
218+ shutdown_hook( ) . await ;
219+ eprintln!( "[runtime] Graceful shutdown completed" ) ;
220+ std:: process:: exit( 0 ) ;
221+ } ,
222+ }
223+ } ;
224+
225+ // TODO: add biased! to always poll the signal handling future first, once supported:
226+ // https://github.com/tokio-rs/tokio/issues/7304
227+ let _: ( _ , ( ) ) = tokio:: join!( graceful_shutdown_future, async {
228+ // we suppress extension errors because we don't actually mind if it crashes,
229+ // all we need to do is kick off the run so that lambda exits the init phase
230+ let _ = extension. run( ) . await ;
231+ } ) ;
232+ } ) ;
233+ }
0 commit comments