cronback-lib 0.1.0

Common utilities for all cronback services
Documentation 
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::task::{Context, Poll};
use std::time::Instant;

use hyper::Body;
use metrics::{histogram, increment_counter};
use tonic::body::BoxBody;
use tower::{Layer, Service};

use crate::consts::{PROJECT_ID_HEADER, REQUEST_ID_HEADER};
use crate::model::{ModelId, ValidShardedId};
use crate::types::{ProjectId, RequestId};

#[derive(Debug, Clone, Default)]
pub struct CronbackRpcMiddleware {
    /// Sets the label "service" in emitted metrics
    service_name: String,
}

impl CronbackRpcMiddleware {
    pub fn new(service_name: &str) -> CronbackRpcMiddleware {
        CronbackRpcMiddleware {
            service_name: service_name.into(),
        }
    }
}

impl<S> Layer<S> for CronbackRpcMiddleware {
    type Service = InnerMiddleware<S>;

    fn layer(&self, service: S) -> Self::Service {
        InnerMiddleware::new(&self.service_name, service)
    }
}

#[derive(Debug, Clone)]
pub struct InnerMiddleware<S> {
    inner: S,
    service_name: String,
}

impl<S> InnerMiddleware<S> {
    pub fn new(service_name: &str, inner: S) -> Self {
        InnerMiddleware {
            inner,
            service_name: service_name.to_owned(),
        }
    }
}

impl<S> Service<hyper::Request<Body>> for InnerMiddleware<S>
where
    S: Service<hyper::Request<Body>, Response = hyper::Response<BoxBody>>
        + Clone
        + Send
        + 'static,
    S::Future: Send + 'static,
{
    type Error = S::Error;
    type Future = futures::future::BoxFuture<
        'static,
        Result<Self::Response, Self::Error>,
    >;
    type Response = S::Response;

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

    fn call(&mut self, mut req: hyper::Request<Body>) -> Self::Future {
        // This is necessary because tonic internally uses
        // `tower::buffer::Buffer`. See https://github.com/tower-rs/tower/issues/547#issuecomment-767629149
        // for details on why this is necessary
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);

        // Do we have a x-cronback-request-id header? Only used in grpc
        // services. The api-server will set the request-id header with
        // a random value.
        if let Some(cronback_request_id) = req.headers().get(REQUEST_ID_HEADER)
        {
            // If so, set the request id to the value of the header
            let cronback_request_id = cronback_request_id.to_str().unwrap();
            let cronback_request_id =
                RequestId::from(cronback_request_id.to_owned());
            req.extensions_mut().insert(cronback_request_id);
        }

        // Do we have a x-cronback-project-id header?
        // If project-id is set, it must be valid. We store the result in
        // extensions.
        if let Some(project_id) = req.headers().get(PROJECT_ID_HEADER) {
            // If so, set the project id to the value of the header
            let project_id = project_id.to_str().unwrap();
            let maybe_project_id =
                ProjectId::from(project_id.to_owned()).validated();
            req.extensions_mut().insert(maybe_project_id);
        }

        // Removes the leading '/' in the path.
        let endpoint = req.uri().path()[1..].to_owned();
        let service_name = self.service_name.clone();
        let start = Instant::now();
        increment_counter!(
            "rpc.requests_total",
            "service" => service_name.clone(),
            "endpoint" => endpoint.clone()
        );
        Box::pin(async move {
            let mut response = inner.call(req).await?;
            let latency_s = (Instant::now() - start).as_secs_f64();
            histogram!(
                "rpc.duration_seconds",
                latency_s,
                "service" => service_name.clone(),
                "endpoint" => endpoint.clone(),
            );

            // Inject request_id into response headers
            if let Some(request_id) =
                response.extensions().get::<RequestId>().cloned()
            {
                response.headers_mut().insert(
                    REQUEST_ID_HEADER,
                    request_id.to_string().parse().unwrap(),
                );
            }

            // Inject project_id into response headers
            if let Some(project_id) = response
                .extensions()
                .get::<ValidShardedId<ProjectId>>()
                .cloned()
            {
                response.headers_mut().insert(
                    PROJECT_ID_HEADER,
                    project_id.to_string().parse().unwrap(),
                );
            }

            Ok(response)
        })
    }
}