Skip to main content

obeli_sk_grpc/
lib.rs

1pub mod grpc_gen;
2pub mod grpc_mapping;
3
4use http::{Uri, uri::Scheme};
5use tonic::transport::{Channel, ClientTlsConfig};
6
7pub type TonicResult<T> = Result<T, tonic::Status>;
8
9pub type TonicRespResult<T> = TonicResult<tonic::Response<T>>;
10
11#[derive(Debug, thiserror::Error)]
12pub enum ToChannelError {
13    #[error(transparent)]
14    InvalidUri(http::uri::InvalidUri),
15    #[error(transparent)]
16    Transport(tonic::transport::Error),
17    #[error("unknown schame {0}")]
18    UnknownScheme(String),
19}
20
21pub async fn to_channel(url: &str) -> Result<Channel, ToChannelError> {
22    let tls = ClientTlsConfig::new().with_native_roots();
23    let url: Uri = url.parse().map_err(ToChannelError::InvalidUri)?;
24    if url.scheme() == Some(&Scheme::HTTP) {
25        Channel::builder(url)
26            .connect()
27            .await
28            .map_err(ToChannelError::Transport)
29    } else if url.scheme() == Some(&Scheme::HTTPS) {
30        Channel::builder(url)
31            .tls_config(tls)
32            .map_err(ToChannelError::Transport)?
33            .connect()
34            .await
35            .map_err(ToChannelError::Transport)
36    } else {
37        Err(ToChannelError::UnknownScheme(format!("{:?}", url.scheme())))
38    }
39}
40
41#[cfg(feature = "otlp")]
42// Source: https://github.com/hseeberger/hello-tracing-rs/blob/b411f8b192b7d585c42b5928ea635b2bd8bde29c/hello-tracing-common/src/otel/grpc.rs
43pub mod injector {
44    use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
45    use tracing::{Span, warn};
46
47    struct MetadataInjector<'a>(&'a mut MetadataMap);
48
49    impl opentelemetry::propagation::Injector for MetadataInjector<'_> {
50        fn set(&mut self, key: &str, value: String) {
51            match MetadataKey::from_bytes(key.as_bytes()) {
52                Ok(key) => match MetadataValue::try_from(&value) {
53                    Ok(value) => {
54                        self.0.insert(key, value);
55                    }
56
57                    Err(error) => warn!(
58                        value,
59                        error = format!("{error:?}"),
60                        "cannot parse metadata value"
61                    ),
62                },
63
64                Err(error) => warn!(
65                    key,
66                    error = format!("{error:?}"),
67                    "cannot parse metadata key"
68                ),
69            }
70        }
71    }
72
73    /// Client interceptor that injects current span IDs to the request's metadata.
74    pub struct TracingInjector;
75
76    impl tonic::service::Interceptor for TracingInjector {
77        fn call(
78            &mut self,
79            mut request: tonic::Request<()>,
80        ) -> Result<tonic::Request<()>, tonic::Status> {
81            use tracing_opentelemetry::OpenTelemetrySpanExt as _;
82            opentelemetry::global::get_text_map_propagator(|propagator| {
83                let context = Span::current().context();
84                propagator.inject_context(&context, &mut MetadataInjector(request.metadata_mut()));
85            });
86            Ok(request)
87        }
88    }
89}
90
91#[cfg(feature = "otlp")]
92pub mod extractor {
93    use opentelemetry_http::HeaderExtractor;
94    use tracing::Span;
95    use tracing_opentelemetry::OpenTelemetrySpanExt as _;
96
97    /// Uses [`HeaderExtractor`] to fetch trace and span IDs from the request headers,
98    /// and sets the extracted context as the current span's parent, if valid.
99    pub fn accept_trace<B>(request: http::Request<B>) -> http::Request<B> {
100        // Current context, if no or invalid data is received.
101        let parent_context = opentelemetry::global::get_text_map_propagator(|propagator| {
102            propagator.extract(&HeaderExtractor(request.headers()))
103        });
104        let _ = Span::current().set_parent(parent_context);
105        request
106    }
107}