feat: ⚡ Improved Logging when Lambda Runtime Api fails, this will be useful for the customer for timeouts

Author: darklight3it

lambda-runtime/src/layers/api_client.rs

@@ -1,5 +1,7 @@
 use crate::LambdaInvocation;
 use futures::{future::BoxFuture, ready, FutureExt, TryFutureExt};
+use http::StatusCode;
+use http_body_util::BodyExt;
 use hyper::body::Incoming;
 use lambda_runtime_api_client::{body::Body, BoxError, Client};
 use pin_project::pin_project;
@@ -40,7 +42,7 @@ where
     fn call(&mut self, req: LambdaInvocation) -> Self::Future {
         let request_fut = self.inner.call(req);
         let client = self.client.clone();
-        RuntimeApiClientFuture::First(request_fut, client)
+        RuntimeApiClientFuture::BuildingRequest(request_fut, client)
     }
 }
 
@@ -58,8 +60,12 @@ where
 
 #[pin_project(project = RuntimeApiClientFutureProj)]
 pub enum RuntimeApiClientFuture<F> {
-    First(#[pin] F, Arc<Client>),
-    Second(#[pin] BoxFuture<'static, Result<http::Response<Incoming>, BoxError>>),
+    /// The inner service is building the HTTP request to send to the Lambda Runtime API.
+    BuildingRequest(#[pin] F, Arc<Client>),
+    /// The HTTP request is in-flight; waiting for the Lambda Runtime API response.
+    SendingRequest(#[pin] BoxFuture<'static, Result<http::Response<Incoming>, BoxError>>),
+    /// The Lambda Runtime API returned a non-2xx response; reading the response body for logging.
+    ReadingErrorBody(StatusCode, #[pin] BoxFuture<'static, Result<Vec<u8>, BoxError>>),
 }
 
 impl<F> Future for RuntimeApiClientFuture<F>
@@ -72,7 +78,7 @@ where
         // NOTE: We loop here to directly poll the second future once the first has finished.
         task::Poll::Ready(loop {
             match self.as_mut().project() {
-                RuntimeApiClientFutureProj::First(fut, client) => match ready!(fut.poll(cx)) {
+                RuntimeApiClientFutureProj::BuildingRequest(fut, client) => match ready!(fut.poll(cx)) {
                     Ok(ok) => {
                         // NOTE: We use 'client.call_boxed' here to obtain a future with static
                         // lifetime. Otherwise, this future would need to be self-referential...
@@ -83,12 +89,127 @@ where
                                 err
                             })
                             .boxed();
-                        self.set(RuntimeApiClientFuture::Second(next_fut));
+                        self.set(RuntimeApiClientFuture::SendingRequest(next_fut));
                     }
                     Err(err) => break Err(err),
                 },
-                RuntimeApiClientFutureProj::Second(fut) => break ready!(fut.poll(cx)).map(|_| ()),
+                RuntimeApiClientFutureProj::SendingRequest(fut) => match ready!(fut.poll(cx)) {
+                    Ok(resp) if !resp.status().is_success() => {
+                        let status = resp.status();
+                        let body_fut = resp
+                            .into_body()
+                            .collect()
+                            .map_ok(|b| b.to_bytes().to_vec())
+                            .map_err(BoxError::from)
+                            .boxed();
+                        self.set(RuntimeApiClientFuture::ReadingErrorBody(status, body_fut));
+                    }
+                    Ok(_) => break Ok(()),
+                    Err(err) => break Err(err),
+                },
+                RuntimeApiClientFutureProj::ReadingErrorBody(status, fut) => {
+                    match ready!(fut.poll(cx)) {
+                        Ok(body) => {
+                            let body_str = String::from_utf8_lossy(&body);
+                            error!(status = %status, body = %body_str, "Lambda Runtime API returned non-200 response");
+                        }
+                        Err(err) => {
+                            error!(status = %status, error = ?err, "Lambda Runtime API returned non-200 response; failed to read body");
+                        }
+                    }
+                    break Ok(());
+                }
             }
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::Context;
+    use bytes::Bytes;
+    use httpmock::prelude::*;
+    use std::future;
+    use tower::Service;
+
+    fn make_invocation() -> LambdaInvocation {
+        let (parts, _) = http::Response::new(()).into_parts();
+        LambdaInvocation {
+            parts,
+            body: Bytes::new(),
+            context: Context::default(),
+        }
+    }
+
+    fn make_client(server: &MockServer) -> Arc<Client> {
+        let uri = format!("http://{}:{}", server.host(), server.port()).parse().unwrap();
+        Arc::new(Client::builder().with_endpoint(uri).build().unwrap())
+    }
+
+    /// Builds a mock inner service that always returns the given request.
+    fn mock_inner(
+        req: http::Request<Body>,
+    ) -> impl Service<
+        LambdaInvocation,
+        Response = http::Request<Body>,
+        Error = BoxError,
+        Future = impl Future<Output = Result<http::Request<Body>, BoxError>>,
+    > {
+        tower::service_fn(move |_inv: LambdaInvocation| {
+            let req = http::Request::builder()
+                .uri(req.uri().clone())
+                .method(req.method().clone())
+                .body(Body::empty())
+                .unwrap();
+            future::ready(Ok::<_, BoxError>(req))
+        })
+    }
+
+    #[tokio::test]
+    async fn test_2xx_response_succeeds() {
+        let server = MockServer::start();
+        server.mock(|when, then| {
+            when.any_request();
+            then.status(200).body("ok");
+        });
+
+        let client = make_client(&server);
+        let req = http::Request::builder().uri("/some/path").body(Body::empty()).unwrap();
+        let mut svc = RuntimeApiClientService::new(mock_inner(req), client);
+
+        let result = svc.call(make_invocation()).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_non_2xx_response_is_logged_and_succeeds() {
+        let server = MockServer::start();
+        server.mock(|when, then| {
+            when.any_request();
+            then.status(500).body("internal server error");
+        });
+
+        let client = make_client(&server);
+        let req = http::Request::builder().uri("/some/path").body(Body::empty()).unwrap();
+        let mut svc = RuntimeApiClientService::new(mock_inner(req), client);
+
+        // Non-2xx should still resolve Ok(()) — the error is logged, not propagated.
+        let result = svc.call(make_invocation()).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_inner_service_error_propagates() {
+        let server = MockServer::start();
+        let client = make_client(&server);
+
+        let failing_inner = tower::service_fn(|_: LambdaInvocation| {
+            future::ready(Err::<http::Request<Body>, BoxError>("inner error".into()))
+        });
+        let mut svc = RuntimeApiClientService::new(failing_inner, client);
+
+        let result = svc.call(make_invocation()).await;
+        assert!(result.is_err());
+    }
+}