Undeprecate `OperationExtension` and add `Plugin` (#2157)

// This program is exported as a binary named `pokemon-service`.

use std::{net::SocketAddr, sync::Arc};

use aws_smithy_http_server::{plugin::PluginPipeline, AddExtensionLayer};

use aws_smithy_http_server::{extension::OperationExtensionExt, plugin::PluginPipeline, AddExtensionLayer};

use clap::Parser;

use pokemon_service::{

    capture_pokemon, check_health, do_nothing, get_pokemon_species, get_server_statistics, get_storage,

pub async fn main() {

    let args = Args::parse();

    setup_tracing();

    let plugins = PluginPipeline::new()

        // Apply the `PrintPlugin` defined in `plugin.rs`

    let plugins = PluginPipeline::new().print();

        .print()

        // Apply the `OperationExtensionPlugin` defined in `aws_smithy_http_server::extension`. This allows other

        // plugins or tests to access a `aws_smithy_http_server::extension::OperationExtension` from

        // `Response::extensions`, or infer routing failure when it's missing.

        .insert_operation_extension();

    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

//!

//! [extensions]: https://docs.rs/http/latest/http/struct.Extensions.html

use std::ops::Deref;

use std::{fmt, future::Future, ops::Deref, pin::Pin, task::Context, task::Poll};

use futures_util::ready;

use futures_util::TryFuture;

use thiserror::Error;

use tower::{layer::util::Stack, Layer, Service};

use crate::operation::{Operation, OperationShape};

use crate::plugin::{plugin_from_operation_name_fn, OperationNameFn, Plugin, PluginPipeline, PluginStack};

#[allow(deprecated)]

use crate::request::RequestParts;

pub use crate::request::extension::Extension;

pub use crate::request::extension::MissingExtension;

pub use crate::request::extension::{Extension, MissingExtension};

/// Extension type used to store information about Smithy operations in HTTP responses.

/// This extension type is set when it has been correctly determined that the request should be

/// routed to a particular operation. The operation handler might not even get invoked because the

/// request fails to deserialize into the modeled operation input.

/// This extension type is inserted, via the [`OperationExtensionPlugin`], whenever it has been correctly determined

/// that the request should be routed to a particular operation. The operation handler might not even get invoked

/// because the request fails to deserialize into the modeled operation input.

///

/// The format given must be the absolute shape ID with `#` replaced with a `.`.

#[derive(Debug, Clone)]

#[deprecated(

    since = "0.52.0",

    note = "This is no longer inserted by the new service builder. Layers should be constructed per operation using the plugin system."

)]

pub struct OperationExtension {

    absolute: &'static str,

    }

}

pin_project_lite::pin_project! {

    /// The [`Service::Future`] of [`OperationExtensionService`] - inserts an [`OperationExtension`] into the

    /// [`http::Response]`.

    pub struct OperationExtensionFuture<Fut> {

        #[pin]

        inner: Fut,

        operation_extension: Option<OperationExtension>

    }

}

impl<Fut, RespB> Future for OperationExtensionFuture<Fut>

where

    Fut: TryFuture<Ok = http::Response<RespB>>,

{

    type Output = Result<http::Response<RespB>, Fut::Error>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {

        let this = self.project();

        let ext = this

            .operation_extension

            .take()

            .expect("futures cannot be polled after completion");

        let resp = ready!(this.inner.try_poll(cx));

        Poll::Ready(resp.map(|mut resp| {

            resp.extensions_mut().insert(ext);

            resp

        }))

    }

}

/// Inserts a [`OperationExtension`] into the extensions of the [`http::Response`].

#[derive(Debug, Clone)]

pub struct OperationExtensionService<S> {

    inner: S,

    operation_extension: OperationExtension,

}

impl<S, B, RespBody> Service<http::Request<B>> for OperationExtensionService<S>

where

    S: Service<http::Request<B>, Response = http::Response<RespBody>>,

{

    type Response = http::Response<RespBody>;

    type Error = S::Error;

    type Future = OperationExtensionFuture<S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

        self.inner.poll_ready(cx)

    }

    fn call(&mut self, req: http::Request<B>) -> Self::Future {

        OperationExtensionFuture {

            inner: self.inner.call(req),

            operation_extension: Some(self.operation_extension.clone()),

        }

    }

}

/// A [`Layer`] applying the [`OperationExtensionService`] to an inner [`Service`].

#[derive(Debug, Clone)]

pub struct OperationExtensionLayer(OperationExtension);

impl<S> Layer<S> for OperationExtensionLayer {

    type Service = OperationExtensionService<S>;

    fn layer(&self, inner: S) -> Self::Service {

        OperationExtensionService {

            inner,

            operation_extension: self.0.clone(),

        }

    }

}

/// A [`Plugin`] which applies [`OperationExtensionLayer`] to every operation.

pub struct OperationExtensionPlugin(OperationNameFn<fn(&'static str) -> OperationExtensionLayer>);

impl fmt::Debug for OperationExtensionPlugin {

    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

        f.debug_tuple("OperationExtensionPlugin").field(&"...").finish()

    }

}

impl<P, Op, S, L> Plugin<P, Op, S, L> for OperationExtensionPlugin

where

    Op: OperationShape,

{

    type Service = S;

    type Layer = Stack<L, OperationExtensionLayer>;

    fn map(&self, input: Operation<S, L>) -> Operation<Self::Service, Self::Layer> {

        <OperationNameFn<fn(&'static str) -> OperationExtensionLayer> as Plugin<P, Op, S, L>>::map(&self.0, input)

    }

}

/// An extension trait on [`PluginPipeline`] allowing the application of [`OperationExtensionPlugin`].

///

/// See [`module`](crate::extension) documentation for more info.

pub trait OperationExtensionExt<P> {

    /// Apply the [`OperationExtensionPlugin`], which inserts the [`OperationExtension`] into every [`http::Response`].

    fn insert_operation_extension(self) -> PluginPipeline<PluginStack<OperationExtensionPlugin, P>>;

}

impl<P> OperationExtensionExt<P> for PluginPipeline<P> {

    fn insert_operation_extension(self) -> PluginPipeline<PluginStack<OperationExtensionPlugin, P>> {

        let plugin = OperationExtensionPlugin(plugin_from_operation_name_fn(|name| {

            let operation_extension = OperationExtension::new(name).expect("Operation name is malformed, this should never happen. Please file an issue against https://github.com/awslabs/smithy-rs");

            OperationExtensionLayer(operation_extension)

        }));

        self.push(plugin)

    }

}

/// Extension type used to store the type of user-modeled error returned by an operation handler.

/// These are modeled errors, defined in the Smithy model.

#[derive(Debug, Clone)]

}

#[cfg(test)]

#[allow(deprecated)]

mod tests {

    use super::*;