Push RequestId setting into NewHandlerService

This will ensure all Handler instances, not just Router are provided with this functionality which is a reasonable default expectation.
gotham-rs · Jun 15, 2017 · 3450180 · 3450180
1 parent d107e8d
commit 3450180
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 36 deletions.
diff --git a/src/handler/mod.rs b/src/handler/mod.rs
@@ -9,7 +9,7 @@ use hyper;
 use hyper::server;
 use hyper::server::Request;
 use futures::{future, Future};
-use state::State;
+use state::{State, set_request_id};
 use std::io;
 use std::sync::Arc;
 
@@ -129,13 +129,16 @@ impl<T> server::Service for NewHandlerService<T>
     type Future = Box<Future<Item = server::Response, Error = hyper::Error>>;
 
     fn call(&self, req: Self::Request) -> Self::Future {
+        let mut state = State::new();
+        set_request_id(&mut state, &req);
+
         // Hyper doesn't allow us to present an affine-typed `Handler` interface directly. We have
         // to emulate the promise given by hyper's documentation, by creating a `Handler` value and
         // immediately consuming it.
         match self.t.new_handler() {
             Ok(handler) => {
                 handler
-                    .handle(State::new(), req)
+                    .handle(state, req)
                     .and_then(|(_, response)| future::ok(response))
                     .or_else(|(_, error)| future::err(error))
                     .boxed()

diff --git a/src/router/mod.rs b/src/router/mod.rs
@@ -16,7 +16,6 @@ use handler::{NewHandler, Handler, HandlerFuture};
 use router::response_extender::ResponseExtender;
 use router::route::Route;
 use router::tree::{SegmentMapping, Tree};
-use state::request_id::set_request_id;
 use state::{State, StateData};
 use http::{request_path, query_string};
 
@@ -96,8 +95,6 @@ impl<'n, P> Handler for Router<'n, P>
     /// `Tree`, storing any path related variables in `State` and dispatching
     /// appropriately to the configured `Handler`.
     fn handle(self, mut state: State, req: Request) -> Box<HandlerFuture> {
-        set_request_id(&mut state, &req);
-
         let uri = req.uri().clone();
         self.populate_state(&mut state, &req);
         let response = self.route(uri, state, req);
@@ -196,13 +193,12 @@ mod tests {
     use std::str::FromStr;
     use hyper::{Request, Method, Uri, Body};
     use hyper::header::{ContentType, ContentLength};
-    use uuid::Uuid;
 
     use borrow_bag;
 
     use router::tree::TreeBuilder;
     use router::response_extender::ResponseExtenderBuilder;
-    use state::request_id;
+    use state::set_request_id;
 
     #[test]
     fn executes_response_extender_when_present() {
@@ -223,7 +219,8 @@ mod tests {
         let uri = Uri::from_str("https://test.gotham.rs").unwrap();
         let request: Request<Body> = Request::new(method, uri);
 
-        let state = State::new();
+        let mut state = State::new();
+        set_request_id(&mut state, &request);
         let result = router.handle(state, request).wait();
 
         match result {
@@ -250,7 +247,8 @@ mod tests {
         request.set_version(version.clone());
         request.headers_mut().set(ContentType::json());
 
-        let state = State::new();
+        let mut state = State::new();
+        set_request_id(&mut state, &request);
         let result = router.handle(state, request).wait();
 
         match result {
@@ -268,30 +266,4 @@ mod tests {
             Err(_) => panic!("Router should have correctly handled request"),
         };
     }
-
-    #[test]
-    fn populates_request_id() {
-        let tree_builder = TreeBuilder::new();
-        let tree = tree_builder.finalize();
-        let pipelines = borrow_bag::new_borrow_bag();
-        let response_extender = ResponseExtenderBuilder::new().finalize();
-
-        let router = Router::new(tree, pipelines, response_extender);
-        let method = Method::Get;
-        let uri = Uri::from_str("https://test.gotham.rs").unwrap();
-        let request: Request<Body> = Request::new(method, uri);
-
-        let state = State::new();
-        let result = router.handle(state, request).wait();
-
-        match result {
-            Ok((state, _res)) => {
-                assert_eq!(4,
-                           Uuid::parse_str(request_id(&state))
-                               .unwrap()
-                               .get_version_num())
-            }
-            Err(_) => panic!("Router should have correctly handled request"),
-        }
-    }
 }
diff --git a/src/state/mod.rs b/src/state/mod.rs
@@ -6,6 +6,7 @@ use std::collections::HashMap;
 use std::any::{Any, TypeId};
 
 pub use self::request_id::request_id;
+pub use self::request_id::set_request_id;
 
 /// Provides storage for request state, and stores one item of each type. The types used for
 /// storage must implement the `gotham::state::StateData` trait to allow its storage.

diff --git a/src/state/request_id.rs b/src/state/request_id.rs
@@ -26,7 +26,7 @@ header! {
 /// 1. If the header X-Request-ID is provided this value is used as is;
 /// 2. Alternatively creates and stores a UUID v4 value.
 ///
-/// This method MUST be invoked by Gotham, specifically by the `Router`, before handing control to
+/// This method MUST be invoked by Gotham, before handing control to
 /// pipelines or Handlers to ensure that a value for `RequestId` is always available.
 pub fn set_request_id<'a>(state: &'a mut State, req: &Request) -> &'a str {
     if !state.has::<RequestId>() {