smithy-lang · jjant · Apr 24, 2023 · Apr 4, 2023 · Apr 4, 2023 · Apr 4, 2023
@@ -8,8 +8,8 @@ mod plugin;
 use std::{net::SocketAddr, sync::Arc};
 
 use aws_smithy_http_server::{
-    extension::OperationExtensionExt, instrumentation::InstrumentExt, plugin::PluginPipeline,
-    request::request_id::ServerRequestIdProviderLayer, AddExtensionLayer,
+    extension::OperationExtensionExt, instrumentation::InstrumentExt, plugin::CheckHealthLayer,
+    plugin::PluginPipeline, request::request_id::ServerRequestIdProviderLayer, AddExtensionLayer,
 };
 use clap::Parser;
 
@@ -63,6 +63,7 @@ pub async fn main() {
     let app = app
         // Setup shared state and middlewares.
         .layer(&AddExtensionLayer::new(Arc::new(State::default())))
+        .layer(&CheckHealthLayer::with_default_handler())
         // Add request IDs
         .layer(&ServerRequestIdProviderLayer::new());
 

diff --git a/rust-runtime/aws-smithy-http-server/src/plugin/check_health.rs b/rust-runtime/aws-smithy-http-server/src/plugin/check_health.rs
@@ -0,0 +1,96 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use std::future::Ready;
+
+use futures_util::Future;
+use http::StatusCode;
+use hyper::{Body, Request, Response};
+use tower::Layer;
+use tower::Service;
+
+use crate::body;
+use crate::body::BoxBody;
+
+use super::Either;
+
+#[derive(Clone)]
+pub struct CheckHealthLayer<H> {
+    ping_handler: H,
+}
+
+impl<H> CheckHealthLayer<H> {
+    pub fn new(ping_handler: H) -> Self {
+        CheckHealthLayer { ping_handler }
+    }
+}
+
+pub type DefaultHandler<E> = fn(Request<Body>) -> Ready<Result<Response<BoxBody>, E>>;
+
+impl CheckHealthLayer<()> {
+    pub fn with_default_handler<E>() -> CheckHealthLayer<DefaultHandler<E>> {
+        CheckHealthLayer::new(default_ping_handler)
+    }
+}
+
+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(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct CheckHealthService<H, S> {
+    inner: S,
+    layer: CheckHealthLayer<H>,
+}
+
+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 = Result<Response<BoxBody>, S::Error>>,
+    H: Fn(Request<Body>) -> HandlerFuture,
+{
+    type Response = S::Response;
+
+    type Error = S::Error;
+
+    type Future = Either<HandlerFuture, S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Request<Body>) -> Self::Future {
+        let clone = self.inner.clone();
+        let mut service = std::mem::replace(&mut self.inner, clone);
+
+        if req.uri() == "/ping" {
+            Either::Left {
+                value: (self.layer.ping_handler)(req),
+            }
+        } else {
+            Either::Right {
+                value: service.call(req),
+            }
+        }
+    }
+}
+
+/// A handler that returns `200 OK` with an empty body.
+fn default_ping_handler<E>(_req: Request<Body>) -> Ready<Result<Response<BoxBody>, E>> {
+    let response = Response::builder()
+        .status(StatusCode::OK)
+        .body(body::boxed(Body::empty()))
+        .expect("Couldn't construct response");
+
+    std::future::ready(Ok::<_, E>(response))
+}
@@ -118,6 +118,7 @@
 //! ```
 //!
 
+mod check_health;
 mod closure;
 mod either;
 mod filter;
@@ -128,6 +129,7 @@ mod stack;
 
 use crate::operation::Operation;
 
+pub use check_health::CheckHealthLayer;
 pub use closure::{plugin_from_operation_name_fn, OperationNameFn};
 pub use either::Either;
 pub use filter::{filter_by_operation_name, FilterByOperationName};