matrix/authentication-service
1
0
Fork 0
mirror of https://github.com/matrix-org/matrix-authentication-service.git synced 2025-07-31 09:24:31 +03:00
Code Activity

Replace the OTEL-based tracing layer with tracing based layers

Browse Source
This commit is contained in:
Quentin Gliech
2023-04-15 14:21:12 +02:00
parent 4a4fbc7a16
commit 08f58db08b
47 changed files with 1703 additions and 1520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/axum-utils/src/client_authorization.rs
View File

@ -191,7 +191,7 @@ async fn fetch_jwks(
.unwrap();
let mut client = http_client_factory
.client("fetch-jwks")
.client()
.await?
.response_body_to_bytes()
.json_response::<PublicJsonWebKeySet>();

14
crates/axum-utils/src/http_client_factory.rs
View File

@ -43,16 +43,13 @@ impl HttpClientFactory {
/// # Errors
///
/// Returns an error if the client failed to initialise
pub async fn client<B>(
&self,
operation: &'static str,
) -> Result<ClientService<TracedClient<B>>, ClientInitError>
pub async fn client<B>(&self) -> Result<ClientService<TracedClient<B>>, ClientInitError>
where
B: axum::body::HttpBody + Send,
B::Data: Send,
{
let client = mas_http::make_traced_client::<B>().await?;
let layer = ClientLayer::with_semaphore(operation, self.semaphore.clone());
let layer = ClientLayer::with_semaphore(self.semaphore.clone());
Ok(layer.layer(client))
}
@ -61,11 +58,8 @@ impl HttpClientFactory {
/// # Errors
///
/// Returns an error if the client failed to initialise
pub async fn http_service(
&self,
operation: &'static str,
) -> Result<HttpService, ClientInitError> {
let client = self.client(operation).await?;
pub async fn http_service(&self) -> Result<HttpService, ClientInitError> {
let client = self.client().await?;
let client = (
MapErrLayer::new(BoxError::from),
MapRequestLayer::new(|req: http::Request<_>| req.map(Full::new)),

1
crates/cli/Cargo.toml
View File

@ -57,6 +57,7 @@ mas-storage = { path = "../storage" }
mas-storage-pg = { path = "../storage-pg" }
mas-tasks = { path = "../tasks" }
mas-templates = { path = "../templates" }
mas-tower = { path = "../tower" }
oauth2-types = { path = "../oauth2-types" }
[dev-dependencies]

4
crates/cli/src/commands/debug.rs
View File

@ -75,7 +75,7 @@ impl Options {
url,
} => {
let _span = info_span!("cli.debug.http").entered();
let mut client = http_client_factory.client("cli-debug-http").await?;
let mut client = http_client_factory.client().await?;
let request = hyper::Request::builder()
.uri(url)
.body(hyper::Body::empty())?;
@ -101,7 +101,7 @@ impl Options {
} => {
let _span = info_span!("cli.debug.http").entered();
let mut client = http_client_factory
.client("cli-debug-http")
.client()
.await?
.response_body_to_bytes()
.json_response();

166
crates/cli/src/server.rs
View File

@ -13,27 +13,153 @@
// limitations under the License.
use std::{
borrow::Cow,
future::ready,
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener, ToSocketAddrs},
os::unix::net::UnixListener,
};
use anyhow::Context;
use axum::{body::HttpBody, error_handling::HandleErrorLayer, extract::FromRef, Extension, Router};
use hyper::StatusCode;
use axum::{
body::HttpBody,
error_handling::HandleErrorLayer,
extract::{FromRef, MatchedPath},
Extension, Router,
};
use hyper::{Method, Request, Response, StatusCode, Version};
use listenfd::ListenFd;
use mas_config::{HttpBindConfig, HttpResource, HttpTlsConfig, UnixOrTcp};
use mas_handlers::AppState;
use mas_http::otel::TraceLayer;
use mas_listener::{unix_or_tcp::UnixOrTcpListener, ConnectionInfo};
use mas_router::Route;
use mas_spa::ViteManifestService;
use mas_templates::Templates;
use opentelemetry::KeyValue;
use mas_tower::{
make_span_fn, metrics_attributes_fn, DurationRecorderLayer, InFlightCounterLayer, TraceLayer,
KV,
};
use opentelemetry::{Key, KeyValue};
use opentelemetry_http::HeaderExtractor;
use opentelemetry_semantic_conventions::trace::{
HTTP_METHOD, HTTP_ROUTE, HTTP_SCHEME, HTTP_STATUS_CODE,
};
use rustls::ServerConfig;
use sentry_tower::{NewSentryLayer, SentryHttpLayer};
use tower::Layer;
use tower_http::{compression::CompressionLayer, services::ServeDir};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;
const NET_PROTOCOL_NAME: Key = Key::from_static_str("net.protocol.name");
const NET_PROTOCOL_VERSION: Key = Key::from_static_str("net.protocol.version");
const MAS_LISTENER_NAME: Key = Key::from_static_str("mas.listener.name");
#[inline]
fn otel_http_method<B>(request: &Request<B>) -> Cow<'static, str> {
match request.method() {
&Method::OPTIONS => "OPTIONS".into(),
&Method::GET => "GET".into(),
&Method::POST => "POST".into(),
&Method::PUT => "PUT".into(),
&Method::DELETE => "DELETE".into(),
&Method::HEAD => "HEAD".into(),
&Method::TRACE => "TRACE".into(),
&Method::CONNECT => "CONNECT".into(),
&Method::PATCH => "PATCH".into(),
other => other.to_string().into(),
}
}
#[inline]
fn otel_net_protocol_version<B>(request: &Request<B>) -> Cow<'static, str> {
match request.version() {
Version::HTTP_09 => "0.9".into(),
Version::HTTP_10 => "1.0".into(),
Version::HTTP_11 => "1.1".into(),
Version::HTTP_2 => "2.0".into(),
Version::HTTP_3 => "3.0".into(),
other => format!("{other:?}").into(),
}
}
fn otel_http_route<B>(request: &Request<B>) -> Option<&str> {
request
.extensions()
.get::<MatchedPath>()
.map(MatchedPath::as_str)
}
fn otel_http_target<B>(request: &Request<B>) -> &str {
request.uri().path_and_query().map_or("", |p| p.as_str())
}
fn otel_http_scheme<B>(request: &Request<B>) -> &'static str {
// XXX: maybe we should panic if the connection info was not injected in the
// request extensions
request
.extensions()
.get::<ConnectionInfo>()
.map_or("http", |conn_info| {
if conn_info.get_tls_ref().is_some() {
"https"
} else {
"http"
}
})
}
fn make_http_span<B>(req: &Request<B>) -> Span {
let method = otel_http_method(req);
let route = otel_http_route(req);
let span_name = if let Some(route) = route.as_ref() {
format!("{method} {route}")
} else {
format!("{method}")
};
let span = tracing::info_span!(
"http.server.request",
"otel.kind" = "server",
"otel.name" = span_name,
"otel.status_code" = tracing::field::Empty,
"net.protocol.name" = "http",
"net.protocol.version" = otel_net_protocol_version(req).as_ref(),
"http.scheme" = otel_http_scheme(req),
"http.method" = method.as_ref(),
"http.route" = tracing::field::Empty,
"http.target" = otel_http_target(req),
"http.status_code" = tracing::field::Empty,
);
if let Some(route) = route.as_ref() {
span.record("http.route", route);
}
// Extract the parent span context from the request headers
let parent_context = opentelemetry::global::get_text_map_propagator(|propagator| {
let extractor = HeaderExtractor(req.headers());
let context = opentelemetry::Context::new();
propagator.extract_with_context(&context, &extractor)
});
span.set_parent(parent_context);
span
}
fn on_http_request_labels<B>(request: &Request<B>) -> Vec<KeyValue> {
vec![
NET_PROTOCOL_NAME.string("http"),
NET_PROTOCOL_VERSION.string(otel_net_protocol_version(request)),
HTTP_METHOD.string(otel_http_method(request)),
HTTP_SCHEME.string(otel_http_scheme(request).as_ref()),
HTTP_ROUTE.string(otel_http_route(request).unwrap_or("FALLBACK").to_owned()),
]
}
fn on_http_response_labels<B>(res: &Response<B>) -> Vec<KeyValue> {
vec![HTTP_STATUS_CODE.i64(res.status().as_u16().into())]
}
pub fn build_router<B>(
state: AppState,
@ -112,16 +238,34 @@ where
}
}
let mut trace_layer = TraceLayer::axum();
if let Some(name) = name {
trace_layer = trace_layer.with_static_attribute(KeyValue::new("listener", name.to_owned()));
}
router
.layer(
InFlightCounterLayer::new("http.server.active_requests").on_request((
name.map(|name| MAS_LISTENER_NAME.string(name.to_owned())),
metrics_attributes_fn(on_http_request_labels),
)),
)
.layer(
DurationRecorderLayer::new("http.server.duration")
.on_request((
name.map(|name| MAS_LISTENER_NAME.string(name.to_owned())),
metrics_attributes_fn(on_http_request_labels),
))
.on_response_fn(on_http_response_labels),
)
.layer(
TraceLayer::new((
make_span_fn(make_http_span),
name.map(|name| KV("mas.listener.name", name.to_owned())),
))
.on_response_fn(|span: &Span, response: &Response<_>| {
let status_code = response.status().as_u16();
span.record("http.status_code", status_code);
span.record("otel.status_code", "OK");
}),
)
.layer(SentryHttpLayer::new())
.layer(NewSentryLayer::new_from_top())
.layer(trace_layer)
.layer(CompressionLayer::new())
.with_state(state)
}

3
crates/cli/src/telemetry.rs
View File

@ -375,8 +375,9 @@ fn prometheus_meter() -> anyhow::Result<BasicController> {
fn prometheus_meter() -> anyhow::Result<BasicController> {
let controller = sdk::metrics::controllers::basic(
sdk::metrics::processors::factory(
// All histogram metrics are in milliseconds. Each bucket is ~2x the previous one.
sdk::metrics::selectors::simple::histogram([
0.01, 0.02, 0.05, 0.10, 0.20, 0.50, 1.0, 2.0, 5.0,
1.0, 3.0, 5.0, 10.0, 30.0, 50.0, 100.0, 300.0, 1000.0,
]),
sdk::export::metrics::aggregation::cumulative_temporality_selector(),
)

5
crates/email/Cargo.toml
View File

@ -9,14 +9,19 @@ license = "Apache-2.0"
async-trait = "0.1.68"
tracing = "0.1.37"
thiserror = "1.0.40"
headers = "0.3.8"
aws-sdk-sesv2 = { version = "0.24.0", default-features = false }
aws-config = { version = "0.54.1", default-features = false }
aws-smithy-client = { version = "0.54.4", default-features = false, features = ["client-hyper"] }
aws-smithy-async = { version = "0.54.4", default-features = false, features = ["rt-tokio"] }
aws-smithy-http = { version = "0.54.4", default-features = false }
aws-smithy-http-tower = { version = "0.54.4", default-features = false }
aws-types = "0.54.1"
mas-templates = { path = "../templates" }
mas-http = { path = "../http", features = ["aws-sdk", "client"] }
mas-tower = { path = "../tower", features = ["aws-sdk"] }
[dependencies.lettre]
version = "0.10.3"

77
crates/email/src/transport/aws_ses.rs
View File

@ -25,8 +25,11 @@ use aws_sdk_sesv2::{
};
use aws_smithy_async::rt::sleep::TokioSleep;
use aws_smithy_client::erase::{DynConnector, DynMiddleware};
use headers::{ContentLength, HeaderMapExt, Host, UserAgent};
use lettre::{address::Envelope, AsyncTransport};
use mas_http::{otel::TraceLayer, ClientInitError};
use mas_http::ClientInitError;
use mas_tower::{enrich_span_fn, make_span_fn, TraceContextLayer, TraceLayer};
use tracing::{info_span, Span};
pub type Error = aws_smithy_client::SdkError<aws_sdk_sesv2::error::SendEmailError>;
@ -58,7 +61,77 @@ impl Transport {
let http_connector = DynConnector::new(http_connector);
// Middleware to add tracing to AWS SDK operations
let middleware = DynMiddleware::new((TraceLayer::aws_sdk(), DefaultMiddleware::default()));
let middleware = DynMiddleware::new((
DefaultMiddleware::default(),
// TODO: factor this out somewhere else
TraceLayer::new(make_span_fn(|op: &aws_smithy_http::operation::Request| {
let properties = op.properties();
let request = op.http();
let span = info_span!(
"aws.sdk.operation",
"otel.kind" = "client",
"otel.name" = tracing::field::Empty,
"otel.status_code" = tracing::field::Empty,
"rpc.system" = "aws-api",
"rpc.service" = tracing::field::Empty,
"rpc.method" = tracing::field::Empty,
"http.method" = %request.method(),
"http.url" = %request.uri(),
"http.host" = tracing::field::Empty,
"http.request_content_length" = tracing::field::Empty,
"http.response_content_length" = tracing::field::Empty,
"http.status_code" = tracing::field::Empty,
"user_agent.original" = tracing::field::Empty,
);
if let Some(metadata) = properties.get::<aws_smithy_http::operation::Metadata>() {
span.record("rpc.service", metadata.service());
span.record("rpc.method", metadata.name());
let name = format!("{}::{}", metadata.service(), metadata.name());
span.record("otel.name", name);
} else if let Some(service) = properties.get::<aws_types::SigningService>() {
span.record("rpc.service", tracing::field::debug(service));
span.record("otel.name", tracing::field::debug(service));
}
let headers = request.headers();
if let Some(host) = headers.typed_get::<Host>() {
span.record("http.host", tracing::field::display(host));
}
if let Some(user_agent) = headers.typed_get::<UserAgent>() {
span.record("user_agent.original", tracing::field::display(user_agent));
}
if let Some(ContentLength(content_length)) = headers.typed_get() {
span.record("http.request_content_length", content_length);
}
span
}))
.on_response(enrich_span_fn(
|span: &Span, res: &aws_smithy_http::operation::Response| {
span.record("otel.status_code", "OK");
let response = res.http();
let status = response.status();
span.record("http.status_code", status.as_u16());
let headers = response.headers();
if let Some(ContentLength(content_length)) = headers.typed_get() {
span.record("http.response_content_length", content_length);
}
},
))
.on_error(enrich_span_fn(
|span: &Span, err: &aws_smithy_http_tower::SendOperationError| {
span.record("otel.status_code", "ERROR");
span.record("exception.message", tracing::field::debug(err));
},
)),
TraceContextLayer::new(),
));
// Use that connector for discovering the config
let config = ProviderConfig::default().with_http_connector(http_connector.clone());

4
crates/handlers/src/upstream_oauth2/authorize.rs
View File

@ -78,9 +78,7 @@ pub(crate) async fn get(
.await?
.ok_or(RouteError::ProviderNotFound)?;
let http_service = http_client_factory
.http_service("upstream-discover")
.await?;
let http_service = http_client_factory.http_service().await?;
// First, discover the provider
let metadata =

12
crates/handlers/src/upstream_oauth2/callback.rs
View File

@ -184,19 +184,13 @@ pub(crate) async fn get(
CodeOrError::Code { code } => code,
};
let http_service = http_client_factory
.http_service("upstream-discover")
.await?;
let http_service = http_client_factory.http_service().await?;
// XXX: we shouldn't discover on-the-fly
// Discover the provider
let metadata =
mas_oidc_client::requests::discovery::discover(&http_service, &provider.issuer).await?;
let http_service = http_client_factory
.http_service("upstream-fetch-jwks")
.await?;
// Fetch the JWKS
let jwks =
mas_oidc_client::requests::jose::fetch_jwks(&http_service, metadata.jwks_uri()).await?;
@ -227,10 +221,6 @@ pub(crate) async fn get(
client_id: &provider.client_id,
};
let http_service = http_client_factory
.http_service("upstream-exchange-code")
.await?;
let (response, id_token) =
mas_oidc_client::requests::authorization_code::access_token_with_authorization_code(
&http_service,

4
crates/http/Cargo.toml
View File

@ -18,8 +18,6 @@ hyper = "0.14.25"
hyper-rustls = { version = "0.23.2", features = ["http1", "http2"], default-features = false, optional = true }
once_cell = "1.17.1"
opentelemetry = "0.18.0"
opentelemetry-http = "0.7.0"
opentelemetry-semantic-conventions = "0.10.0"
rustls = { version = "0.20.8", optional = true }
rustls-native-certs = { version = "0.6.2", optional = true }
serde = "1.0.158"
@ -34,6 +32,8 @@ tracing-opentelemetry = "0.18.0"
webpki = { version = "0.22.0", optional = true }
webpki-roots = { version = "0.22.6", optional = true }
mas-tower = { path = "../tower" }
[dev-dependencies]
anyhow = "1.0.69"
serde = { version = "1.0.158", features = ["derive"] }

32
crates/http/src/client.rs
View File

@ -15,14 +15,20 @@
use std::convert::Infallible;
use hyper::{
client::{connect::dns::GaiResolver, HttpConnector},
client::{
connect::dns::{GaiResolver, Name},
HttpConnector,
},
Client,
};
use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
use mas_tower::{
DurationRecorderLayer, DurationRecorderService, FnWrapper, InFlightCounterLayer,
InFlightCounterService, TraceLayer, TraceService,
};
use thiserror::Error;
use tower::Layer;
use crate::layers::otel::{TraceDns, TraceLayer};
use tracing::Span;
#[cfg(all(not(feature = "webpki-roots"), not(feature = "native-roots")))]
compile_error!("enabling the 'client' feature requires also enabling the 'webpki-roots' or the 'native-roots' features");
@ -165,8 +171,10 @@ where
Ok(Client::builder().build(https))
}
pub type TraceResolver<S> =
InFlightCounterService<DurationRecorderService<TraceService<S, FnWrapper<fn(&Name) -> Span>>>>;
pub type UntracedConnector = HttpsConnector<HttpConnector<GaiResolver>>;
pub type TracedConnector = HttpsConnector<HttpConnector<TraceDns<GaiResolver>>>;
pub type TracedConnector = HttpsConnector<HttpConnector<TraceResolver<GaiResolver>>>;
/// Create a traced HTTP and HTTPS connector
///
@ -176,8 +184,20 @@ pub type TracedConnector = HttpsConnector<HttpConnector<TraceDns<GaiResolver>>>;
pub async fn make_traced_connector() -> Result<TracedConnector, ClientInitError>
where
{
// Trace DNS requests
let resolver = TraceLayer::dns().layer(GaiResolver::new());
let in_flight_counter = InFlightCounterLayer::new("dns.resolve.active_requests");
let duration_recorder = DurationRecorderLayer::new("dns.resolve.duration");
let trace_layer = TraceLayer::from_fn(
(|request: &Name| {
tracing::info_span!(
"dns.resolve",
"otel.kind" = "client",
"net.host.name" = %request,
)
}) as fn(&Name) -> Span,
);
let resolver = (in_flight_counter, duration_recorder, trace_layer).layer(GaiResolver::new());
let tls_config = make_tls_config().await?;
Ok(make_connector(resolver, tls_config))
}

48
crates/http/src/layers/client.rs
View File

@ -14,7 +14,8 @@
use std::{sync::Arc, time::Duration};
use http::{header::USER_AGENT, HeaderValue};
use http::{header::USER_AGENT, HeaderValue, Request};
use mas_tower::{MakeSpan, TraceContextLayer, TraceContextService, TraceLayer, TraceService};
use tokio::sync::Semaphore;
use tower::{
limit::{ConcurrencyLimit, GlobalConcurrencyLimitLayer},
@ -26,42 +27,61 @@ use tower_http::{
timeout::{Timeout, TimeoutLayer},
};
use super::otel::TraceLayer;
use crate::otel::{TraceHttpClient, TraceHttpClientLayer};
pub type ClientService<S> = SetRequestHeader<
TraceHttpClient<ConcurrencyLimit<FollowRedirect<TraceHttpClient<Timeout<S>>>>>,
ConcurrencyLimit<
FollowRedirect<TraceService<TraceContextService<Timeout<S>>, MakeSpanForRequest>>,
>,
HeaderValue,
>;
#[derive(Debug, Clone)]
pub struct MakeSpanForRequest;
impl<B> MakeSpan<Request<B>> for MakeSpanForRequest {
fn make_span(&self, request: &Request<B>) -> tracing::Span {
// TODO: better attributes
tracing::info_span!(
"http.client.request",
"http.method" = %request.method(),
"http.uri" = %request.uri(),
)
}
}
#[derive(Debug, Clone)]
pub struct ClientLayer {
user_agent_layer: SetRequestHeaderLayer<HeaderValue>,
outer_trace_layer: TraceHttpClientLayer,
concurrency_limit_layer: GlobalConcurrencyLimitLayer,
follow_redirect_layer: FollowRedirectLayer,
inner_trace_layer: TraceHttpClientLayer,
trace_layer: TraceLayer<MakeSpanForRequest>,
trace_context_layer: TraceContextLayer,
timeout_layer: TimeoutLayer,
}
impl Default for ClientLayer {
fn default() -> Self {
Self::new()
}
}
impl ClientLayer {
#[must_use]
pub fn new(operation: &'static str) -> Self {
pub fn new() -> Self {
let semaphore = Arc::new(Semaphore::new(10));
Self::with_semaphore(operation, semaphore)
Self::with_semaphore(semaphore)
}
#[must_use]
pub fn with_semaphore(operation: &'static str, semaphore: Arc<Semaphore>) -> Self {
pub fn with_semaphore(semaphore: Arc<Semaphore>) -> Self {
Self {
user_agent_layer: SetRequestHeaderLayer::overriding(
USER_AGENT,
HeaderValue::from_static("matrix-authentication-service/0.0.1"),
),
outer_trace_layer: TraceLayer::http_client(operation),
concurrency_limit_layer: GlobalConcurrencyLimitLayer::with_semaphore(semaphore),
follow_redirect_layer: FollowRedirectLayer::new(),
inner_trace_layer: TraceLayer::inner_http_client(),
trace_layer: TraceLayer::new(MakeSpanForRequest),
trace_context_layer: TraceContextLayer::new(),
timeout_layer: TimeoutLayer::new(Duration::from_secs(10)),
}
}
@ -76,10 +96,10 @@ where
fn layer(&self, inner: S) -> Self::Service {
(
&self.user_agent_layer,
&self.outer_trace_layer,
&self.concurrency_limit_layer,
&self.follow_redirect_layer,
&self.inner_trace_layer,
&self.trace_layer,
&self.trace_context_layer,
&self.timeout_layer,
)
.layer(inner)

1
crates/http/src/layers/mod.rs
View File

@ -18,7 +18,6 @@ pub mod catch_http_codes;
pub mod form_urlencoded_request;
pub mod json_request;
pub mod json_response;
pub mod otel;
#[cfg(feature = "client")]
pub(crate) mod client;

51
crates/http/src/layers/otel/extract_context.rs
View File

@ -1,51 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use http::Request;
use opentelemetry::{trace::TraceContextExt, Context};
use opentelemetry_http::HeaderExtractor;
pub trait ExtractContext<R> {
fn extract_context(&self, request: &R) -> Context;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultExtractContext;
impl<T> ExtractContext<T> for DefaultExtractContext {
fn extract_context(&self, _request: &T) -> Context {
Context::current()
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ExtractFromHttpRequest;
impl<T> ExtractContext<Request<T>> for ExtractFromHttpRequest {
fn extract_context(&self, request: &Request<T>) -> Context {
let headers = request.headers();
let extractor = HeaderExtractor(headers);
let parent_cx = Context::current();
let cx = opentelemetry::global::get_text_map_propagator(|propagator| {
propagator.extract_with_context(&parent_cx, &extractor)
});
if cx.span().span_context().is_remote() {
cx
} else {
parent_cx
}
}
}

76
crates/http/src/layers/otel/inject_context.rs
View File

@ -1,76 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use http::Request;
use opentelemetry::Context;
use opentelemetry_http::HeaderInjector;
pub trait InjectContext<R> {
type Output;
fn inject_context(&self, cx: &Context, request: R) -> Self::Output;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultInjectContext;
impl<R> InjectContext<R> for DefaultInjectContext {
type Output = R;
fn inject_context(&self, _cx: &Context, request: R) -> Self::Output {
request
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct InjectInHttpRequest;
impl<T> InjectContext<Request<T>> for InjectInHttpRequest {
type Output = Request<T>;
fn inject_context(&self, cx: &Context, mut request: Request<T>) -> Self::Output {
let headers = request.headers_mut();
let mut injector = HeaderInjector(headers);
opentelemetry::global::get_text_map_propagator(|propagator| {
propagator.inject_context(cx, &mut injector);
});
request
}
}
#[cfg(feature = "aws-sdk")]
#[derive(Debug, Clone, Copy, Default)]
pub struct InjectInAwsRequest;
#[cfg(feature = "aws-sdk")]
impl InjectContext<aws_smithy_http::operation::Request> for InjectInAwsRequest {
type Output = aws_smithy_http::operation::Request;
fn inject_context(
&self,
cx: &Context,
mut request: aws_smithy_http::operation::Request,
) -> Self::Output {
let headers = request.http_mut().headers_mut();
let mut injector = HeaderInjector(headers);
opentelemetry::global::get_text_map_propagator(|propagator| {
propagator.inject_context(cx, &mut injector);
});
request
}
}

334
crates/http/src/layers/otel/layer.rs
View File

@ -1,334 +0,0 @@
use std::sync::Arc;
use opentelemetry::{
metrics::{Counter, Histogram, UpDownCounter},
KeyValue,
};
use tower::Layer;
use super::{
extract_context::DefaultExtractContext, inject_context::DefaultInjectContext,
make_metrics_labels::DefaultMakeMetricsLabels, make_span_builder::DefaultMakeSpanBuilder,
on_error::DefaultOnError, on_response::DefaultOnResponse, service::Trace,
};
#[derive(Debug, Clone)]
pub struct TraceLayer<
ExtractContext = DefaultExtractContext,
InjectContext = DefaultInjectContext,
MakeSpanBuilder = DefaultMakeSpanBuilder,
MakeMetricsLabels = DefaultMakeMetricsLabels,
OnResponse = DefaultOnResponse,
OnError = DefaultOnError,
> {
tracer: Arc<opentelemetry::global::BoxedTracer>,
extract_context: ExtractContext,
inject_context: InjectContext,
make_span_builder: MakeSpanBuilder,
make_metrics_labels: MakeMetricsLabels,
on_response: OnResponse,
on_error: OnError,
inflight_requests: UpDownCounter<i64>,
request_counter: Counter<u64>,
request_histogram: Histogram<f64>,
static_attributes: Vec<KeyValue>,
}
impl Default for TraceLayer {
fn default() -> Self {
Self::with_namespace("http")
}
}
impl TraceLayer {
#[must_use]
pub fn with_namespace(namespace: &'static str) -> Self {
let tracer = Arc::new(opentelemetry::global::tracer("mas-http"));
let meter = opentelemetry::global::meter("mas-http");
let inflight_requests = meter
.i64_up_down_counter(format!("{namespace}.inflight_requests"))
.with_description("Number of in-flight requests")
.init();
let request_counter = meter
.u64_counter(format!("{namespace}.requests_total"))
.with_description("Total number of requests made.")
.init();
let request_histogram = meter
.f64_histogram(format!("{namespace}.request_duration_seconds"))
.with_description("The request latencies in seconds.")
.init();
Self::new(
tracer,
inflight_requests,
request_counter,
request_histogram,
)
}
}
impl<ExtractContext, InjectContext, MakeSpanBuilder, MakeMetricsLabels, OnResponse, OnError>
TraceLayer<
ExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
>
{
#[must_use]
pub fn new(
tracer: Arc<opentelemetry::global::BoxedTracer>,
inflight_requests: UpDownCounter<i64>,
request_counter: Counter<u64>,
request_histogram: Histogram<f64>,
) -> Self
where
ExtractContext: Default,
InjectContext: Default,
MakeSpanBuilder: Default,
MakeMetricsLabels: Default,
OnResponse: Default,
OnError: Default,
{
Self {
tracer,
extract_context: ExtractContext::default(),
inject_context: InjectContext::default(),
make_span_builder: MakeSpanBuilder::default(),
make_metrics_labels: MakeMetricsLabels::default(),
on_response: OnResponse::default(),
on_error: OnError::default(),
inflight_requests,
request_counter,
request_histogram,
static_attributes: Vec::new(),
}
}
#[must_use]
pub fn with_static_attribute(mut self, attribute: KeyValue) -> Self {
self.static_attributes.push(attribute);
self
}
#[must_use]
pub fn with_static_attributes(
mut self,
attributes: impl IntoIterator<Item = KeyValue>,
) -> Self {
self.static_attributes.extend(attributes);
self
}
#[must_use]
pub fn extract_context<NewExtractContext>(
self,
extract_context: NewExtractContext,
) -> TraceLayer<
NewExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
> {
TraceLayer {
tracer: self.tracer,
extract_context,
inject_context: self.inject_context,
make_span_builder: self.make_span_builder,
make_metrics_labels: self.make_metrics_labels,
on_response: self.on_response,
on_error: self.on_error,
inflight_requests: self.inflight_requests,
request_counter: self.request_counter,
request_histogram: self.request_histogram,
static_attributes: self.static_attributes,
}
}
#[must_use]
pub fn inject_context<NewInjectContext>(
self,
inject_context: NewInjectContext,
) -> TraceLayer<
ExtractContext,
NewInjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
> {
TraceLayer {
tracer: self.tracer,
extract_context: self.extract_context,
inject_context,
make_span_builder: self.make_span_builder,
make_metrics_labels: self.make_metrics_labels,
on_response: self.on_response,
on_error: self.on_error,
inflight_requests: self.inflight_requests,
request_counter: self.request_counter,
request_histogram: self.request_histogram,
static_attributes: self.static_attributes,
}
}
#[must_use]
pub fn make_span_builder<NewMakeSpanBuilder>(
self,
make_span_builder: NewMakeSpanBuilder,
) -> TraceLayer<
ExtractContext,
InjectContext,
NewMakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
> {
TraceLayer {
tracer: self.tracer,
extract_context: self.extract_context,
inject_context: self.inject_context,
make_span_builder,
make_metrics_labels: self.make_metrics_labels,
on_response: self.on_response,
on_error: self.on_error,
inflight_requests: self.inflight_requests,
request_counter: self.request_counter,
request_histogram: self.request_histogram,
static_attributes: self.static_attributes,
}
}
#[must_use]
pub fn make_metrics_labels<NewMakeMetricsLabels>(
self,
make_metrics_labels: NewMakeMetricsLabels,
) -> TraceLayer<
ExtractContext,
InjectContext,
MakeSpanBuilder,
NewMakeMetricsLabels,
OnResponse,
OnError,
> {
TraceLayer {
tracer: self.tracer,
extract_context: self.extract_context,
inject_context: self.inject_context,
make_span_builder: self.make_span_builder,
make_metrics_labels,
on_response: self.on_response,
on_error: self.on_error,
inflight_requests: self.inflight_requests,
request_counter: self.request_counter,
request_histogram: self.request_histogram,
static_attributes: self.static_attributes,
}
}
#[must_use]
pub fn on_response<NewOnResponse>(
self,
on_response: NewOnResponse,
) -> TraceLayer<
ExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
NewOnResponse,
OnError,
> {
TraceLayer {
tracer: self.tracer,
extract_context: self.extract_context,
inject_context: self.inject_context,
make_span_builder: self.make_span_builder,
make_metrics_labels: self.make_metrics_labels,
on_response,
on_error: self.on_error,
inflight_requests: self.inflight_requests,
request_counter: self.request_counter,
request_histogram: self.request_histogram,
static_attributes: self.static_attributes,
}
}
#[must_use]
pub fn on_error<NewOnError>(
self,
on_error: NewOnError,
) -> TraceLayer<
ExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
NewOnError,
> {
TraceLayer {
tracer: self.tracer,
extract_context: self.extract_context,
inject_context: self.inject_context,
make_span_builder: self.make_span_builder,
make_metrics_labels: self.make_metrics_labels,
on_response: self.on_response,
on_error,
inflight_requests: self.inflight_requests,
request_counter: self.request_counter,
request_histogram: self.request_histogram,
static_attributes: self.static_attributes,
}
}
}
impl<ExtractContext, InjectContext, MakeSpanBuilder, MakeMetricsLabels, OnResponse, OnError, S>
Layer<S>
for TraceLayer<
ExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
>
where
ExtractContext: Clone,
InjectContext: Clone,
MakeSpanBuilder: Clone,
MakeMetricsLabels: Clone,
OnResponse: Clone,
OnError: Clone,
{
type Service = Trace<
ExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
S,
>;
fn layer(&self, inner: S) -> Self::Service {
Trace::new(
inner,
self.tracer.clone(),
self.extract_context.clone(),
self.inject_context.clone(),
self.make_span_builder.clone(),
self.make_metrics_labels.clone(),
self.on_response.clone(),
self.on_error.clone(),
self.inflight_requests.clone(),
self.request_counter.clone(),
self.request_histogram.clone(),
self.static_attributes.clone(),
)
}
}

62
crates/http/src/layers/otel/make_metrics_labels.rs
View File

@ -1,62 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#[cfg(feature = "axum")]
use std::borrow::Cow;
use http::Request;
use opentelemetry::KeyValue;
use super::utils::http_method_str;
pub trait MakeMetricsLabels<R> {
fn make_metrics_labels(&self, request: &R) -> Vec<KeyValue>;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultMakeMetricsLabels;
impl<R> MakeMetricsLabels<R> for DefaultMakeMetricsLabels {
fn make_metrics_labels(&self, _request: &R) -> Vec<KeyValue> {
Vec::new()
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct MetricsLabelsFromHttpRequest;
impl<B> MakeMetricsLabels<Request<B>> for MetricsLabelsFromHttpRequest {
fn make_metrics_labels(&self, request: &Request<B>) -> Vec<KeyValue> {
vec![KeyValue::new("method", http_method_str(request.method()))]
}
}
#[cfg(feature = "axum")]
#[derive(Debug, Clone, Copy, Default)]
pub struct MetricsLabelsFromAxumRequest;
#[cfg(feature = "axum")]
impl<B> MakeMetricsLabels<Request<B>> for MetricsLabelsFromAxumRequest {
fn make_metrics_labels(&self, request: &Request<B>) -> Vec<KeyValue> {
let path: Cow<'static, str> = request
.extensions()
.get::<axum::extract::MatchedPath>()
.map_or("FALLBACK".into(), |path| path.as_str().to_owned().into());
vec![
KeyValue::new("method", http_method_str(request.method())),
KeyValue::new("route", path),
]
}
}

235
crates/http/src/layers/otel/make_span_builder.rs
View File

@ -1,235 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#[cfg(any(feature = "axum", feature = "aws-sdk"))]
use std::borrow::Cow;
#[cfg(feature = "axum")]
use axum::extract::{ConnectInfo, MatchedPath};
use headers::{ContentLength, HeaderMapExt, Host, UserAgent};
use http::Request;
#[cfg(feature = "client")]
use hyper::client::connect::dns::Name;
use opentelemetry::trace::{SpanBuilder, SpanKind};
use opentelemetry_semantic_conventions::trace as SC;
use super::utils::{http_flavor, http_method_str};
pub trait MakeSpanBuilder<R> {
fn make_span_builder(&self, request: &R) -> SpanBuilder;
}
#[derive(Debug, Clone, Copy)]
pub struct DefaultMakeSpanBuilder {
operation: &'static str,
}
impl DefaultMakeSpanBuilder {
#[must_use]
pub fn new(operation: &'static str) -> Self {
Self { operation }
}
}
impl Default for DefaultMakeSpanBuilder {
fn default() -> Self {
Self {
operation: "service",
}
}
}
impl<R> MakeSpanBuilder<R> for DefaultMakeSpanBuilder {
fn make_span_builder(&self, _request: &R) -> SpanBuilder {
SpanBuilder::from_name(self.operation)
}
}
#[derive(Debug, Clone)]
pub struct SpanFromHttpRequest {
operation: &'static str,
span_kind: SpanKind,
}
impl SpanFromHttpRequest {
#[must_use]
pub fn server() -> Self {
Self {
operation: "http-server",
span_kind: SpanKind::Server,
}
}
#[must_use]
pub fn inner_client() -> Self {
Self {
operation: "http-client",
span_kind: SpanKind::Client,
}
}
#[must_use]
pub fn client(operation: &'static str) -> Self {
Self {
operation,
span_kind: SpanKind::Client,
}
}
}
impl<B> MakeSpanBuilder<Request<B>> for SpanFromHttpRequest {
fn make_span_builder(&self, request: &Request<B>) -> SpanBuilder {
let mut attributes = vec![
SC::HTTP_METHOD.string(http_method_str(request.method())),
SC::HTTP_FLAVOR.string(http_flavor(request.version())),
SC::HTTP_TARGET.string(request.uri().to_string()),
];
let headers = request.headers();
if let Some(host) = headers.typed_get::<Host>() {
attributes.push(SC::HTTP_HOST.string(host.to_string()));
}
if let Some(user_agent) = headers.typed_get::<UserAgent>() {
attributes.push(SC::HTTP_USER_AGENT.string(user_agent.to_string()));
}
if let Some(ContentLength(content_length)) = headers.typed_get() {
if let Ok(content_length) = content_length.try_into() {
attributes.push(SC::HTTP_REQUEST_CONTENT_LENGTH.i64(content_length));
}
}
SpanBuilder::from_name(self.operation)
.with_kind(self.span_kind.clone())
.with_attributes(attributes)
}
}
#[cfg(feature = "axum")]
#[derive(Debug, Clone)]
pub struct SpanFromAxumRequest;
#[cfg(feature = "axum")]
impl<B> MakeSpanBuilder<Request<B>> for SpanFromAxumRequest {
fn make_span_builder(&self, request: &Request<B>) -> SpanBuilder {
let (name, route): (String, Cow<'static, str>) =
if let Some(path) = request.extensions().get::<MatchedPath>() {
let path = path.as_str().to_owned();
let name = path.clone();
(name, path.into())
} else {
(request.uri().path().to_owned(), Cow::Borrowed("FALLBACK"))
};
let mut attributes = vec![
SC::HTTP_METHOD.string(http_method_str(request.method())),
SC::HTTP_FLAVOR.string(http_flavor(request.version())),
SC::HTTP_TARGET.string(request.uri().to_string()),
SC::HTTP_ROUTE.string(route),
];
let headers = request.headers();
if let Some(host) = headers.typed_get::<Host>() {
attributes.push(SC::HTTP_HOST.string(host.to_string()));
}
if let Some(user_agent) = headers.typed_get::<UserAgent>() {
attributes.push(SC::HTTP_USER_AGENT.string(user_agent.to_string()));
}
if let Some(ContentLength(content_length)) = headers.typed_get() {
if let Ok(content_length) = content_length.try_into() {
attributes.push(SC::HTTP_REQUEST_CONTENT_LENGTH.i64(content_length));
}
}
if let Some(ConnectInfo(addr)) = request
.extensions()
.get::<ConnectInfo<std::net::SocketAddr>>()
{
attributes.push(SC::NET_TRANSPORT.string("ip_tcp"));
attributes.push(SC::NET_PEER_IP.string(addr.ip().to_string()));
attributes.push(SC::NET_PEER_PORT.i64(addr.port().into()));
}
SpanBuilder::from_name(name)
.with_kind(SpanKind::Server)
.with_attributes(attributes)
}
}
#[cfg(feature = "client")]
#[derive(Debug, Clone, Copy, Default)]
pub struct SpanFromDnsRequest;
#[cfg(feature = "client")]
impl MakeSpanBuilder<Name> for SpanFromDnsRequest {
fn make_span_builder(&self, request: &Name) -> SpanBuilder {
let attributes = vec![SC::NET_HOST_NAME.string(request.as_str().to_owned())];
SpanBuilder::from_name("resolve")
.with_kind(SpanKind::Client)
.with_attributes(attributes)
}
}
#[cfg(feature = "aws-sdk")]
#[derive(Debug, Clone, Copy, Default)]
pub struct SpanFromAwsRequest;
#[cfg(feature = "aws-sdk")]
impl MakeSpanBuilder<aws_smithy_http::operation::Request> for SpanFromAwsRequest {
fn make_span_builder(&self, request: &aws_smithy_http::operation::Request) -> SpanBuilder {
let properties = request.properties();
let request = request.http();
let mut attributes = vec![
SC::RPC_SYSTEM.string("aws-api"),
SC::HTTP_METHOD.string(http_method_str(request.method())),
SC::HTTP_FLAVOR.string(http_flavor(request.version())),
SC::HTTP_TARGET.string(request.uri().to_string()),
];
let mut name = Cow::Borrowed("aws_sdk");
if let Some(metadata) = properties.get::<aws_smithy_http::operation::Metadata>() {
attributes.push(SC::RPC_SERVICE.string(metadata.service().to_owned()));
attributes.push(SC::RPC_METHOD.string(metadata.name().to_owned()));
name = Cow::Owned(metadata.name().to_owned());
} else if let Some(service) = properties.get::<aws_types::SigningService>() {
attributes.push(SC::RPC_SERVICE.string(service.as_ref().to_owned()));
}
let headers = request.headers();
if let Some(host) = headers.typed_get::<Host>() {
attributes.push(SC::HTTP_HOST.string(host.to_string()));
}
if let Some(user_agent) = headers.typed_get::<UserAgent>() {
attributes.push(SC::HTTP_USER_AGENT.string(user_agent.to_string()));
}
if let Some(ContentLength(content_length)) = headers.typed_get() {
if let Ok(content_length) = content_length.try_into() {
attributes.push(SC::HTTP_REQUEST_CONTENT_LENGTH.i64(content_length));
}
}
SpanBuilder::from_name(name)
.with_kind(SpanKind::Client)
.with_attributes(attributes)
}
}

197
crates/http/src/layers/otel/mod.rs
View File

@ -1,197 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
mod extract_context;
mod inject_context;
mod layer;
mod make_metrics_labels;
mod make_span_builder;
mod on_error;
mod on_response;
mod service;
mod utils;
pub type TraceHttpServerLayer = TraceLayer<
ExtractFromHttpRequest,
DefaultInjectContext,
SpanFromHttpRequest,
MetricsLabelsFromHttpRequest,
OnHttpResponse,
DefaultOnError,
>;
pub type TraceHttpServer<S> = Trace<
ExtractFromHttpRequest,
DefaultInjectContext,
SpanFromHttpRequest,
MetricsLabelsFromHttpRequest,
OnHttpResponse,
DefaultOnError,
S,
>;
#[cfg(feature = "axum")]
pub type TraceAxumServerLayer = TraceLayer<
ExtractFromHttpRequest,
DefaultInjectContext,
SpanFromAxumRequest,
MetricsLabelsFromAxumRequest,
OnHttpResponse,
DefaultOnError,
>;
#[cfg(feature = "axum")]
pub type TraceAxumServer<S> = Trace<
ExtractFromHttpRequest,
DefaultInjectContext,
SpanFromAxumRequest,
MetricsLabelsFromAxumRequest,
OnHttpResponse,
DefaultOnError,
S,
>;
pub type TraceHttpClientLayer = TraceLayer<
DefaultExtractContext,
InjectInHttpRequest,
SpanFromHttpRequest,
MetricsLabelsFromHttpRequest,
OnHttpResponse,
DefaultOnError,
>;
pub type TraceHttpClient<S> = Trace<
DefaultExtractContext,
InjectInHttpRequest,
SpanFromHttpRequest,
MetricsLabelsFromHttpRequest,
OnHttpResponse,
DefaultOnError,
S,
>;
#[cfg(feature = "client")]
pub type TraceDnsLayer = TraceLayer<
DefaultExtractContext,
DefaultInjectContext,
SpanFromDnsRequest,
DefaultMakeMetricsLabels,
DefaultOnResponse,
DefaultOnError,
>;
#[cfg(feature = "client")]
pub type TraceDns<S> = Trace<
DefaultExtractContext,
DefaultInjectContext,
SpanFromDnsRequest,
DefaultMakeMetricsLabels,
DefaultOnResponse,
DefaultOnError,
S,
>;
#[cfg(feature = "aws-sdk")]
pub type TraceAwsSdkClientLayer = TraceLayer<
DefaultExtractContext,
InjectInAwsRequest,
SpanFromAwsRequest,
DefaultMakeMetricsLabels,
OnAwsResponse,
DebugOnError,
>;
#[cfg(feature = "aws-sdk")]
pub type TraceAwsSdkClient<S> = Trace<
DefaultExtractContext,
InjectInAwsRequest,
SpanFromAwsRequest,
DefaultMakeMetricsLabels,
OnAwsResponse,
DebugOnError,
S,
>;
impl TraceHttpServerLayer {
#[must_use]
pub fn http_server() -> Self {
TraceLayer::with_namespace("http_server")
.make_span_builder(SpanFromHttpRequest::server())
.make_metrics_labels(MetricsLabelsFromHttpRequest::default())
.on_response(OnHttpResponse)
.extract_context(ExtractFromHttpRequest)
}
}
#[cfg(feature = "axum")]
impl TraceAxumServerLayer {
#[must_use]
pub fn axum() -> Self {
TraceLayer::with_namespace("http_server")
.make_span_builder(SpanFromAxumRequest)
.make_metrics_labels(MetricsLabelsFromAxumRequest::default())
.on_response(OnHttpResponse)
.extract_context(ExtractFromHttpRequest)
}
}
impl TraceHttpClientLayer {
#[must_use]
pub fn http_client(operation: &'static str) -> Self {
TraceLayer::with_namespace("http_client")
.make_span_builder(SpanFromHttpRequest::client(operation))
.make_metrics_labels(MetricsLabelsFromHttpRequest::default())
.on_response(OnHttpResponse)
.inject_context(InjectInHttpRequest)
}
#[must_use]
pub fn inner_http_client() -> Self {
TraceLayer::with_namespace("inner_http_client")
.make_span_builder(SpanFromHttpRequest::inner_client())
.make_metrics_labels(MetricsLabelsFromHttpRequest::default())
.on_response(OnHttpResponse)
.inject_context(InjectInHttpRequest)
}
}
#[cfg(feature = "client")]
impl TraceDnsLayer {
#[must_use]
pub fn dns() -> Self {
TraceLayer::with_namespace("dns").make_span_builder(SpanFromDnsRequest)
}
}
#[cfg(feature = "aws-sdk")]
impl TraceAwsSdkClientLayer {
#[must_use]
pub fn aws_sdk() -> Self {
TraceLayer::with_namespace("aws_sdk")
.make_span_builder(SpanFromAwsRequest)
.on_response(OnAwsResponse)
.on_error(DebugOnError)
.inject_context(InjectInAwsRequest)
}
}
#[cfg(feature = "client")]
use self::make_metrics_labels::DefaultMakeMetricsLabels;
#[cfg(feature = "axum")]
use self::make_metrics_labels::MetricsLabelsFromAxumRequest;
use self::make_metrics_labels::MetricsLabelsFromHttpRequest;
pub use self::{
extract_context::*, inject_context::*, layer::*, make_span_builder::*, on_error::*,
on_response::*, service::*,
};

46
crates/http/src/layers/otel/on_error.rs
View File

@ -1,46 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use opentelemetry::{trace::SpanRef, KeyValue};
use opentelemetry_semantic_conventions::trace::EXCEPTION_MESSAGE;
pub trait OnError<E> {
fn on_error(&self, span: &SpanRef<'_>, metrics_labels: &mut Vec<KeyValue>, err: &E);
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultOnError;
impl<E> OnError<E> for DefaultOnError
where
E: std::fmt::Display,
{
fn on_error(&self, span: &SpanRef<'_>, _metrics_labels: &mut Vec<KeyValue>, err: &E) {
let attributes = vec![EXCEPTION_MESSAGE.string(err.to_string())];
span.add_event("exception".to_owned(), attributes);
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DebugOnError;
impl<E> OnError<E> for DebugOnError
where
E: std::fmt::Debug,
{
fn on_error(&self, span: &SpanRef<'_>, _metrics_labels: &mut Vec<KeyValue>, err: &E) {
let attributes = vec![EXCEPTION_MESSAGE.string(format!("{err:?}"))];
span.add_event("exception".to_owned(), attributes);
}
}

99
crates/http/src/layers/otel/on_response.rs
View File

@ -1,99 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use headers::{ContentLength, HeaderMapExt};
use http::Response;
#[cfg(feature = "client")]
use hyper::client::connect::HttpInfo;
use opentelemetry::{trace::SpanRef, KeyValue};
use opentelemetry_semantic_conventions::trace as SC;
pub trait OnResponse<R> {
fn on_response(&self, span: &SpanRef<'_>, metrics_labels: &mut Vec<KeyValue>, response: &R);
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultOnResponse;
impl<R> OnResponse<R> for DefaultOnResponse {
fn on_response(&self, _span: &SpanRef<'_>, _metrics_labels: &mut Vec<KeyValue>, _response: &R) {
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct OnHttpResponse;
impl<B> OnResponse<Response<B>> for OnHttpResponse {
fn on_response(
&self,
span: &SpanRef<'_>,
metrics_labels: &mut Vec<KeyValue>,
response: &Response<B>,
) {
let status_code = i64::from(response.status().as_u16());
span.set_attribute(SC::HTTP_STATUS_CODE.i64(status_code));
metrics_labels.push(KeyValue::new("status_code", status_code));
if let Some(ContentLength(content_length)) = response.headers().typed_get() {
if let Ok(content_length) = content_length.try_into() {
span.set_attribute(SC::HTTP_RESPONSE_CONTENT_LENGTH.i64(content_length));
}
}
#[cfg(feature = "client")]
// Get local and remote address from hyper's HttpInfo injected by the
// HttpConnector
if let Some(info) = response.extensions().get::<HttpInfo>() {
span.set_attribute(SC::NET_PEER_IP.string(info.remote_addr().ip().to_string()));
span.set_attribute(SC::NET_PEER_PORT.i64(info.remote_addr().port().into()));
span.set_attribute(SC::NET_HOST_IP.string(info.local_addr().ip().to_string()));
span.set_attribute(SC::NET_HOST_PORT.i64(info.local_addr().port().into()));
}
}
}
#[cfg(feature = "aws-sdk")]
#[derive(Debug, Clone, Copy, Default)]
pub struct OnAwsResponse;
#[cfg(feature = "aws-sdk")]
impl OnResponse<aws_smithy_http::operation::Response> for OnAwsResponse {
fn on_response(
&self,
span: &SpanRef<'_>,
metrics_labels: &mut Vec<KeyValue>,
response: &aws_smithy_http::operation::Response,
) {
let response = response.http();
let status_code = i64::from(response.status().as_u16());
span.set_attribute(SC::HTTP_STATUS_CODE.i64(status_code));
metrics_labels.push(KeyValue::new("status_code", status_code));
if let Some(ContentLength(content_length)) = response.headers().typed_get() {
if let Ok(content_length) = content_length.try_into() {
span.set_attribute(SC::HTTP_RESPONSE_CONTENT_LENGTH.i64(content_length));
}
}
#[cfg(feature = "client")]
// Get local and remote address from hyper's HttpInfo injected by the
// HttpConnector
if let Some(info) = response.extensions().get::<HttpInfo>() {
span.set_attribute(SC::NET_PEER_IP.string(info.remote_addr().ip().to_string()));
span.set_attribute(SC::NET_PEER_PORT.i64(info.remote_addr().port().into()));
span.set_attribute(SC::NET_HOST_IP.string(info.local_addr().ip().to_string()));
span.set_attribute(SC::NET_HOST_PORT.i64(info.local_addr().port().into()));
}
}
}

208
crates/http/src/layers/otel/service.rs
View File

@ -1,208 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::{sync::Arc, task::Poll, time::SystemTime};
use futures_util::{future::BoxFuture, FutureExt as _};
use opentelemetry::{
metrics::{Counter, Histogram, UpDownCounter},
trace::{FutureExt as _, TraceContextExt},
Context, KeyValue,
};
use tower::Service;
use super::{
extract_context::ExtractContext, inject_context::InjectContext,
make_metrics_labels::MakeMetricsLabels, make_span_builder::MakeSpanBuilder, on_error::OnError,
on_response::OnResponse,
};
#[derive(Debug, Clone)]
pub struct Trace<
ExtractContext,
InjectContext,
MakeSpanBuilder,
MakeMetricsLabels,
OnResponse,
OnError,
S,
> {
inner: S,
tracer: Arc<opentelemetry::global::BoxedTracer>,
extract_context: ExtractContext,
inject_context: InjectContext,
make_span_builder: MakeSpanBuilder,
make_metrics_labels: MakeMetricsLabels,
on_response: OnResponse,
on_error: OnError,
inflight_requests: UpDownCounter<i64>,
request_counter: Counter<u64>,
request_histogram: Histogram<f64>,
static_attributes: Vec<KeyValue>,
}
impl<ExtractContext, InjectContext, MakeSpanBuilder, MakeMetricsLabels, OnResponse, OnError, S>
Trace<ExtractContext, InjectContext, MakeSpanBuilder, MakeMetricsLabels, OnResponse, OnError, S>
{
#[allow(clippy::too_many_arguments)]
pub fn new(
service: S,
tracer: Arc<opentelemetry::global::BoxedTracer>,
extract_context: ExtractContext,
inject_context: InjectContext,
make_span_builder: MakeSpanBuilder,
make_metrics_labels: MakeMetricsLabels,
on_response: OnResponse,
on_error: OnError,
inflight_requests: UpDownCounter<i64>,
request_counter: Counter<u64>,
request_histogram: Histogram<f64>,
static_attributes: Vec<KeyValue>,
) -> Self {
Self {
inner: service,
tracer,
extract_context,
inject_context,
make_span_builder,
make_metrics_labels,
on_response,
on_error,
inflight_requests,
request_counter,
request_histogram,
static_attributes,
}
}
}
struct InFlightGuard {
context: Context,
meter: UpDownCounter<i64>,
attributes: Vec<KeyValue>,
}
impl InFlightGuard {
fn increment(context: &Context, meter: &UpDownCounter<i64>, attributes: &[KeyValue]) -> Self {
meter.add(context, 1, attributes);
Self {
context: context.clone(),
meter: meter.clone(),
attributes: attributes.to_vec(),
}
}
}
impl Drop for InFlightGuard {
fn drop(&mut self) {
self.meter.add(&self.context, -1, &self.attributes);
}
}
impl<
Req,
S,
ExtractContextT,
InjectContextT,
MakeSpanBuilderT,
MakeMetricsLabelsT,
OnResponseT,
OnErrorT,
> Service<Req>
for Trace<
ExtractContextT,
InjectContextT,
MakeSpanBuilderT,
MakeMetricsLabelsT,
OnResponseT,
OnErrorT,
S,
>
where
ExtractContextT: ExtractContext<Req> + Send,
InjectContextT: InjectContext<Req> + Send,
S: Service<InjectContextT::Output> + Send,
OnResponseT: OnResponse<S::Response> + Send + Clone + 'static,
OnErrorT: OnError<S::Error> + Send + Clone + 'static,
MakeSpanBuilderT: MakeSpanBuilder<Req> + Send,
MakeMetricsLabelsT: MakeMetricsLabels<Req> + Send,
S::Future: Send + 'static,
{
type Response = S::Response;
type Error = S::Error;
type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, request: Req) -> Self::Future {
let request_counter = self.request_counter.clone();
let request_histogram = self.request_histogram.clone();
let start_time = SystemTime::now();
let cx = self.extract_context.extract_context(&request);
let mut span_builder = self.make_span_builder.make_span_builder(&request);
let mut metrics_labels = self.make_metrics_labels.make_metrics_labels(&request);
// Add the static attributes to the metrics and the span
metrics_labels.extend_from_slice(&self.static_attributes[..]);
let mut span_attributes = span_builder.attributes.unwrap_or_default();
span_attributes.extend(self.static_attributes.iter().cloned());
span_builder.attributes = Some(span_attributes);
let span = span_builder.start_with_context(self.tracer.as_ref(), &cx);
let cx = cx.with_span(span);
let request = self.inject_context.inject_context(&cx, request);
let guard = InFlightGuard::increment(&cx, &self.inflight_requests, &metrics_labels);
let on_response = self.on_response.clone();
let on_error = self.on_error.clone();
let attachment = cx.clone().attach();
let ret = self
.inner
.call(request)
.with_context(cx.clone())
.inspect(move |r| {
// This ensures the guard gets moved to the future. In case the future panics,
// it will be dropped anyway, ensuring the in-flight counter stays accurate
let _guard = guard;
let span = cx.span();
match r {
Ok(response) => on_response.on_response(&span, &mut metrics_labels, response),
Err(err) => on_error.on_error(&span, &mut metrics_labels, err),
};
request_counter.add(&cx, 1, &metrics_labels);
request_histogram.record(
&cx,
start_time.elapsed().map_or(0.0, |d| d.as_secs_f64()),
&metrics_labels,
);
span.end();
})
.boxed();
drop(attachment);
ret
}
}

45
crates/http/src/layers/otel/utils.rs
View File

@ -1,45 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::borrow::Cow;
use http::{Method, Version};
#[inline]
pub(super) fn http_method_str(method: &Method) -> Cow<'static, str> {
match method {
&Method::OPTIONS => "OPTIONS".into(),
&Method::GET => "GET".into(),
&Method::POST => "POST".into(),
&Method::PUT => "PUT".into(),
&Method::DELETE => "DELETE".into(),
&Method::HEAD => "HEAD".into(),
&Method::TRACE => "TRACE".into(),
&Method::CONNECT => "CONNECT".into(),
&Method::PATCH => "PATCH".into(),
other => other.to_string().into(),
}
}
#[inline]
pub(super) fn http_flavor(version: Version) -> Cow<'static, str> {
match version {
Version::HTTP_09 => "0.9".into(),
Version::HTTP_10 => "1.0".into(),
Version::HTTP_11 => "1.1".into(),
Version::HTTP_2 => "2.0".into(),
Version::HTTP_3 => "3.0".into(),
other => format!("{other:?}").into(),
}
}

3
crates/http/src/lib.rs
View File

@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
//! [`tower`] layers and services to help building HTTP client and servers
//! [`tower`] layers and services to help building HTTP client and servers
#![forbid(unsafe_code)]
#![deny(
@ -47,7 +47,6 @@ pub use self::{
form_urlencoded_request::{self, FormUrlencodedRequest, FormUrlencodedRequestLayer},
json_request::{self, JsonRequest, JsonRequestLayer},
json_response::{self, JsonResponse, JsonResponseLayer},
otel,
},
service::{BoxCloneSyncService, HttpService},
};

8
crates/tasks/Cargo.toml
View File

@ -19,13 +19,15 @@ thiserror = "1.0.40"
tower = "0.4.13"
tracing = "0.1.37"
tracing-opentelemetry = "0.18.0"
opentelemetry = "0.18.0"
ulid = "1.0.0"
url = "2.3.1"
serde = { version = "1.0.159", features = ["derive"] }
mas-axum-utils = { path = "../axum-utils" }
mas-storage = { path = "../storage" }
mas-storage-pg = { path = "../storage-pg" }
mas-data-model = { path = "../data-model" }
mas-email = { path = "../email" }
mas-http = { path = "../http" }
mas-data-model = { path = "../data-model" }
mas-storage = { path = "../storage" }
mas-storage-pg = { path = "../storage-pg" }
mas-tower = { path = "../tower" }

3
crates/tasks/src/database.rs
View File

@ -28,7 +28,7 @@ use chrono::{DateTime, Utc};
use mas_storage::{oauth2::OAuth2AccessTokenRepository, RepositoryAccess};
use tracing::{debug, info};
use crate::{JobContextExt, State};
use crate::{utils::metrics_layer, JobContextExt, State};
#[derive(Default, Clone)]
pub struct CleanupExpiredTokensJob {
@ -77,6 +77,7 @@ pub(crate) fn register(
let worker = WorkerBuilder::new(worker_name)
.stream(CronStream::new(schedule).to_stream())
.layer(state.inject())
.layer(metrics_layer::<CleanupExpiredTokensJob>())
.build_fn(cleanup_expired_tokens);
monitor.register(worker)

8
crates/tasks/src/email.rs
View File

@ -27,7 +27,10 @@ use mas_storage::job::{JobWithSpanContext, VerifyEmailJob};
use rand::{distributions::Uniform, Rng};
use tracing::info;
use crate::{layers::TracingLayer, JobContextExt, State};
use crate::{
utils::{metrics_layer, trace_layer},
JobContextExt, State,
};
#[tracing::instrument(
name = "job.verify_email",
@ -98,7 +101,8 @@ pub(crate) fn register(
let worker_name = format!("{job}-{suffix}", job = VerifyEmailJob::NAME);
let worker = WorkerBuilder::new(worker_name)
.layer(state.inject())
.layer(TracingLayer::new())
.layer(trace_layer::<VerifyEmailJob>())
.layer(metrics_layer::<JobWithSpanContext<VerifyEmailJob>>())
.with_storage(storage)
.build_fn(verify_email);
monitor.register(worker)

70
crates/tasks/src/layers.rs
View File

@ -1,70 +0,0 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::task::{Context, Poll};
use apalis_core::{job::Job, request::JobRequest};
use mas_storage::job::JobWithSpanContext;
use tower::{Layer, Service};
use tracing::{info_span, instrument::Instrumented, Instrument};
use tracing_opentelemetry::OpenTelemetrySpanExt;
pub struct TracingLayer;
impl TracingLayer {
pub fn new() -> Self {
Self
}
}
impl<S> Layer<S> for TracingLayer {
type Service = TracingService<S>;
fn layer(&self, inner: S) -> Self::Service {
TracingService { inner }
}
}
pub struct TracingService<S> {
inner: S,
}
impl<J, S> Service<JobRequest<JobWithSpanContext<J>>> for TracingService<S>
where
J: Job,
S: Service<JobRequest<JobWithSpanContext<J>>>,
{
type Response = S::Response;
type Error = S::Error;
type Future = Instrumented<S::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: JobRequest<JobWithSpanContext<J>>) -> Self::Future {
let span = info_span!(
"job.run",
job.id = %req.id(),
job.attempts = req.attempts(),
job.name = J::NAME,
);
if let Some(context) = req.inner().span_context() {
span.add_link(context);
}
self.inner.call(req).instrument(span)
}
}

9
crates/tasks/src/lib.rs
View File

@ -31,8 +31,8 @@ use tracing::debug;
mod database;
mod email;
mod layers;
mod matrix;
mod utils;
pub use self::matrix::HomeserverConnection;
@ -103,15 +103,12 @@ impl State {
&self.homeserver
}
pub async fn http_client<B>(
&self,
operation: &'static str,
) -> Result<ClientService<TracedClient<B>>, ClientInitError>
pub async fn http_client<B>(&self) -> Result<ClientService<TracedClient<B>>, ClientInitError>
where
B: mas_axum_utils::axum::body::HttpBody + Send,
B::Data: Send,
{
self.http_client_factory.client(operation).await
self.http_client_factory.client().await
}
}

43
crates/tasks/src/matrix.rs
View File

@ -33,10 +33,13 @@ use mas_storage::{
};
use serde::{Deserialize, Serialize};
use tower::{Service, ServiceExt};
use tracing::info;
use tracing::{info, info_span, Instrument};
use url::Url;
use crate::{layers::TracingLayer, JobContextExt, State};
use crate::{
utils::{metrics_layer, trace_layer},
JobContextExt, State,
};
pub struct HomeserverConnection {
homeserver: String,
@ -101,7 +104,7 @@ async fn provision_user(
let state = ctx.state();
let matrix = state.matrix_connection();
let mut client = state
.http_client("matrix.provision_user")
.http_client()
.await?
.request_bytes_to_body()
.json_request();
@ -158,7 +161,12 @@ async fn provision_user(
let req = req.body(body).context("Failed to build request")?;
let response = client.ready().await?.call(req).await?;
let response = client
.ready()
.await?
.call(req)
.instrument(info_span!("matrix.provision_user"))
.await?;
match response.status() {
StatusCode::CREATED => info!(%user.id, %mxid, "User created"),
@ -194,7 +202,7 @@ async fn provision_device(
let state = ctx.state();
let matrix = state.matrix_connection();
let mut client = state
.http_client("matrix.provision_device")
.http_client()
.await?
.request_bytes_to_body()
.json_request();
@ -225,7 +233,12 @@ async fn provision_device(
})
.context("Failed to build request")?;
let response = client.ready().await?.call(req).await?;
let response = client
.ready()
.await?
.call(req)
.instrument(info_span!("matrix.create_device"))
.await?;
match response.status() {
StatusCode::CREATED => {
@ -255,7 +268,7 @@ async fn delete_device(
) -> Result<(), anyhow::Error> {
let state = ctx.state();
let matrix = state.matrix_connection();
let mut client = state.http_client("matrix.delete_device").await?;
let mut client = state.http_client().await?;
let mut repo = state.repository().await?;
let user = repo
@ -284,7 +297,12 @@ async fn delete_device(
.body(EmptyBody::new())
.context("Failed to build request")?;
let response = client.ready().await?.call(req).await?;
let response = client
.ready()
.await?
.call(req)
.instrument(info_span!("matrix.delete_device"))
.await?;
match response.status() {
StatusCode::OK => info!(%user.id, %mxid, "Device deleted"),
@ -303,7 +321,8 @@ pub(crate) fn register(
let worker_name = format!("{job}-{suffix}", job = ProvisionUserJob::NAME);
let provision_user_worker = WorkerBuilder::new(worker_name)
.layer(state.inject())
.layer(TracingLayer::new())
.layer(trace_layer())
.layer(metrics_layer::<JobWithSpanContext<ProvisionUserJob>>())
.with_storage(storage)
.build_fn(provision_user);
@ -311,7 +330,8 @@ pub(crate) fn register(
let worker_name = format!("{job}-{suffix}", job = ProvisionDeviceJob::NAME);
let provision_device_worker = WorkerBuilder::new(worker_name)
.layer(state.inject())
.layer(TracingLayer::new())
.layer(trace_layer())
.layer(metrics_layer::<JobWithSpanContext<ProvisionDeviceJob>>())
.with_storage(storage)
.build_fn(provision_device);
@ -319,7 +339,8 @@ pub(crate) fn register(
let worker_name = format!("{job}-{suffix}", job = DeleteDeviceJob::NAME);
let delete_device_worker = WorkerBuilder::new(worker_name)
.layer(state.inject())
.layer(TracingLayer::new())
.layer(trace_layer())
.layer(metrics_layer::<JobWithSpanContext<DeleteDeviceJob>>())
.with_storage(storage)
.build_fn(delete_device);

79
crates/tasks/src/utils.rs Normal file
View File

@ -0,0 +1,79 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use apalis_core::{job::Job, request::JobRequest};
use mas_storage::job::JobWithSpanContext;
use mas_tower::{
make_span_fn, DurationRecorderLayer, FnWrapper, InFlightCounterLayer, TraceLayer, KV,
};
use opentelemetry::{Key, KeyValue};
use tracing::info_span;
use tracing_opentelemetry::OpenTelemetrySpanExt;
const JOB_NAME: Key = Key::from_static_str("job.name");
const JOB_STATUS: Key = Key::from_static_str("job.status");
fn make_span_for_job_request<J>(req: &JobRequest<JobWithSpanContext<J>>) -> tracing::Span
where
J: Job,
{
let span = info_span!(
"job.run",
"otel.kind" = "consumer",
"otel.status_code" = tracing::field::Empty,
"job.id" = %req.id(),
"job.attempts" = req.attempts(),
"job.name" = J::NAME,
);
if let Some(context) = req.inner().span_context() {
span.add_link(context);
}
span
}
type TraceLayerForJob<J> = TraceLayer<
FnWrapper<fn(&JobRequest<JobWithSpanContext<J>>) -> tracing::Span>,
KV<&'static str>,
KV<&'static str>,
>;
pub(crate) fn trace_layer<J>() -> TraceLayerForJob<J>
where
J: Job,
{
TraceLayer::new(make_span_fn(
make_span_for_job_request::<J> as fn(&JobRequest<JobWithSpanContext<J>>) -> tracing::Span,
))
.on_response(KV("otel.status_code", "OK"))
.on_error(KV("otel.status_code", "ERROR"))
}
pub(crate) fn metrics_layer<J>() -> (
DurationRecorderLayer<KeyValue, KeyValue, KeyValue>,
InFlightCounterLayer<KeyValue>,
)
where
J: Job,
{
let duration_recorder = DurationRecorderLayer::new("job.run.duration")
.on_request(JOB_NAME.string(J::NAME))
.on_response(JOB_STATUS.string("success"))
.on_error(JOB_STATUS.string("error"));
let in_flight_counter =
InFlightCounterLayer::new("job.run.active").on_request(JOB_NAME.string(J::NAME));
(duration_recorder, in_flight_counter)
}

20
crates/tower/Cargo.toml Normal file
View File

@ -0,0 +1,20 @@
[package]
name = "mas-tower"
version = "0.1.0"
authors = ["Quentin Gliech <quenting@element.io>"]
edition = "2021"
license = "Apache-2.0"
[dependencies]
aws-smithy-http = { version = "0.54.4", optional = true }
http = "0.2.9"
tracing = "0.1.37"
tracing-opentelemetry = "0.18.0"
tower = "0.4.13"
tokio = { version = "1.27.0", features = ["time"] }
opentelemetry = { version = "0.18.0", features = ["metrics"] }
opentelemetry-http = "0.7.0"
pin-project-lite = "0.2.9"
[features]
aws-sdk = ["dep:aws-smithy-http"]

32
crates/tower/src/lib.rs Normal file
View File

@ -0,0 +1,32 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#![deny(clippy::all)]
#![warn(clippy::pedantic)]
#![allow(clippy::module_name_repetitions)]
mod metrics;
mod trace_context;
mod tracing;
mod utils;
pub use self::{metrics::*, trace_context::*, tracing::*, utils::*};
fn meter() -> opentelemetry::metrics::Meter {
opentelemetry::global::meter_with_version(
env!("CARGO_PKG_NAME"),
Some(env!("CARGO_PKG_VERSION")),
None,
)
}

237
crates/tower/src/metrics/duration.rs Normal file
View File

@ -0,0 +1,237 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::future::Future;
use opentelemetry::{metrics::Histogram, Context, KeyValue};
use pin_project_lite::pin_project;
use tokio::time::Instant;
use tower::{Layer, Service};
use crate::{utils::FnWrapper, MetricsAttributes};
/// A [`Layer`] that records the duration of requests in milliseconds.
#[derive(Clone, Debug)]
pub struct DurationRecorderLayer<OnRequest = (), OnResponse = (), OnError = ()> {
histogram: Histogram<u64>,
on_request: OnRequest,
on_response: OnResponse,
on_error: OnError,
}
impl DurationRecorderLayer {
/// Create a new [`DurationRecorderLayer`].
#[must_use]
pub fn new(name: &'static str) -> Self {
let histogram = crate::meter().u64_histogram(name).init();
Self {
histogram,
on_request: (),
on_response: (),
on_error: (),
}
}
}
impl<OnRequest, OnResponse, OnError> DurationRecorderLayer<OnRequest, OnResponse, OnError> {
/// Set the [`MetricsAttributes`] to use on request.
#[must_use]
pub fn on_request<NewOnRequest>(
self,
on_request: NewOnRequest,
) -> DurationRecorderLayer<NewOnRequest, OnResponse, OnError> {
DurationRecorderLayer {
histogram: self.histogram,
on_request,
on_response: self.on_response,
on_error: self.on_error,
}
}
#[must_use]
pub fn on_request_fn<F, T>(
self,
on_request: F,
) -> DurationRecorderLayer<FnWrapper<F>, OnResponse, OnError>
where
F: Fn(&T) -> Vec<KeyValue>,
{
self.on_request(FnWrapper(on_request))
}
/// Set the [`MetricsAttributes`] to use on response.
#[must_use]
pub fn on_response<NewOnResponse>(
self,
on_response: NewOnResponse,
) -> DurationRecorderLayer<OnRequest, NewOnResponse, OnError> {
DurationRecorderLayer {
histogram: self.histogram,
on_request: self.on_request,
on_response,
on_error: self.on_error,
}
}
#[must_use]
pub fn on_response_fn<F, T>(
self,
on_response: F,
) -> DurationRecorderLayer<OnRequest, FnWrapper<F>, OnError>
where
F: Fn(&T) -> Vec<KeyValue>,
{
self.on_response(FnWrapper(on_response))
}
/// Set the [`MetricsAttributes`] to use on error.
#[must_use]
pub fn on_error<NewOnError>(
self,
on_error: NewOnError,
) -> DurationRecorderLayer<OnRequest, OnResponse, NewOnError> {
DurationRecorderLayer {
histogram: self.histogram,
on_request: self.on_request,
on_response: self.on_response,
on_error,
}
}
#[must_use]
pub fn on_error_fn<F, T>(
self,
on_error: F,
) -> DurationRecorderLayer<OnRequest, OnResponse, FnWrapper<F>>
where
F: Fn(&T) -> Vec<KeyValue>,
{
self.on_error(FnWrapper(on_error))
}
}
impl<S, OnRequest, OnResponse, OnError> Layer<S>
for DurationRecorderLayer<OnRequest, OnResponse, OnError>
where
OnRequest: Clone,
OnResponse: Clone,
OnError: Clone,
{
type Service = DurationRecorderService<S, OnRequest, OnResponse, OnError>;
fn layer(&self, inner: S) -> Self::Service {
DurationRecorderService {
inner,
histogram: self.histogram.clone(),
on_request: self.on_request.clone(),
on_response: self.on_response.clone(),
on_error: self.on_error.clone(),
}
}
}
/// A middleware that records the duration of requests in milliseconds.
#[derive(Clone, Debug)]
pub struct DurationRecorderService<S, OnRequest = (), OnResponse = (), OnError = ()> {
inner: S,
histogram: Histogram<u64>,
on_request: OnRequest,
on_response: OnResponse,
on_error: OnError,
}
pin_project! {
/// The future returned by the [`DurationRecorderService`].
pub struct DurationRecorderFuture<F, OnResponse = (), OnError = ()> {
#[pin]
inner: F,
start: Instant,
histogram: Histogram<u64>,
attributes_from_request: Vec<KeyValue>,
from_response: OnResponse,
from_error: OnError,
}
}
impl<F, R, E, OnResponse, OnError> Future for DurationRecorderFuture<F, OnResponse, OnError>
where
F: Future<Output = Result<R, E>>,
OnResponse: MetricsAttributes<R>,
OnError: MetricsAttributes<E>,
{
type Output = F::Output;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let this = self.project();
let result = std::task::ready!(this.inner.poll(cx));
// Measure the duration of the request.
let duration = this.start.elapsed();
let duration_ms = duration.as_millis().try_into().unwrap_or(u64::MAX);
// Collect the attributes from the request, response and error.
let mut attributes = this.attributes_from_request.clone();
match &result {
Ok(response) => {
attributes.extend(this.from_response.attributes(response));
}
Err(error) => {
attributes.extend(this.from_error.attributes(error));
}
}
this.histogram
.record(&Context::new(), duration_ms, &attributes);
std::task::Poll::Ready(result)
}
}
impl<S, R, OnRequest, OnResponse, OnError> Service<R>
for DurationRecorderService<S, OnRequest, OnResponse, OnError>
where
S: Service<R>,
OnRequest: MetricsAttributes<R>,
OnResponse: MetricsAttributes<S::Response> + Clone,
OnError: MetricsAttributes<S::Error> + Clone,
{
type Response = S::Response;
type Error = S::Error;
type Future = DurationRecorderFuture<S::Future, OnResponse, OnError>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, request: R) -> Self::Future {
let start = Instant::now();
let attributes_from_request = self.on_request.attributes(&request).collect();
let inner = self.inner.call(request);
DurationRecorderFuture {
inner,
start,
histogram: self.histogram.clone(),
attributes_from_request,
from_response: self.on_response.clone(),
from_error: self.on_error.clone(),
}
}
}

168
crates/tower/src/metrics/in_flight.rs Normal file
View File

@ -0,0 +1,168 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::future::Future;
use opentelemetry::{
metrics::{Unit, UpDownCounter},
Context, KeyValue,
};
use pin_project_lite::pin_project;
use tower::{Layer, Service};
use crate::MetricsAttributes;
/// A [`Layer`] that records the number of in-flight requests.
///
/// # Generic Parameters
///
/// * `OnRequest`: A type that can extract attributes from a request.
#[derive(Clone, Debug)]
pub struct InFlightCounterLayer<OnRequest = ()> {
counter: UpDownCounter<i64>,
on_request: OnRequest,
}
impl InFlightCounterLayer {
/// Create a new [`InFlightCounterLayer`].
#[must_use]
pub fn new(name: &'static str) -> Self {
let counter = crate::meter()
.i64_up_down_counter(name)
.with_unit(Unit::new("ms"))
.with_description("The number of in-flight requests")
.init();
Self {
counter,
on_request: (),
}
}
}
impl<F> InFlightCounterLayer<F> {
/// Set the [`MetricsAttributes`] to use.
#[must_use]
pub fn on_request<OnRequest>(self, on_request: OnRequest) -> InFlightCounterLayer<OnRequest> {
InFlightCounterLayer {
counter: self.counter,
on_request,
}
}
}
impl<S, OnRequest> Layer<S> for InFlightCounterLayer<OnRequest>
where
OnRequest: Clone,
{
type Service = InFlightCounterService<S, OnRequest>;
fn layer(&self, inner: S) -> Self::Service {
InFlightCounterService {
inner,
counter: self.counter.clone(),
on_request: self.on_request.clone(),
}
}
}
/// A middleware that records the number of in-flight requests.
///
/// # Generic Parameters
///
/// * `S`: The type of the inner service.
/// * `OnRequest`: A type that can extract attributes from a request.
#[derive(Clone, Debug)]
pub struct InFlightCounterService<S, OnRequest = ()> {
inner: S,
counter: UpDownCounter<i64>,
on_request: OnRequest,
}
/// A guard that decrements the in-flight request count when dropped.
struct InFlightGuard {
counter: UpDownCounter<i64>,
attributes: Vec<KeyValue>,
}
impl InFlightGuard {
fn new(counter: UpDownCounter<i64>, attributes: Vec<KeyValue>) -> Self {
counter.add(&Context::new(), 1, &attributes);
Self {
counter,
attributes,
}
}
}
impl Drop for InFlightGuard {
fn drop(&mut self) {
self.counter.add(&Context::new(), -1, &self.attributes);
}
}
pin_project! {
/// The future returned by [`InFlightCounterService`]
pub struct InFlightFuture<F> {
guard: InFlightGuard,
#[pin]
inner: F,
}
}
impl<F> Future for InFlightFuture<F>
where
F: Future,
{
type Output = F::Output;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
self.project().inner.poll(cx)
}
}
impl<R, S, OnRequest> Service<R> for InFlightCounterService<S, OnRequest>
where
S: Service<R>,
OnRequest: MetricsAttributes<R>,
{
type Response = S::Response;
type Error = S::Error;
type Future = InFlightFuture<S::Future>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: R) -> Self::Future {
// Extract attributes from the request.
let attributes = self.on_request.attributes(&req).collect();
// Increment the in-flight request count.
let guard = InFlightGuard::new(self.counter.clone(), attributes);
// Call the inner service, and return a future that decrements the in-flight
// when dropped.
let inner = self.inner.call(req);
InFlightFuture { guard, inner }
}
}

154
crates/tower/src/metrics/make_attributes.rs Normal file
View File

@ -0,0 +1,154 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use opentelemetry::{KeyValue, Value};
use crate::{utils::KV, FnWrapper};
/// Make metrics attributes from a type.
pub trait MetricsAttributes<T> {
type Iter<'a>: Iterator<Item = KeyValue>
where
Self: 'a,
T: 'a;
fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a>;
}
pub fn metrics_attributes_fn<T, F>(f: F) -> FnWrapper<F>
where
F: Fn(&T) -> Vec<KeyValue> + 'static,
T: 'static,
{
FnWrapper(f)
}
impl<T, F> MetricsAttributes<T> for FnWrapper<F>
where
F: Fn(&T) -> Vec<KeyValue> + 'static,
T: 'static,
{
type Iter<'a> = std::vec::IntoIter<KeyValue>;
fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
let values: Vec<KeyValue> = self.0(t);
values.into_iter()
}
}
impl<T> MetricsAttributes<T> for ()
where
T: 'static,
{
type Iter<'a> = std::iter::Empty<KeyValue>;
fn attributes(&self, _t: &T) -> Self::Iter<'_> {
std::iter::empty()
}
}
impl<V, T> MetricsAttributes<T> for Vec<V>
where
V: MetricsAttributes<T> + 'static,
T: 'static,
{
type Iter<'a> = Box<dyn Iterator<Item = KeyValue> + 'a>;
fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'_> {
Box::new(self.iter().flat_map(|v| v.attributes(t)))
}
}
impl<V, T> MetricsAttributes<T> for KV<V>
where
V: Into<Value> + Clone + 'static,
T: 'static,
{
type Iter<'a> = std::iter::Once<KeyValue>;
fn attributes(&self, _t: &T) -> Self::Iter<'_> {
std::iter::once(KeyValue::new(self.0, self.1.clone().into()))
}
}
impl<T> MetricsAttributes<T> for KeyValue
where
T: 'static,
{
type Iter<'a> = std::iter::Once<KeyValue>;
fn attributes(&self, _t: &T) -> Self::Iter<'_> {
std::iter::once(self.clone())
}
}
impl<V, T> MetricsAttributes<T> for Option<V>
where
V: MetricsAttributes<T> + 'static,
T: 'static,
{
type Iter<'a> = std::iter::Flatten<std::option::IntoIter<V::Iter<'a>>>;
fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'_> {
self.as_ref().map(|v| v.attributes(t)).into_iter().flatten()
}
}
macro_rules! chain_for {
// Sub-macro for reversing the list of types.
(@reverse ($( $reversed:ident ,)*)) => {
chain_for!(@build_chain $($reversed),*)
};
(@reverse ($($reversed:ident,)*) $head:ident $(, $tail:ident)*) => {
chain_for!(@reverse ($head, $($reversed,)*) $($tail),*)
};
// Sub-macro for building the chain of iterators.
(@build_chain $last:ident) => {
$last::Iter<'a>
};
(@build_chain $head:ident, $($tail:ident),*) => {
std::iter::Chain<chain_for!(@build_chain $($tail),*), $head::Iter<'a>>
};
($($idents:ident),+) => {
chain_for!(@reverse () $($idents),+)
};
}
macro_rules! impl_for_tuple {
($first:ident $(,$rest:ident)*) => {
impl<T, $first, $($rest,)*> MetricsAttributes<T> for ($first, $($rest,)*)
where
T: 'static,
$first: MetricsAttributes<T> + 'static,
$($rest: MetricsAttributes<T> + 'static,)*
{
type Iter<'a> = chain_for!($first $(, $rest)*);
fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
#[allow(non_snake_case)]
let (head, $($rest,)*) = self;
head.attributes(t)
$(.chain($rest.attributes(t)))*
}
}
};
}
impl_for_tuple!(V1);
impl_for_tuple!(V1, V2);
impl_for_tuple!(V1, V2, V3);
impl_for_tuple!(V1, V2, V3, V4);
impl_for_tuple!(V1, V2, V3, V4, V5);
impl_for_tuple!(V1, V2, V3, V4, V5, V6);
impl_for_tuple!(V1, V2, V3, V4, V5, V6, V7);
impl_for_tuple!(V1, V2, V3, V4, V5, V6, V7, V8);
impl_for_tuple!(V1, V2, V3, V4, V5, V6, V7, V8, V9);

23
crates/tower/src/metrics/mod.rs Normal file
View File

@ -0,0 +1,23 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
mod duration;
mod in_flight;
mod make_attributes;
pub use self::{
duration::{DurationRecorderFuture, DurationRecorderLayer, DurationRecorderService},
in_flight::{InFlightCounterLayer, InFlightCounterService, InFlightFuture},
make_attributes::{metrics_attributes_fn, MetricsAttributes},
};

115
crates/tower/src/trace_context.rs Normal file
View File

@ -0,0 +1,115 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use http::Request;
use opentelemetry::propagation::Injector;
use opentelemetry_http::HeaderInjector;
use tower::{Layer, Service};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;
/// A trait to get an [`Injector`] from a request.
trait AsInjector {
type Injector<'a>: Injector
where
Self: 'a;
fn as_injector(&mut self) -> Self::Injector<'_>;
}
impl<B> AsInjector for Request<B> {
type Injector<'a> = HeaderInjector<'a> where Self: 'a;
fn as_injector(&mut self) -> Self::Injector<'_> {
HeaderInjector(self.headers_mut())
}
}
#[cfg(feature = "aws-sdk")]
impl AsInjector for aws_smithy_http::operation::Request {
type Injector<'a> = HeaderInjector<'a> where Self: 'a;
fn as_injector(&mut self) -> Self::Injector<'_> {
HeaderInjector(self.http_mut().headers_mut())
}
}
/// A [`Layer`] that adds a trace context to the request.
#[derive(Debug, Clone, Copy, Default)]
pub struct TraceContextLayer {
_private: (),
}
impl TraceContextLayer {
/// Create a new [`TraceContextLayer`].
#[must_use]
pub fn new() -> Self {
Self::default()
}
}
impl<S> Layer<S> for TraceContextLayer {
type Service = TraceContextService<S>;
fn layer(&self, inner: S) -> Self::Service {
TraceContextService::new(inner)
}
}
/// A [`Service`] that adds a trace context to the request.
#[derive(Debug, Clone)]
pub struct TraceContextService<S> {
inner: S,
}
impl<S> TraceContextService<S> {
/// Create a new [`TraceContextService`].
pub fn new(inner: S) -> Self {
Self { inner }
}
}
impl<S, R> Service<R> for TraceContextService<S>
where
S: Service<R>,
R: AsInjector,
{
type Response = S::Response;
type Error = S::Error;
type Future = S::Future;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: R) -> Self::Future {
// Get the `opentelemetry` context out of the `tracing` span.
let context = Span::current().context();
// Inject the trace context into the request. The block is there to ensure that
// the injector is dropped before calling the inner service, to avoid borrowing
// issues.
{
let mut injector = req.as_injector();
opentelemetry::global::get_text_map_propagator(|propagator| {
propagator.inject_context(&context, &mut injector);
});
}
self.inner.call(req)
}
}

114
crates/tower/src/tracing/enrich_span.rs Normal file
View File

@ -0,0 +1,114 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use tracing::{Span, Value};
use crate::utils::{FnWrapper, KV};
/// A trait for enriching a span with information a structure.
pub trait EnrichSpan<T> {
fn enrich_span(&self, span: &Span, t: &T);
}
impl<T, F> EnrichSpan<T> for FnWrapper<F>
where
F: Fn(&Span, &T),
{
fn enrich_span(&self, span: &Span, t: &T) {
(self.0)(span, t);
}
}
/// Enrich span from a function.
#[must_use]
pub fn enrich_span_fn<T, F>(f: F) -> FnWrapper<F>
where
F: Fn(&Span, &T),
{
FnWrapper(f)
}
impl<T> EnrichSpan<T> for () {
fn enrich_span(&self, _span: &Span, _t: &T) {}
}
impl<V, T> EnrichSpan<T> for KV<V>
where
V: Value,
{
fn enrich_span(&self, span: &Span, _t: &T) {
span.record(self.0, &self.1);
}
}
/// A macro to implement [`EnrichSpan`] for a tuple of types that implement
/// [`EnrichSpan`].
macro_rules! impl_for_tuple {
($($T:ident),+) => {
impl<T, $($T),+> EnrichSpan<T> for ($($T,)+)
where
$($T: EnrichSpan<T>),+
{
fn enrich_span(&self, span: &Span, t: &T) {
#[allow(non_snake_case)]
let ($(ref $T,)+) = *self;
$(
$T.enrich_span(span, t);
)+
}
}
};
}
impl_for_tuple!(T1);
impl_for_tuple!(T1, T2);
impl_for_tuple!(T1, T2, T3);
impl_for_tuple!(T1, T2, T3, T4);
impl_for_tuple!(T1, T2, T3, T4, T5);
impl_for_tuple!(T1, T2, T3, T4, T5, T6);
impl_for_tuple!(T1, T2, T3, T4, T5, T6, T7);
impl_for_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
impl<T, R> EnrichSpan<R> for Option<T>
where
T: EnrichSpan<R>,
{
fn enrich_span(&self, span: &Span, request: &R) {
if let Some(ref t) = *self {
t.enrich_span(span, request);
}
}
}
impl<T, R, const N: usize> EnrichSpan<R> for [T; N]
where
T: EnrichSpan<R>,
{
fn enrich_span(&self, span: &Span, request: &R) {
for t in self {
t.enrich_span(span, request);
}
}
}
impl<T, R> EnrichSpan<R> for Vec<T>
where
T: EnrichSpan<R>,
{
fn enrich_span(&self, span: &Span, request: &R) {
for t in self {
t.enrich_span(span, request);
}
}
}

70
crates/tower/src/tracing/future.rs Normal file
View File

@ -0,0 +1,70 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::{future::Future, task::ready};
use pin_project_lite::pin_project;
use tracing::Span;
pin_project! {
pub struct TraceFuture<F, OnResponse, OnError> {
#[pin]
inner: F,
span: Span,
on_response: OnResponse,
on_error: OnError,
}
}
impl<F, OnResponse, OnError> TraceFuture<F, OnResponse, OnError> {
pub fn new(inner: F, span: Span, on_response: OnResponse, on_error: OnError) -> Self {
Self {
inner,
span,
on_response,
on_error,
}
}
}
impl<F, R, E, OnResponse, OnError> Future for TraceFuture<F, OnResponse, OnError>
where
F: Future<Output = Result<R, E>>,
OnResponse: super::enrich_span::EnrichSpan<R>,
OnError: super::enrich_span::EnrichSpan<E>,
{
type Output = Result<R, E>;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let this = self.project();
// Poll the inner future, with the span entered. This is effectively what
// [`tracing::Instrumented`] does.
let result = ready!(this.span.in_scope(|| this.inner.poll(cx)));
match &result {
Ok(response) => {
this.on_response.enrich_span(this.span, response);
}
Err(error) => {
this.on_error.enrich_span(this.span, error);
}
}
std::task::Poll::Ready(result)
}
}

105
crates/tower/src/tracing/layer.rs Normal file
View File

@ -0,0 +1,105 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use tower::Layer;
use tracing::Span;
use crate::{enrich_span_fn, make_span_fn, utils::FnWrapper};
#[derive(Clone, Debug)]
pub struct TraceLayer<MakeSpan, OnResponse = (), OnError = ()> {
make_span: MakeSpan,
on_response: OnResponse,
on_error: OnError,
}
impl<F> TraceLayer<FnWrapper<F>> {
#[must_use]
pub fn from_fn<T>(f: F) -> TraceLayer<FnWrapper<F>>
where
F: Fn(&T) -> Span,
{
TraceLayer::new(make_span_fn(f))
}
}
impl<MakeSpan> TraceLayer<MakeSpan> {
#[must_use]
pub fn new(make_span: MakeSpan) -> Self {
Self {
make_span,
on_response: (),
on_error: (),
}
}
}
impl<MakeSpan, OnResponse, OnError> TraceLayer<MakeSpan, OnResponse, OnError> {
#[must_use]
pub fn on_response<NewOnResponse>(
self,
on_response: NewOnResponse,
) -> TraceLayer<MakeSpan, NewOnResponse, OnError> {
TraceLayer {
make_span: self.make_span,
on_response,
on_error: self.on_error,
}
}
#[must_use]
pub fn on_response_fn<F, T>(self, f: F) -> TraceLayer<MakeSpan, FnWrapper<F>, OnError>
where
F: Fn(&Span, &T),
{
self.on_response(enrich_span_fn(f))
}
#[must_use]
pub fn on_error<NewOnError>(
self,
on_error: NewOnError,
) -> TraceLayer<MakeSpan, OnResponse, NewOnError> {
TraceLayer {
make_span: self.make_span,
on_response: self.on_response,
on_error,
}
}
pub fn on_error_fn<F, E>(self, f: F) -> TraceLayer<MakeSpan, OnResponse, FnWrapper<F>>
where
F: Fn(&Span, &E),
{
self.on_error(enrich_span_fn(f))
}
}
impl<S, MakeSpan, OnResponse, OnError> Layer<S> for TraceLayer<MakeSpan, OnResponse, OnError>
where
MakeSpan: Clone,
OnResponse: Clone,
OnError: Clone,
{
type Service = super::service::TraceService<S, MakeSpan, OnResponse, OnError>;
fn layer(&self, inner: S) -> Self::Service {
super::service::TraceService::new(
inner,
self.make_span.clone(),
self.on_response.clone(),
self.on_error.clone(),
)
}
}

72
crates/tower/src/tracing/make_span.rs Normal file
View File

@ -0,0 +1,72 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use tracing::Span;
use super::enrich_span::EnrichSpan;
use crate::utils::FnWrapper;
/// A trait for creating a span for a request.
pub trait MakeSpan<R> {
fn make_span(&self, request: &R) -> Span;
}
impl<R, F> MakeSpan<R> for FnWrapper<F>
where
F: Fn(&R) -> Span,
{
fn make_span(&self, request: &R) -> Span {
(self.0)(request)
}
}
/// Make span from a function.
pub fn make_span_fn<R, F>(f: F) -> FnWrapper<F>
where
F: Fn(&R) -> Span,
{
FnWrapper(f)
}
/// A macro to implement [`MakeSpan`] for a tuple of types, where the first type
/// implements [`MakeSpan`] and the rest implement [`EnrichSpan`].
macro_rules! impl_for_tuple {
(M, $($T:ident),+) => {
impl<R, M, $($T),+> MakeSpan<R> for (M, $($T),+)
where
M: MakeSpan<R>,
$($T: EnrichSpan<R>),+
{
fn make_span(&self, request: &R) -> Span {
#[allow(non_snake_case)]
let (ref m, $(ref $T),+) = *self;
let span = m.make_span(request);
$(
$T.enrich_span(&span, request);
)+
span
}
}
};
}
impl_for_tuple!(M, T1);
impl_for_tuple!(M, T1, T2);
impl_for_tuple!(M, T1, T2, T3);
impl_for_tuple!(M, T1, T2, T3, T4);
impl_for_tuple!(M, T1, T2, T3, T4, T5);
impl_for_tuple!(M, T1, T2, T3, T4, T5, T6);
impl_for_tuple!(M, T1, T2, T3, T4, T5, T6, T7);
impl_for_tuple!(M, T1, T2, T3, T4, T5, T6, T7, T8);

27
crates/tower/src/tracing/mod.rs Normal file
View File

@ -0,0 +1,27 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
mod enrich_span;
mod future;
mod layer;
mod make_span;
mod service;
pub use self::{
enrich_span::{enrich_span_fn, EnrichSpan},
future::TraceFuture,
layer::TraceLayer,
make_span::{make_span_fn, MakeSpan},
service::TraceService,
};

67
crates/tower/src/tracing/service.rs Normal file
View File

@ -0,0 +1,67 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use tower::Service;
use super::future::TraceFuture;
#[derive(Clone, Debug)]
pub struct TraceService<S, MakeSpan, OnResponse = (), OnError = ()> {
inner: S,
make_span: MakeSpan,
on_response: OnResponse,
on_error: OnError,
}
impl<S, MakeSpan, OnResponse, OnError> TraceService<S, MakeSpan, OnResponse, OnError> {
/// Create a new [`TraceService`].
#[must_use]
pub fn new(inner: S, make_span: MakeSpan, on_response: OnResponse, on_error: OnError) -> Self {
Self {
inner,
make_span,
on_response,
on_error,
}
}
}
impl<R, S, MakeSpan, OnResponse, OnError> Service<R>
for TraceService<S, MakeSpan, OnResponse, OnError>
where
S: Service<R>,
MakeSpan: super::make_span::MakeSpan<R>,
OnResponse: super::enrich_span::EnrichSpan<S::Response> + Clone,
OnError: super::enrich_span::EnrichSpan<S::Error> + Clone,
{
type Response = S::Response;
type Error = S::Error;
type Future = TraceFuture<S::Future, OnResponse, OnError>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, request: R) -> Self::Future {
let span = self.make_span.make_span(&request);
let guard = span.enter();
let inner = self.inner.call(request);
drop(guard);
TraceFuture::new(inner, span, self.on_response.clone(), self.on_error.clone())
}
}

33
crates/tower/src/utils.rs Normal file
View File

@ -0,0 +1,33 @@
// Copyright 2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use opentelemetry::{KeyValue, Value};
/// A simple static key-value pair.
#[derive(Clone, Debug)]
pub struct KV<V>(pub &'static str, pub V);
impl<V> From<KV<V>> for KeyValue
where
V: Into<Value>,
{
fn from(value: KV<V>) -> Self {
Self::new(value.0, value.1.into())
}
}
/// A wrapper around a function that can be used to generate a key-value pair,
/// make or enrich spans.
#[derive(Clone, Debug)]
pub struct FnWrapper<F>(pub F);