Add AlbHealthCheckLayer

examples/pokemon-service/Cargo.toml

@@ -24,6 +24,9 @@ async-stream = "0.3"
 rand = "0.8.5"
 serial_test = "1.0.0"
 
+# We use hyper client on tests
+hyper = {version = "0.14.25", features = ["server", "client"] }
+
 # This dependency is only required for testing the `pokemon-service-tls` program.
 hyper-rustls = { version = "0.23.2", features = ["http2"] }
 

examples/pokemon-service/src/main.rs

@@ -8,11 +8,16 @@ mod plugin;
 use std::{net::SocketAddr, sync::Arc};
 
 use aws_smithy_http_server::{
-    extension::OperationExtensionExt, instrumentation::InstrumentExt, plugin::PluginPipeline,
-    request::request_id::ServerRequestIdProviderLayer, AddExtensionLayer,
+    body,
+    extension::OperationExtensionExt,
+    instrumentation::InstrumentExt,
+    plugin::{check_health::CheckHealthLayer, PluginPipeline},
+    request::request_id::ServerRequestIdProviderLayer,
+    AddExtensionLayer,
 };
 use clap::Parser;
 
+use hyper::{Body, Response, StatusCode};
 use plugin::PrintExt;
 use pokemon_service::{
     do_nothing_but_log_request_ids, get_storage_with_local_approved, DEFAULT_ADDRESS, DEFAULT_PORT,
@@ -46,7 +51,17 @@ pub async fn main() {
         // `Response::extensions`, or infer routing failure when it's missing.
         .insert_operation_extension()
         // Adds `tracing` spans and events to the request lifecycle.
-        .instrument();
+        .instrument()
+        // Handle `/ping` health check requests.
+        .http_layer(CheckHealthLayer::new("/ping", |_req| {
+            let response = Response::builder()
+                .status(StatusCode::OK)
+                .body(body::boxed(Body::empty()))
+                .expect("Couldn't construct response");
+
+            std::future::ready(response)
+        }));
+
     let app = PokemonService::builder_with_plugins(plugins)
         // Build a registry containing implementations to all the operations in the service. These
         // are async functions or async closures that take as input the operation's input and

examples/pokemon-service/tests/common/mod.rs

@@ -24,11 +24,14 @@ pub async fn run_server() -> ChildDrop {
     ChildDrop(child)
 }
 
+pub fn base_url() -> String {
+    format!("http://{DEFAULT_ADDRESS}:{DEFAULT_PORT}")
+}
+
 pub fn client() -> Client<DynConnector, DynMiddleware<DynConnector>> {
-    let base_url = format!("http://{DEFAULT_ADDRESS}:{DEFAULT_PORT}");
     let raw_client = Builder::new()
         .rustls_connector(Default::default())
-        .middleware_fn(rewrite_base_url(base_url))
+        .middleware_fn(rewrite_base_url(base_url()))
         .build_dyn();
     let config = Config::builder().build();
     Client::with_config(raw_client, config)

examples/pokemon-service/tests/simple.rs

@@ -81,4 +81,11 @@ async fn simple_integration_test() {
 
     let service_statistics_out = client.get_server_statistics().send().await.unwrap();
     assert_eq!(2, service_statistics_out.calls_count.unwrap());
+
+    let hyper_client = hyper::Client::new();
+    let health_check_url = format!("{}/ping", common::base_url());
+    let health_check_url = hyper::Uri::try_from(health_check_url).unwrap();
+    let result = hyper_client.get(health_check_url).await.unwrap();
+
+    assert_eq!(result.status(), 200);
 }

rust-runtime/aws-smithy-http-server/src/plugin/check_health.rs

@@ -0,0 +1,115 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use std::future::Ready;
+use std::marker::PhantomData;
+use std::task::{Context, Poll};
+
+use futures_util::Future;
+use hyper::{Body, Request, Response};
+use pin_project_lite::pin_project;
+use tower::{util::Oneshot, Layer, Service, ServiceExt};
+
+use crate::body::BoxBody;
+
+use super::Either;
+
+/// A [`tower::Layer`] used to apply [`CheckHealthService`].
+#[derive(Clone, Debug)]
+pub struct CheckHealthLayer<PingHandler> {
+    health_check_uri: &'static str,
+    ping_handler: PingHandler,
+}
+
+impl CheckHealthLayer<()> {
+    pub fn new<HandlerFuture: Future<Output = Response<BoxBody>>, H: Fn(Request<Body>) -> HandlerFuture>(
+        health_check_uri: &'static str,
+        ping_handler: H,
+    ) -> CheckHealthLayer<H> {
+        CheckHealthLayer {
+            health_check_uri,
+            ping_handler,
+        }
+    }
+}
+
+pub type DefaultHandler = fn(Request<Body>) -> Ready<Response<BoxBody>>;
+
+impl<S, H: Clone> Layer<S> for CheckHealthLayer<H> {
+    type Service = CheckHealthService<H, S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        CheckHealthService {
+            inner,
+            layer: self.clone(),
+        }
+    }
+}
+
+/// A middleware [`Service`] responsible for handling health check requests.
+#[derive(Clone, Debug)]
+pub struct CheckHealthService<H, S> {
+    inner: S,
+    layer: CheckHealthLayer<H>,
+}
+
+pin_project! {
+    /// A future that converts `F` into a compatible `S::Future`.
+    pub struct MappedHandlerFuture<R, S, F> {
+        #[pin]
+        inner: F,
+        pd: PhantomData<fn(R) -> S>,
+    }
+}
+
+impl<R, S, F> MappedHandlerFuture<R, S, F> {
+    fn new(inner: F) -> MappedHandlerFuture<R, S, F> {
+        Self { inner, pd: PhantomData }
+    }
+}
+
+impl<R, S: Service<R>, F: Future<Output = S::Response>> Future for MappedHandlerFuture<R, S, F> {
+    type Output = Result<S::Response, S::Error>;
+
+    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.project().inner.poll(cx).map(Ok)
+    }
+}
+
+impl<H, HandlerFuture, S> Service<Request<Body>> for CheckHealthService<H, S>
+where
+    S: Service<Request<Body>, Response = Response<BoxBody>> + Clone,
+    S::Future: std::marker::Send + 'static,
+    HandlerFuture: Future<Output = Response<BoxBody>>,
+    H: Fn(Request<Body>) -> HandlerFuture,
+{
+    type Response = S::Response;
+
+    type Error = S::Error;
+
+    type Future = Either<MappedHandlerFuture<Request<Body>, S, HandlerFuture>, Oneshot<S, Request<Body>>>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        // The check that the service is ready is done by `Oneshot` below.
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, req: Request<Body>) -> Self::Future {
+        if req.uri() == self.layer.health_check_uri {
+            let handler_future = (self.layer.ping_handler)(req);
+
+            Either::Left {
+                value: MappedHandlerFuture::new(handler_future),
+            }
+        } else {
+            let clone = self.inner.clone();
+            let service = std::mem::replace(&mut self.inner, clone);
+
+            Either::Right {
+                value: service.oneshot(req),
+            }
+        }
+    }
+}

