dipper/mw/
request_id.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
//! [trace specification]: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/api.md
//! [opentelemetry specification]: https://github.com/open-telemetry/opentelemetry-specification/blob/213b91992ff11eb7a0116c91985dc4f429be72fe/specification/logs/data-model.md

use crate::dyn_mod::*;
use crate::handler;
use crate::http::HeaderValue;
use std::sync::Arc;
use sulid::SulidGenerator;

pub const HEADER_REQUEST_ID: &'static str = "x-request-id";

pub struct RequestIdHoop {
    sulid: Arc<SulidGenerator>,
}

#[handler]
impl RequestIdHoop {
    pub fn new(worker_id: u16) -> Self {
        Self {
            sulid: Arc::new(SulidGenerator::v2_new(worker_id)),
        }
    }
    async fn handle(&self, request: &mut Request, depot: &mut Depot, response: &mut Response, ctrl: &mut FlowCtrl) {
        let mut is_gen = false;
        let request_id = match request.headers().get(HEADER_REQUEST_ID) {
            Some(v) => v.to_str().map(|s| s.to_owned()).unwrap_or_else(|_| {
                is_gen = true;
                self.sulid.generate().to_string()
            }),
            None => {
                is_gen = true;
                self.sulid.generate().to_string()
            },
        };

        if is_gen {
            let header_request_id = HeaderValue::from_str(request_id.as_str()).unwrap();
            request.add_header(HEADER_REQUEST_ID, header_request_id.clone(), true).ok();
            response.add_header(HEADER_REQUEST_ID, header_request_id, true).ok();
        }

        depot.insert(HEADER_REQUEST_ID, request_id.clone());

        ctrl.call_next(request, depot, response).await;
    }
}