mas_context/
service.rs

1// Copyright 2025 New Vector Ltd.
2//
3// SPDX-License-Identifier: AGPL-3.0-only
4// Please see LICENSE in the repository root for full details.
5
6use std::{
7    borrow::Cow,
8    task::{Context, Poll},
9};
10
11use tower_service::Service;
12
13use crate::{LogContext, LogContextFuture};
14
15/// A service which wraps another service and creates a log context for
16/// each request.
17pub struct LogContextService<S, R> {
18    inner: S,
19    tagger: fn(&R) -> Cow<'static, str>,
20}
21
22impl<S: Clone, R> Clone for LogContextService<S, R> {
23    fn clone(&self) -> Self {
24        Self {
25            inner: self.inner.clone(),
26            tagger: self.tagger,
27        }
28    }
29}
30
31impl<S, R> LogContextService<S, R> {
32    pub fn new(inner: S, tagger: fn(&R) -> Cow<'static, str>) -> Self {
33        Self { inner, tagger }
34    }
35}
36
37impl<S, R> Service<R> for LogContextService<S, R>
38where
39    S: Service<R>,
40{
41    type Response = S::Response;
42    type Error = S::Error;
43    type Future = LogContextFuture<S::Future>;
44
45    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
46        self.inner.poll_ready(cx)
47    }
48
49    fn call(&mut self, req: R) -> Self::Future {
50        let tag = (self.tagger)(&req);
51        let log_context = LogContext::new(tag);
52        log_context.run(|| self.inner.call(req))
53    }
54}