Allow handler extensions

generated/examples/usage.rs

@@ -1,34 +1,35 @@
 use std::{convert::Infallible, future::Ready, task::Poll};
 
 use generated::{operations::*, services::*, structures::*};
-use runtime::operation::{AdjoinState, OperationError, OperationShapeExt};
+use runtime::operation::{Extension, OperationError, OperationShapeExt};
 use tower::{util::MapResponseLayer, Service};
 
-/// Fallible handler with state
+// Fallible handler with extensions.
 async fn get_pokemon_species_stateful(
     _input: GetPokemonSpeciesInput,
-    _state: usize,
+    _ext_a: Extension<usize>,
+    _ext_b: Extension<String>,
 ) -> Result<GetPokemonSpeciesOutput, ResourceNotFoundException> {
     todo!()
 }
 
-/// Fallible handler without state
+// Fallible handler without extensions.
 async fn get_pokemon_species(
     _input: GetPokemonSpeciesInput,
 ) -> Result<GetPokemonSpeciesOutput, ResourceNotFoundException> {
     todo!()
 }
 
-/// Infallible handler without state
+// Infallible handler without extensions.
 async fn empty_operation(_input: EmptyOperationInput) -> EmptyOperationOutput {
     todo!()
 }
 
-/// Bespoke implementation of `EmptyOperation`.
+// Bespoke implementation of `EmptyOperation`.
 #[derive(Clone)]
-struct EmptyOperationService;
+struct EmptyOperationServiceA;
 
