Replace the OTEL-based tracing layer with tracing based layers

crates/tower/Cargo.toml

@@ -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"]

crates/tower/src/lib.rs

@@ -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,
+    )
+}

crates/tower/src/metrics/duration.rs

@@ -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(),
+        }
+    }
+}

crates/tower/src/metrics/in_flight.rs

@@ -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 }
+    }
+}

crates/tower/src/metrics/make_attributes.rs

@@ -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);

crates/tower/src/metrics/mod.rs

@@ -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},
+};

crates/tower/src/trace_context.rs

@@ -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)
+    }
+}

crates/tower/src/tracing/enrich_span.rs

@@ -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);
+        }
+    }
+}

crates/tower/src/tracing/future.rs

@@ -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)
+    }
+}

crates/tower/src/tracing/layer.rs

@@ -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(),
+        )
+    }
+}

crates/tower/src/tracing/make_span.rs

@@ -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);

crates/tower/src/tracing/mod.rs

@@ -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,
+};

crates/tower/src/tracing/service.rs

@@ -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())
+    }
+}

crates/tower/src/utils.rs

@@ -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);