mas_tower/
utils.rs

1// Copyright 2023 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use opentelemetry::{KeyValue, Value};
16use tower::{Layer, Service};
17
18/// A simple static key-value pair.
19#[derive(Clone, Debug)]
20pub struct KV<V>(pub &'static str, pub V);
21
22impl<V> From<KV<V>> for KeyValue
23where
24    V: Into<Value>,
25{
26    fn from(value: KV<V>) -> Self {
27        Self::new(value.0, value.1.into())
28    }
29}
30
31/// A wrapper around a function that can be used to generate a key-value pair,
32/// make or enrich spans.
33#[derive(Clone, Debug)]
34pub struct FnWrapper<F>(pub F);
35
36/// A no-op layer that has the request type bound.
37#[derive(Clone, Copy, Debug)]
38pub struct IdentityLayer<R> {
39    _request: std::marker::PhantomData<R>,
40}
41
42impl<R> Default for IdentityLayer<R> {
43    fn default() -> Self {
44        Self {
45            _request: std::marker::PhantomData,
46        }
47    }
48}
49
50/// A no-op service that has the request type bound.
51#[derive(Clone, Copy, Debug)]
52pub struct IdentityService<R, S> {
53    _request: std::marker::PhantomData<R>,
54    inner: S,
55}
56
57impl<R, S> Default for IdentityService<R, S>
58where
59    S: Default,
60{
61    fn default() -> Self {
62        Self {
63            _request: std::marker::PhantomData,
64            inner: S::default(),
65        }
66    }
67}
68
69impl<R, S> Layer<S> for IdentityLayer<R> {
70    type Service = IdentityService<R, S>;
71
72    fn layer(&self, inner: S) -> Self::Service {
73        IdentityService {
74            _request: std::marker::PhantomData,
75            inner,
76        }
77    }
78}
79
80impl<S, R> Service<R> for IdentityService<R, S>
81where
82    S: Service<R>,
83{
84    type Response = S::Response;
85    type Error = S::Error;
86    type Future = S::Future;
87
88    fn poll_ready(
89        &mut self,
90        cx: &mut std::task::Context<'_>,
91    ) -> std::task::Poll<Result<(), Self::Error>> {
92        self.inner.poll_ready(cx)
93    }
94
95    fn call(&mut self, req: R) -> Self::Future {
96        self.inner.call(req)
97    }
98}