-impl Service<EmptyOperationInput> for EmptyOperationService {
+impl Service<EmptyOperationInput> for EmptyOperationServiceA {
     type Response = EmptyOperationOutput;
     type Error = OperationError<String, Infallible>;
     type Future = Ready<Result<Self::Response, Self::Error>>;
@@ -42,13 +43,31 @@ impl Service<EmptyOperationInput> for EmptyOperationService {
     }
 }
 
+// Bespoke implementation of `EmptyOperation` with an extension.
+#[derive(Clone)]
+struct EmptyOperationServiceB;
+
+impl Service<(EmptyOperationInput, Extension<String>)> for EmptyOperationServiceB {
+    type Response = EmptyOperationOutput;
+    type Error = OperationError<String, Infallible>;
+    type Future = Ready<Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        todo!()
+    }
+
+    fn call(&mut self, _req: (EmptyOperationInput, Extension<String>)) -> Self::Future {
+        todo!()
+    }
+}
+
 fn main() {
     // Various ways of constructing operations
     let _get_pokemon_species = GetPokemonSpecies::from_handler(get_pokemon_species);
     let _empty_operation = EmptyOperation::from_handler(empty_operation);
-    let empty_operation = EmptyOperation::from_service(EmptyOperationService);
-    let get_pokemon_species =
-        GetPokemonSpecies::from_handler(get_pokemon_species_stateful.with_state(29));
+    let _empty_operation = EmptyOperation::from_service(EmptyOperationServiceA);
+    let empty_operation = EmptyOperation::from_service(EmptyOperationServiceB);
+    let get_pokemon_species = GetPokemonSpecies::from_handler(get_pokemon_species_stateful);
 
     // We can apply a layer to them
     let get_pokemon_species = get_pokemon_species.layer(MapResponseLayer::new(|resp| resp));

generated/src/operations/empty_operation.rs

@@ -32,7 +32,7 @@ impl<B> FromRequest<AWSRestJsonV1, EmptyOperation, B> for EmptyOperationInput {
 
     type Future = Ready<Result<Self, Self::Error>>;
 
-    fn from_request(_request: http::Request<B>) -> Self::Future {
+    fn from_request(_request: &mut http::Request<B>) -> Self::Future {
         todo!()
     }
 }

generated/src/operations/get_pokemon_species.rs

@@ -34,7 +34,7 @@ impl<B> FromRequest<AWSRestJsonV1, GetPokemonSpecies, B> for GetPokemonSpeciesIn
 
     type Future = Ready<Result<Self, Self::Error>>;
 
-    fn from_request(_request: http::Request<B>) -> Self::Future {
+    fn from_request(_request: &mut http::Request<B>) -> Self::Future {
         todo!()
     }
 }

generated/src/services/pokemon_service.rs

@@ -112,22 +112,22 @@ impl<Op1, Op2> PokemonServiceBuilder<Op1, Op2> {
 }
 
 impl<S1, S2, L1, L2> PokemonServiceBuilder<Operation<S1, L1>, Operation<S2, L2>> {
-    pub fn build<B>(self) -> PokemonService<RouteService<B>>
+    pub fn build<B, Exts1, Exts2>(self) -> PokemonService<RouteService<B>>
     where
         // GetPokemonSpecies composition
-        UpgradeLayer<AWSRestJsonV1, GetPokemonSpecies, B>: Layer<S1>,
-        L1: Layer<UpgradedService<AWSRestJsonV1, GetPokemonSpecies, B, S1>>,
-        S1: Service<<GetPokemonSpecies as OperationShape>::Input>,
+        UpgradeLayer<AWSRestJsonV1, GetPokemonSpecies, Exts1, B>: Layer<S1>,
+        L1: Layer<UpgradedService<AWSRestJsonV1, GetPokemonSpecies, Exts1, B, S1>>,
+        S1: Service<(<GetPokemonSpecies as OperationShape>::Input, Exts1)>,
         L1::Service: Service<http::Request<B>, Response = http::Response<BoxBody<Bytes, Error>>>,
         L1::Service: Clone + Send + 'static,
         <L1::Service as Service<http::Request<B>>>::Future: Send + 'static,
         <L1::Service as Service<http::Request<B>>>::Error:
             Into<Box<dyn std::error::Error + Send + Sync>> + 'static,
 
         // EmptyOperation composition
-        UpgradeLayer<AWSRestJsonV1, EmptyOperation, B>: Layer<S2>,
-        L2: Layer<UpgradedService<AWSRestJsonV1, EmptyOperation, B, S2>>,
-        S2: Service<<EmptyOperation as OperationShape>::Input>,
+        UpgradeLayer<AWSRestJsonV1, EmptyOperation, Exts2, B>: Layer<S2>,
+        L2: Layer<UpgradedService<AWSRestJsonV1, EmptyOperation, Exts2, B, S2>>,
+        S2: Service<(<EmptyOperation as OperationShape>::Input, Exts2)>,
         L2::Service: Service<http::Request<B>, Response = http::Response<BoxBody<Bytes, Error>>>,
         L2::Service: Clone + Send + 'static,
         <L2::Service as Service<http::Request<B>>>::Future: Send + 'static,

runtime/src/operation/flattened.rs

@@ -0,0 +1,135 @@
+use std::{
+    marker::PhantomData,
+    task::{Context, Poll},
+};
+
+use tower::Service;
+
+use super::{OperationError, OperationShape};
+
+/// A utility trait used to provide an even interface for all operation services.
+///
+/// This serves to take [`Service`]s of the form `Service<(Input, Arg0, Arg1, ...)>` to the canonical representation of
+/// `Service<(Input, (Arg0, Arg1, ...))>` inline with [`IntoService`](super::IntoService).
+pub trait Flattened<Op, Exts, PollError>:
+    Service<Self::Flattened, Response = Op::Output, Error = OperationError<PollError, Op::Error>>
+where
+    Op: OperationShape,
+{
+    type Flattened;
+
+    // Unflatten the request type.
+    fn unflatten(input: Op::Input, exts: Exts) -> Self::Flattened;
+}
+
+// `Service<Op::Input>`
+impl<Op, S, PollError> Flattened<Op, (), PollError> for S
+where
+    Op: OperationShape,
+    S: Service<Op::Input, Response = Op::Output, Error = OperationError<PollError, Op::Error>>,
+{
+    type Flattened = Op::Input;
+
+    fn unflatten(input: Op::Input, _exts: ()) -> Self::Flattened {
+        input
+    }
+}
+
+// `Service<(Op::Input, Arg0)>`
+impl<Op, Arg0, S, PollError> Flattened<Op, (Arg0,), PollError> for S
+where
+    Op: OperationShape,
+    S: Service<
+        (Op::Input, Arg0),
+        Response = Op::Output,
+        Error = OperationError<PollError, Op::Error>,
+    >,
+{
+    type Flattened = (Op::Input, Arg0);
+
+    fn unflatten(input: Op::Input, exts: (Arg0,)) -> Self::Flattened {
+        (input, exts.0)
+    }
+}
+
+// `Service<(Op::Input, Arg0, Arg1)>`
+impl<Op, Arg0, Arg1, S, PollError> Flattened<Op, (Arg0, Arg1), PollError> for S
+where
+    Op: OperationShape,
+    S: Service<
+        (Op::Input, Arg0, Arg1),
+        Response = Op::Output,
+        Error = OperationError<PollError, Op::Error>,
+    >,
+{
+    type Flattened = (Op::Input, Arg0, Arg1);
+
+    fn unflatten(input: Op::Input, exts: (Arg0, Arg1)) -> Self::Flattened {
+        (input, exts.0, exts.1)
+    }
+}
+
+/// An extension trait of [`Flattened`].
+pub trait FlattenedExt<Op, Exts, PollError>: Flattened<Op, Exts, PollError>
+where
+    Op: OperationShape,
+{
+    /// Convert the [`Flattened`] into a canonicalized [`Service`].
+    fn into_unflatten(self) -> IntoUnflattened<Op, Self, PollError>
+    where
+        Self: Sized,
+    {
+        IntoUnflattened {
+            inner: self,
+            _operation: PhantomData,
+            _poll_error: PhantomData,
+        }
+    }
+}
+
+impl<F, Op, Exts, PollError> FlattenedExt<Op, Exts, PollError> for F
+where
+    Op: OperationShape,
+    F: Flattened<Op, Exts, PollError>,
+{
+}
+
+/// A [`Service`] canonicalizing the request type of a [`Flattened`].
+#[derive(Debug)]
+pub struct IntoUnflattened<Op, S, PollError> {
+    inner: S,
+    _operation: PhantomData<Op>,
+    _poll_error: PhantomData<PollError>,
+}
+
+impl<Op, S, PollError> Clone for IntoUnflattened<Op, S, PollError>
+where
+    S: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            _operation: PhantomData,
+            _poll_error: PhantomData,
+        }
+    }
+}
+
+impl<Op, S, Exts, PollError> Service<(Op::Input, Exts)> for IntoUnflattened<Op, S, PollError>
+where
+    Op: OperationShape,
+    S: Flattened<Op, Exts, PollError>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = <S as Service<S::Flattened>>::Future;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, (input, exts): (Op::Input, Exts)) -> Self::Future {
+        let req = S::unflatten(input, exts);
+        self.inner.call(req)
+    }
+}