@@ -1,5 +1,11 @@
-use crate::{RefRuntimeState, error::ServerError, requests::*};
-use axum::{Json, body::Body, extract::State, http::Request, response::Response};
+use crate::{RefRuntimeState, error::ServerError, requests::*, state::ExtensionType};
+use axum::{
+    Json,
+    body::Body,
+    extract::{Path, State},
+    http::Request,
+    response::Response,
+};
 use http_body_util::BodyExt;
 use hyper::HeaderMap;
 use serde::{Serialize, de::DeserializeOwned};
@@ -19,6 +25,7 @@ struct RegisterResponse {
 
 pub(crate) async fn register_extension(
     State(state): State<RefRuntimeState>,
+    function_name: Option<Path<String>>,
     req: Request<Body>,
 ) -> Result<Response<Body>, ServerError> {
     let response_body = serde_json::to_vec(&RegisterResponse {
@@ -40,10 +47,20 @@ pub(crate) async fn register_extension(
         )),
     })?;
 
+    // we know that internal extensions are registered from within a function,
+    // therefore have the function name prefixing their path
+    let extension_type = match function_name {
+        Some(_) => ExtensionType::Internal,
+        _ => ExtensionType::External,
+    };
+
     let payload: EventsRequest = extract_json(req).await?;
     debug!(?payload, "registering extension");
 
-    let extension_id = state.ext_cache.register(payload.events).await;
+    let extension_id = state
+        .ext_cache
+        .register(payload.events, extension_type)
+        .await;
     let resp = Response::builder()
         .status(200)
         .header(EXTENSION_ID_HEADER, extension_id)
 
@@ -17,21 +17,17 @@ pub(crate) fn routes() -> Router<RefRuntimeState> {
     Router::new()
         .route("/2020-01-01/extension/register", post(register_extension))
         // secondary route is for internal extensions
-        // that have the function name in their path
-        // (which we just discard currently)
         .route(
-            "/:_function_name/2020-01-01/extension/register",
+            "/:function_name/2020-01-01/extension/register",
             post(register_extension),
         )
         .route(
             "/2020-01-01/extension/event/next",
             get(next_extension_event),
         )
         // secondary route is for internal extensions
-        // that have the function name in their path
-        // (which we just discard currently)
         .route(
-            "/:_function_name/2020-01-01/extension/event/next",
+            "/:function_name/2020-01-01/extension/event/next",
             get(next_extension_event),
         )
         .route("/2020-08-15/logs", put(subcribe_extension_events))
 
@@ -10,7 +10,11 @@ use std::{
     collections::{HashMap, HashSet, hash_map::Entry},
     net::SocketAddr,
     path::PathBuf,
-    sync::Arc,
+    sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    },
+    time::Duration,
 };
 use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
 use tracing::debug;
@@ -226,14 +230,20 @@ pub(crate) struct ExtensionCache {
     extensions: Arc<Mutex<HashMap<String, Vec<String>>>>,
     events: Arc<Mutex<HashMap<String, Vec<String>>>>,
     senders: Arc<Mutex<HashMap<String, mpsc::Sender<NextEvent>>>>,
+    has_internal_extension: Arc<AtomicBool>,
+    has_external_extension: Arc<AtomicBool>,
 }
 
 impl ExtensionCache {
-    pub async fn register(&self, events: Vec<String>) -> String {
+    pub async fn register(&self, events: Vec<String>, extension_type: ExtensionType) -> String {
         let mut extensions = self.extensions.lock().await;
         let extension_id = Uuid::new_v4();
 
         extensions.insert(extension_id.to_string(), events.clone());
+        match extension_type {
+            ExtensionType::Internal => self.has_internal_extension.store(true, Ordering::Relaxed),
+            ExtensionType::External => self.has_external_extension.store(true, Ordering::Relaxed),
+        };
 
         let mut list = self.events.lock().await;
         for event in events {
@@ -287,4 +297,28 @@ impl ExtensionCache {
             }
         }
     }
+
+    /// Returns the allowed duration between SIGTERM and SIGKILL for the function,
+    /// if any. This depends on whether there are any internal or external functions registered.
+    ///
+    /// We don't currently simulate extension shutdown delay.
+    ///
+    /// Ref:
+    /// - https://github.com/aws-samples/graceful-shutdown-with-aws-lambda
+    /// - https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html#runtimes-lifecycle-shutdown
+    pub async fn function_shutdown_delay(&self) -> Option<Duration> {
+        if self.has_internal_extension.load(Ordering::Relaxed) {
+            Some(Duration::from_millis(500))
+        } else if self.has_external_extension.load(Ordering::Relaxed) {
+            Some(Duration::from_millis(300))
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum ExtensionType {
+    Internal,
+    External,
 }
@@ -3,7 +3,6 @@ use cargo_lambda_metadata::{
     cargo::load_metadata,
     config::{ConfigOptions, load_config_without_cli_flags},
 };
-// use cargo_lambda_metadata::cargo::function_environment_metadata;
 use ignore::create_filter;
 use ignore_files::IgnoreFile;
 use std::{collections::HashMap, convert::Infallible, path::PathBuf, sync::Arc, time::Duration};
@@ -128,11 +127,28 @@ async fn runtime(
         let ext_cache = ext_cache.clone();
         async move {
             if signals.contains(&MainSignal::Terminate) {
+                let function_shutdown_delay = ext_cache.function_shutdown_delay().await;
+                if let Some(delay) = function_shutdown_delay {
+                    function_graceful_shutdown_or_else_sigkill(
+                        action,
+                        MainSignal::Terminate,
+                        delay,
+                    );
+                    return Ok(());
+                }
                 action.outcome(Outcome::both(Outcome::Stop, Outcome::Exit));
                 return Ok(());
             }
-
             if signals.contains(&MainSignal::Interrupt) {
+                let function_shutdown_delay = ext_cache.function_shutdown_delay().await;
+                if let Some(delay) = function_shutdown_delay {
+                    function_graceful_shutdown_or_else_sigkill(
+                        action,
+                        MainSignal::Interrupt,
+                        delay,
+                    );
+                    return Ok(());
+                }
                 action.outcome(Outcome::both(Outcome::Stop, Outcome::Exit));
                 return Ok(());
             }
@@ -211,6 +227,32 @@ async fn runtime(
     Ok(config)
 }
 
+fn function_graceful_shutdown_or_else_sigkill(
+    action: Action,
+    signal_type: MainSignal,
+    max_delay: Duration,
+) {
+    tracing::debug!(
+        ?signal_type,
+        ?max_delay,
+        "attempting graceful function shutdown"
+    );
+    action.outcome(Outcome::both(
+        // send sigterm
+        Outcome::Signal(signal_type),
+        // race graceful shutdown against forced shutdown following max delay
+        Outcome::race(
+            // happy path, process exits then watchexec exits
+            Outcome::both(Outcome::Wait, Outcome::Exit),
+            // unhappy path, we sleep max delay then SIGKILL and exit watchexec
+            Outcome::both(
+                Outcome::Sleep(max_delay),
+                Outcome::both(Outcome::Stop, Outcome::Exit),
+            ),
+        ),
+    ));
+}
+
 fn reload_env(manifest_path: &PathBuf, bin_name: &Option<String>) -> HashMap<String, String> {
     let metadata = match load_metadata(manifest_path) {
         Ok(metadata) => metadata,
 
@@ -275,6 +275,24 @@ The following video shows you how to use the watch subcommand with Lambda extens
 
 <iframe width="560" height="315" src="https://www.youtube.com/embed/z2sv41ukHTE" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen></iframe>
 
+## Graceful shutdown support
+
+The watch subcommand will propagate `SIGTERM` and `SIGINT` signals to your function, with a delay to allow for graceful shutdown logic to fire.
+
+This makes a best-effort attempt to replicate signal handling in AWS Lambda, including allowed shutdown delay, but does not strictly match AWS orchestration. Note that graceful shutdown is only enabled in AWS if extensions are registered. This means that if your function has no registered extensions, it will only receive a `SIGKILL`.
+
+One easy way to trigger a graceful shutdown in a Posix shell is by using `Control-C` in the session
+running your `cargo lambda watch` process.
+
+Some caveats:
+- Graceful shutdown signals are only propagated to the function's process, not that of external extensions.
+- Graceful shutdown signals are not sent properly on hot reload, only direct signaling to the `cargo watch` process.
+- No guarantees are made around sequencing of runtime shutdown and extensions receiving `SHUTDOWN` events.
+
+For more information on graceful shutdown handling, refer to:
+- [AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html#runtimes-lifecycle-shutdown)
+- [AWS Samples repository](https://github.com/aws-samples/graceful-shutdown-with-aws-lambda)
+
 ## TLS support
 
 The watch subcommand supports TLS connections to the runtime if you want to send requests to the runtime securely.