rust-runtime/aws-smithy-http-server/src/plugin/mod.rs

@@ -118,6 +118,7 @@
 //! ```
 //!
 
+pub mod check_health;
 mod closure;
 mod either;
 mod filter;

rust-runtime/aws-smithy-http-server/src/plugin/pipeline.rs

@@ -6,6 +6,8 @@
 use crate::operation::Operation;
 use crate::plugin::{IdentityPlugin, Plugin, PluginStack};
 
+use super::HttpLayer;
+
 /// A wrapper struct for composing [`Plugin`]s.
 /// It is used as input for the `builder_with_plugins` method on the generate service struct
 /// (e.g. `PokemonService::builder_with_plugins`).
@@ -168,6 +170,11 @@ impl<P> PluginPipeline<P> {
     pub fn push<NewPlugin>(self, new_plugin: NewPlugin) -> PluginPipeline<PluginStack<NewPlugin, P>> {
         PluginPipeline(PluginStack::new(new_plugin, self.0))
     }
+
+    /// Applies a single [`Layer`] to all operations _before_ they are deserialized.
+    pub fn http_layer<L>(self, layer: L) -> PluginPipeline<PluginStack<HttpLayer<L>, P>> {
+        PluginPipeline(PluginStack::new(HttpLayer(layer), self.0))
+    }
 }
 
 impl<P, Op, S, L, InnerPlugin> Plugin<P, Op, S, L> for PluginPipeline<InnerPlugin>