rama-http 0.3.0-rc1

use crate::Request;
use crate::header::USER_AGENT;
use crate::opentelemetry::version_as_protocol_version;
use rama_core::telemetry::tracing::{self, Level, Span};
use rama_net::Protocol;
use rama_utils::str::arcstr::{ArcStr, arcstr};

use super::DEFAULT_MESSAGE_LEVEL;

/// Trait used to generate [`Span`]s from requests. [`Trace`] wraps all request handling in this
/// span.
///
/// [`Span`]: tracing::Span
/// [`Trace`]: super::Trace
pub trait MakeSpan<B>: Send + Sync + 'static {
    /// Make a span from a request.
    fn make_span(&self, request: &Request<B>) -> Span;
}

impl<B> MakeSpan<B> for Span {
    fn make_span(&self, _request: &Request<B>) -> Span {
        self.clone()
    }
}

impl<F, B> MakeSpan<B> for F
where
    F: Fn(&Request<B>) -> Span + Send + Sync + 'static,
{
    fn make_span(&self, request: &Request<B>) -> Span {
        self(request)
    }
}

/// The default way [`Span`]s will be created for [`Trace`].
///
/// [`Span`]: tracing::Span
/// [`Trace`]: super::Trace
#[derive(Debug, Clone)]
pub struct DefaultMakeSpan {
    level: Level,
    include_headers: bool,
    otel_name: ArcStr,
}

impl DefaultMakeSpan {
    /// Create a new `DefaultMakeSpan`.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            level: DEFAULT_MESSAGE_LEVEL,
            include_headers: false,
            otel_name: arcstr!("request"),
        }
    }

    rama_utils::macros::generate_set_and_with! {
        /// Set the [`Level`] used for the [tracing span].
        ///
        /// Defaults to [`Level::DEBUG`].
        ///
        /// [tracing span]: https://docs.rs/tracing/latest/tracing/#spans
        pub fn level(mut self, level: Level) -> Self {
            self.level = level;
            self
        }
    }

    rama_utils::macros::generate_set_and_with! {
        /// Include request headers on the [`Span`].
        ///
        /// By default headers are not included.

        pub fn include_headers(mut self, include_headers: bool) -> Self {
            self.include_headers = include_headers;
            self
        }
    }

    rama_utils::macros::generate_set_and_with! {
        /// Set the otel.name field of this [`Span`].
        pub fn name(mut self, otel_name: ArcStr) -> Self {
            self.otel_name = otel_name;
            self
        }
    }
}

impl Default for DefaultMakeSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl<B> MakeSpan<B> for DefaultMakeSpan {
    fn make_span(&self, request: &Request<B>) -> Span {
        let full_uri = request.request_uri();

        // to ensure that we always log authority even if not included in full protocol
        let protocol = full_uri.scheme().unwrap_or(&Protocol::HTTP);

        // according to OTEL spec domain can be domain or IP, so host is fine
        let (found_domain, found_port, found_scheme) = if let Some(host) = full_uri.host() {
            (
                Some(host),
                full_uri.port_u16().or_else(|| protocol.default_port()),
                Some(protocol.as_str()),
            )
        } else {
            tracing::debug!("no authority could be resolved for request");
            (None, None, None)
        };

        let found_domain_cow_str = found_domain.as_ref().map(|d| d.to_str());
        let found_domain_str = found_domain_cow_str.as_deref();

        let url_path = request.uri().path_or_root();
        let url_query = request.uri().query_or_empty();

        // This ugly macro is needed, unfortunately, because `tracing::span!`
        // required the level argument to be static. Meaning we can't just pass
        // `self.level`.
        macro_rules! make_span {
            ($level:expr) => {
                if self.include_headers {
                    tracing::span!(
                        $level,
                        "request",
                        otel.name = self.otel_name.as_str(),
                        http.request.method = %request.method(),
                        url.full = %request.request_uri(),
                        url.domain = found_domain_str,
                        url.port = found_port,
                        url.path = %url_path.as_ref(),
                        url.query = %url_query.as_ref(),
                        url.scheme = found_scheme,
                        network.protocol.name = "http",
                        network.protocol.version = version_as_protocol_version(request.version()),
                        user_agent.original = %request.headers().get(USER_AGENT).and_then(|v| v.to_str().ok()).unwrap_or_default(),
                        headers = ?request.headers(),
                    )
                } else {
                    tracing::span!(
                        $level,
                        "request",
                        otel.name = self.otel_name.as_str(),
                        http.request.method = %request.method(),
                        url.full = %request.request_uri(),
                        url.domain = found_domain_str,
                        url.port = found_port,
                        url.path = %url_path.as_ref(),
                        url.query = %url_query.as_ref(),
                        url.scheme = found_scheme,
                        network.protocol.name = "http",
                        network.protocol.version = version_as_protocol_version(request.version()),
                        user_agent.original = %request.headers().get(USER_AGENT).and_then(|v| v.to_str().ok()).unwrap_or_default(),
                    )
                }
            }
        }

        match self.level {
            Level::ERROR => make_span!(Level::ERROR),
            Level::WARN => make_span!(Level::WARN),
            Level::INFO => make_span!(Level::INFO),
            Level::DEBUG => make_span!(Level::DEBUG),
            Level::TRACE => make_span!(Level::TRACE),
        }
    }
}