Always store state in an Arc (#1270)

* Add extension and state benchmarks

* wip

* Arc the state everywhere

* don't require `S: Clone`

* fix example

axum/src/docs/extract.md

@@ -424,7 +424,7 @@ struct ExtractUserAgent(HeaderValue);
 impl<S, B> FromRequest<S, B> for ExtractUserAgent
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = (StatusCode, &'static str);
 
@@ -476,7 +476,7 @@ struct AuthenticatedUser {
 impl<S, B> FromRequest<S, B> for AuthenticatedUser
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = Response;
 

axum/src/extension.rs

@@ -77,7 +77,7 @@ impl<T, S, B> FromRequest<S, B> for Extension<T>
 where
     T: Clone + Send + Sync + 'static,
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = ExtensionRejection;
 

axum/src/extract/connect_info.rs

@@ -131,7 +131,7 @@ pub struct ConnectInfo<T>(pub T);
 impl<S, B, T> FromRequest<S, B> for ConnectInfo<T>
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
     T: Clone + Send + Sync + 'static,
 {
     type Rejection = <Extension<Self> as FromRequest<S, B>>::Rejection;

axum/src/extract/content_length_limit.rs

@@ -41,7 +41,7 @@ where
     T: FromRequest<S, B>,
     T::Rejection: IntoResponse,
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = ContentLengthLimitRejection<T::Rejection>;
 

axum/src/extract/host.rs

@@ -24,7 +24,7 @@ pub struct Host(pub String);
 impl<S, B> FromRequest<S, B> for Host
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = HostRejection;
 

axum/src/extract/matched_path.rs

@@ -67,7 +67,7 @@ impl MatchedPath {
 impl<S, B> FromRequest<S, B> for MatchedPath
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = MatchedPathRejection;
 

axum/src/extract/multipart.rs

@@ -54,7 +54,7 @@ impl<S, B> FromRequest<S, B> for Multipart
 where
     B: HttpBody<Data = Bytes> + Default + Unpin + Send + 'static,
     B::Error: Into<BoxError>,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = MultipartRejection;
 

axum/src/extract/path/mod.rs

@@ -167,7 +167,7 @@ impl<T, S, B> FromRequest<S, B> for Path<T>
 where
     T: DeserializeOwned + Send,
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = PathRejection;
 

axum/src/extract/query.rs

@@ -53,7 +53,7 @@ impl<T, S, B> FromRequest<S, B> for Query<T>
 where
     T: DeserializeOwned,
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = QueryRejection;
 

axum/src/extract/raw_query.rs

@@ -30,7 +30,7 @@ pub struct RawQuery(pub Option<String>);
 impl<S, B> FromRequest<S, B> for RawQuery
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = Infallible;
 

axum/src/extract/request_parts.rs

@@ -89,7 +89,7 @@ pub struct OriginalUri(pub Uri);
 impl<S, B> FromRequest<S, B> for OriginalUri
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = Infallible;
 
@@ -146,7 +146,7 @@ where
     B: HttpBody + Send + 'static,
     B::Data: Into<Bytes>,
     B::Error: Into<BoxError>,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = BodyAlreadyExtracted;
 
@@ -201,7 +201,7 @@ pub struct RawBody<B = Body>(pub B);
 impl<S, B> FromRequest<S, B> for RawBody<B>
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = BodyAlreadyExtracted;
 

axum/src/extract/state.rs

@@ -153,7 +153,7 @@ use std::{
 ///     // keep `S` generic but require that it can produce a `MyLibraryState`
 ///     // this means users will have to implement `FromRef<UserState> for MyLibraryState`
 ///     MyLibraryState: FromRef<S>,
-///     S: Send,
+///     S: Send + Sync,
 /// {
 ///     type Rejection = Infallible;
 ///
@@ -182,7 +182,7 @@ impl<B, OuterState, InnerState> FromRequest<OuterState, B> for State<InnerState>
 where
     B: Send,
     InnerState: FromRef<OuterState>,
-    OuterState: Send,
+    OuterState: Send + Sync,
 {
     type Rejection = Infallible;
 

axum/src/extract/ws.rs

@@ -278,7 +278,7 @@ impl WebSocketUpgrade {
 impl<S, B> FromRequest<S, B> for WebSocketUpgrade
 where
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = WebSocketUpgradeRejection;
 

axum/src/form.rs

@@ -62,7 +62,7 @@ where
     B: HttpBody + Send,
     B::Data: Send,
     B::Error: Into<BoxError>,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = FormRejection;
 

axum/src/handler/into_service.rs

@@ -5,6 +5,7 @@ use std::{
     convert::Infallible,
     fmt,
     marker::PhantomData,
+    sync::Arc,
     task::{Context, Poll},
 };
 use tower_service::Service;
@@ -16,7 +17,7 @@ use tower_service::Service;
 /// [`HandlerWithoutStateExt::into_service`]: super::HandlerWithoutStateExt::into_service
 pub struct IntoService<H, T, S, B> {
     handler: H,
-    state: S,
+    state: Arc<S>,
     _marker: PhantomData<fn() -> (T, B)>,
 }
 
@@ -35,7 +36,7 @@ fn traits() {
 }
 
 impl<H, T, S, B> IntoService<H, T, S, B> {
-    pub(super) fn new(handler: H, state: S) -> Self {
+    pub(super) fn new(handler: H, state: Arc<S>) -> Self {
         Self {
             handler,
             state,
@@ -55,12 +56,11 @@ impl<H, T, S, B> fmt::Debug for IntoService<H, T, S, B> {
 impl<H, T, S, B> Clone for IntoService<H, T, S, B>
 where
     H: Clone,
-    S: Clone,
 {
     fn clone(&self) -> Self {
         Self {
             handler: self.handler.clone(),
-            state: self.state.clone(),
+            state: Arc::clone(&self.state),
             _marker: PhantomData,
         }
     }
@@ -70,7 +70,7 @@ impl<H, T, S, B> Service<Request<B>> for IntoService<H, T, S, B>
 where
     H: Handler<T, S, B> + Clone + Send + 'static,
     B: Send + 'static,
-    S: Clone,
+    S: Send + Sync,
 {
     type Response = Response;
     type Error = Infallible;
@@ -88,7 +88,7 @@ where
         use futures_util::future::FutureExt;
 
         let handler = self.handler.clone();
-        let future = Handler::call(handler, self.state.clone(), req);
+        let future = Handler::call(handler, Arc::clone(&self.state), req);
         let future = future.map(Ok as _);
 
         super::future::IntoServiceFuture::new(future)

axum/src/handler/into_service_state_in_extension.rs

@@ -5,6 +5,7 @@ use std::{
     convert::Infallible,
     fmt,
     marker::PhantomData,
+    sync::Arc,
     task::{Context, Poll},
 };
 use tower_service::Service;
@@ -54,7 +55,7 @@ impl<H, T, S, B> Service<Request<B>> for IntoServiceStateInExtension<H, T, S, B>
 where
     H: Handler<T, S, B> + Clone + Send + 'static,
     B: Send + 'static,
-    S: Clone + Send + Sync + 'static,
+    S: Send + Sync + 'static,
 {
     type Response = Response;
     type Error = Infallible;
@@ -73,7 +74,7 @@ where
 
         let state = req
             .extensions_mut()
-            .remove::<S>()
+            .remove::<Arc<S>>()
             .expect("state extension missing. This is a bug in axum, please file an issue");
 
         let handler = self.handler.clone();

axum/src/handler/mod.rs

@@ -42,7 +42,7 @@ use crate::{
     routing::IntoMakeService,
 };
 use http::Request;
-use std::{convert::Infallible, fmt, future::Future, marker::PhantomData, pin::Pin};
+use std::{convert::Infallible, fmt, future::Future, marker::PhantomData, pin::Pin, sync::Arc};
 use tower::ServiceExt;
 use tower_layer::Layer;
 use tower_service::Service;
@@ -100,7 +100,7 @@ pub trait Handler<T, S = (), B = Body>: Clone + Send + Sized + 'static {
     type Future: Future<Output = Response> + Send + 'static;
 
     /// Call the handler with the given request.
-    fn call(self, state: S, req: Request<B>) -> Self::Future;
+    fn call(self, state: Arc<S>, req: Request<B>) -> Self::Future;
 
     /// Apply a [`tower::Layer`] to the handler.
     ///
@@ -151,6 +151,11 @@ pub trait Handler<T, S = (), B = Body>: Clone + Send + Sized + 'static {
 
     /// Convert the handler into a [`Service`] by providing the state
     fn with_state(self, state: S) -> WithState<Self, T, S, B> {
+        self.with_state_arc(Arc::new(state))
+    }
+
+    /// Convert the handler into a [`Service`] by providing the state
+    fn with_state_arc(self, state: Arc<S>) -> WithState<Self, T, S, B> {
         WithState {
             service: IntoService::new(self, state),
         }
@@ -166,7 +171,7 @@ where
 {
     type Future = Pin<Box<dyn Future<Output = Response> + Send>>;
 
-    fn call(self, _state: S, _req: Request<B>) -> Self::Future {
+    fn call(self, _state: Arc<S>, _req: Request<B>) -> Self::Future {
         Box::pin(async move { self().await.into_response() })
     }
 }
@@ -179,15 +184,15 @@ macro_rules! impl_handler {
             F: FnOnce($($ty,)*) -> Fut + Clone + Send + 'static,
             Fut: Future<Output = Res> + Send,
             B: Send + 'static,
-            S: Send + 'static,
+            S: Send + Sync + 'static,
             Res: IntoResponse,
             $( $ty: FromRequest<S, B> + Send,)*
         {
             type Future = Pin<Box<dyn Future<Output = Response> + Send>>;
 
-            fn call(self, state: S, req: Request<B>) -> Self::Future {
+            fn call(self, state: Arc<S>, req: Request<B>) -> Self::Future {
                 Box::pin(async move {
-                    let mut req = RequestParts::with_state(state, req);
+                    let mut req = RequestParts::with_state_arc(state, req);
 
                     $(
                         let $ty = match $ty::from_request(&mut req).await {
@@ -254,10 +259,10 @@ where
 {
     type Future = future::LayeredFuture<B, L::Service>;
 
-    fn call(self, state: S, req: Request<B>) -> Self::Future {
+    fn call(self, state: Arc<S>, req: Request<B>) -> Self::Future {
         use futures_util::future::{FutureExt, Map};
 
-        let svc = self.handler.with_state(state);
+        let svc = self.handler.with_state_arc(state);
         let svc = self.layer.layer(svc);
 
         let future: Map<

axum/src/handler/with_state.rs

@@ -106,7 +106,7 @@ impl<H, T, S, B> Service<Request<B>> for WithState<H, T, S, B>
 where
     H: Handler<T, S, B> + Clone + Send + 'static,
     B: Send + 'static,
-    S: Clone,
+    S: Send + Sync,
 {
     type Response = <IntoService<H, T, S, B> as Service<Request<B>>>::Response;
     type Error = <IntoService<H, T, S, B> as Service<Request<B>>>::Error;
@@ -134,7 +134,6 @@ impl<H, T, S, B> std::fmt::Debug for WithState<H, T, S, B> {
 impl<H, T, S, B> Clone for WithState<H, T, S, B>
 where
     H: Clone,
-    S: Clone,
 {
     fn clone(&self) -> Self {
         Self {

axum/src/json.rs

@@ -100,7 +100,7 @@ where
     B: HttpBody + Send,
     B::Data: Send,
     B::Error: Into<BoxError>,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = JsonRejection;
 

axum/src/middleware/from_extractor.rs

@@ -48,7 +48,7 @@ use tower_service::Service;
 /// impl<S, B> FromRequest<S, B> for RequireAuth
 /// where
 ///     B: Send,
-///     S: Send,
+///     S: Send + Sync,
 /// {
 ///     type Rejection = StatusCode;
 ///
@@ -283,7 +283,7 @@ mod tests {
         impl<S, B> FromRequest<S, B> for RequireAuth
         where
             B: Send,
-            S: Send,
+            S: Send + Sync,
         {
             type Rejection = StatusCode;
 

axum/src/routing/method_routing.rs

@@ -16,6 +16,7 @@ use std::{
     convert::Infallible,
     fmt,
     marker::PhantomData,
+    sync::Arc,
     task::{Context, Poll},
 };
 use tower::{service_fn, util::MapResponseLayer};
@@ -143,7 +144,7 @@ macro_rules! top_level_handler_fn {
             H: Handler<T, S, B>,
             B: Send + 'static,
             T: 'static,
-            S: Clone + Send + Sync + 'static,
+            S: Send + Sync + 'static,
         {
             on(MethodFilter::$method, handler)
         }
@@ -279,7 +280,7 @@ macro_rules! chained_handler_fn {
         where
             H: Handler<T, S, B>,
             T: 'static,
-            S: Clone + Send + Sync + 'static,
+            S: Send + Sync + 'static,
         {
             self.on(MethodFilter::$method, handler)
         }
@@ -428,7 +429,7 @@ where
     H: Handler<T, S, B>,
     B: Send + 'static,
     T: 'static,
-    S: Clone + Send + Sync + 'static,
+    S: Send + Sync + 'static,
 {
     MethodRouter::new().on(filter, handler)
 }
@@ -475,7 +476,7 @@ where
     H: Handler<T, S, B>,
     B: Send + 'static,
     T: 'static,
-    S: Clone + Send + Sync + 'static,
+    S: Send + Sync + 'static,
 {
     MethodRouter::new()
         .fallback_boxed_response_body(IntoServiceStateInExtension::new(handler))
@@ -599,7 +600,7 @@ where
     where
         H: Handler<T, S, B>,
         T: 'static,
-        S: Clone + Send + Sync + 'static,
+        S: Send + Sync + 'static,
     {
         self.on_service_boxed_response_body(filter, IntoServiceStateInExtension::new(handler))
     }
@@ -618,7 +619,7 @@ where
     where
         H: Handler<T, S, B>,
         T: 'static,
-        S: Clone + Send + Sync + 'static,
+        S: Send + Sync + 'static,
     {
         self.fallback_service(IntoServiceStateInExtension::new(handler))
     }
@@ -727,6 +728,13 @@ where
     ///
     /// See [`State`](crate::extract::State) for more details about accessing state.
     pub fn with_state(self, state: S) -> WithState<S, B, E> {
+        self.with_state_arc(Arc::new(state))
+    }
+
+    /// Provide the [`Arc`]'ed state.
+    ///
+    /// See [`State`](crate::extract::State) for more details about accessing state.
+    pub fn with_state_arc(self, state: Arc<S>) -> WithState<S, B, E> {
         WithState {
             method_router: self,
             state,
@@ -1127,7 +1135,7 @@ where
 /// Created with [`MethodRouter::with_state`]
 pub struct WithState<S, B, E> {
     method_router: MethodRouter<S, B, E>,
-    state: S,
+    state: Arc<S>,
 }
 
 impl<S, B, E> WithState<S, B, E> {
@@ -1156,14 +1164,11 @@ impl<S, B, E> WithState<S, B, E> {
     }
 }
 
-impl<S, B, E> Clone for WithState<S, B, E>
-where
-    S: Clone,
-{
+impl<S, B, E> Clone for WithState<S, B, E> {
     fn clone(&self) -> Self {
         Self {
             method_router: self.method_router.clone(),
-            state: self.state.clone(),
+            state: Arc::clone(&self.state),
         }
     }
 }
@@ -1183,7 +1188,7 @@ where
 impl<S, B, E> Service<Request<B>> for WithState<S, B, E>
 where
     B: HttpBody,
-    S: Clone + Send + Sync + 'static,
+    S: Send + Sync + 'static,
 {
     type Response = Response;
     type Error = E;
@@ -1232,7 +1237,7 @@ where
                 },
         } = self;
 
-        req.extensions_mut().insert(state.clone());
+        req.extensions_mut().insert(Arc::clone(state));
 
         call!(req, method, HEAD, head);
         call!(req, method, HEAD, get);

axum/src/routing/mod.rs

@@ -62,19 +62,16 @@ impl RouteId {
 
 /// The router type for composing handlers and services.
 pub struct Router<S = (), B = Body> {
-    state: S,
+    state: Arc<S>,
     routes: HashMap<RouteId, Endpoint<S, B>>,
     node: Arc<Node>,
     fallback: Fallback<B>,
 }
 
-impl<S, B> Clone for Router<S, B>
-where
-    S: Clone,
-{
+impl<S, B> Clone for Router<S, B> {
     fn clone(&self) -> Self {
         Self {
-            state: self.state.clone(),
+            state: Arc::clone(&self.state),
             routes: self.routes.clone(),
             node: Arc::clone(&self.node),
             fallback: self.fallback.clone(),
@@ -85,7 +82,7 @@ where
 impl<S, B> Default for Router<S, B>
 where
     B: HttpBody + Send + 'static,
-    S: Default + Clone + Send + Sync + 'static,
+    S: Default + Send + Sync + 'static,
 {
     fn default() -> Self {
         Self::with_state(S::default())
@@ -125,7 +122,7 @@ where
 impl<S, B> Router<S, B>
 where
     B: HttpBody + Send + 'static,
-    S: Clone + Send + Sync + 'static,
+    S: Send + Sync + 'static,
 {
     /// Create a new `Router` with the given state.
     ///
@@ -134,6 +131,16 @@ where
     /// Unless you add additional routes this will respond with `404 Not Found` to
     /// all requests.
     pub fn with_state(state: S) -> Self {
+        Self::with_state_arc(Arc::new(state))
+    }
+
+    /// Create a new `Router` with the given [`Arc`]'ed state.
+    ///
+    /// See [`State`](crate::extract::State) for more details about accessing state.
+    ///
+    /// Unless you add additional routes this will respond with `404 Not Found` to
+    /// all requests.
+    pub fn with_state_arc(state: Arc<S>) -> Self {
         Self {
             state,
             routes: Default::default(),
@@ -262,7 +269,7 @@ where
     pub fn merge<S2, R>(mut self, other: R) -> Self
     where
         R: Into<Router<S2, B>>,
-        S2: Clone + Send + Sync + 'static,
+        S2: Send + Sync + 'static,
     {
         let Router {
             state,
@@ -282,7 +289,7 @@ where
                     method_router
                         // this will set the state for each route
                         // such we don't override the inner state later in `MethodRouterWithState`
-                        .layer(Extension(state.clone()))
+                        .layer(Extension(Arc::clone(&state)))
                         .downcast_state(),
                 ),
                 Endpoint::Route(route) => self.route_service(path, route),
@@ -383,8 +390,8 @@ where
         H: Handler<T, S, B>,
         T: 'static,
     {
-        let state = self.state.clone();
-        self.fallback_service(handler.with_state(state))
+        let state = Arc::clone(&self.state);
+        self.fallback_service(handler.with_state_arc(state))
     }
 
     /// Add a fallback [`Service`] to the router.
@@ -484,7 +491,10 @@ where
             .clone();
 
         match &mut route {
-            Endpoint::MethodRouter(inner) => inner.clone().with_state(self.state.clone()).call(req),
+            Endpoint::MethodRouter(inner) => inner
+                .clone()
+                .with_state_arc(Arc::clone(&self.state))
+                .call(req),
             Endpoint::Route(inner) => inner.call(req),
         }
     }
@@ -498,7 +508,7 @@ where
 impl<S, B> Service<Request<B>> for Router<S, B>
 where
     B: HttpBody + Send + 'static,
-    S: Clone + Send + Sync + 'static,
+    S: Send + Sync + 'static,
 {
     type Response = Response;
     type Error = Infallible;
@@ -618,10 +628,7 @@ enum Endpoint<S, B> {
     Route(Route<B>),
 }
 
-impl<S, B> Clone for Endpoint<S, B>
-where
-    S: Clone,
-{
+impl<S, B> Clone for Endpoint<S, B> {
     fn clone(&self) -> Self {
         match self {
             Endpoint::MethodRouter(inner) => Endpoint::MethodRouter(inner.clone()),

axum/src/typed_header.rs

@@ -56,7 +56,7 @@ impl<T, S, B> FromRequest<S, B> for TypedHeader<T>
 where
     T: headers::Header,
     B: Send,
-    S: Send,
+    S: Send + Sync,
 {
     type Rejection = TypedHeaderRejection;