Is there any way to update the name of the outer transaction span at the end of a request? about tracing-elastic-apm HOT 4 CLOSED

pub struct ApmAugmentLayer {
    inner_layer: ApmLayer,
}

impl ApmAugmentLayer {
    pub fn new(layer: ApmLayer) -> Self {
        ApmAugmentLayer { inner_layer: layer }
    }
}

struct MatchStrVisitor {
    field: String,
    value: Option<String>,
}

impl Visit for MatchStrVisitor {
    fn record_debug(&mut self, _field: &Field, _value: &dyn Debug) {}
    fn record_str(&mut self, field: &Field, value: &str) {
        if field.name() == self.field {
            self.value = Some(value.to_string());
        }
    }
}

struct NameOverrideValue(String);

impl<S> Layer<S> for ApmAugmentLayer
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
        self.inner_layer.on_new_span(attrs, id, ctx);
    }

    fn on_record(&self, span: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
        {
            let spanref = ctx.span(span).expect("Span not found!");
            let mut extensions = spanref.extensions_mut();
            let mut visitor = MatchStrVisitor {
                field: "name_override".to_string(),
                value: None,
            };
            values.record(&mut visitor);
            if let Some(name_override) = visitor.value {
                extensions.insert(NameOverrideValue(name_override.to_string()));
            };
        };
        self.inner_layer.on_record(span, values, ctx);
    }

    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
        self.inner_layer.on_event(event, ctx);
    }

    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
        self.inner_layer.on_enter(id, ctx);
    }

    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
        self.inner_layer.on_exit(id, ctx);
    }

    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
        {
            let span = ctx.span(&id).expect("Span not found!");
            let mut extensions = span.extensions_mut();
            if let Some(name_override) = extensions.remove::<NameOverrideValue>() {
                if let Some(transaction) = extensions.get_mut::<Transaction>() {
                    transaction.name = Some(name_override.0)
                }
            }
        };
        self.inner_layer.on_close(id, ctx);
    }
}

pub async fn wrap_request(
    req: Request<axum::body::Body>,
    next: Next,
) -> Result<Response<axum::body::Body>, Infallible> {
    let request_id = req
        .headers()
        .get(X_PURE_REQUESTID)
        .and_then(|value| value.to_str().ok())
        .map(ToString::to_string);

    let response = match request_id {
        Some(request_id) => {
            let span = info_span!("request", %request_id, name_override = field::Empty);
            let _enter = span.enter();
            let (path, method) = {
                let path = if let Some(path) = req.extensions().get::<MatchedPath>() {
                    path.as_str()
                } else {
                    req.uri().path()
                };
                let method = req.method().clone();
                (path.to_owned(), method)
            };
            let res = next.run(req).in_current_span().await;
            let summ_status = match res.status().as_u16() {
                200..=299 => "2xx",
                300..=399 => "3xx",
                400..=499 => "4xx",
                500..=599 => "5xx",
                _ => "Unk",
            };

            span.record(
                "name_override",
                format!("HTTP {} {} {}", summ_status, method, path),
            );

            res
        }
        None => next.run(req).await,
    };

    Ok(response)
}

    let apm_layer = tracing_elastic_apm::new_layer(
        service_name.to_owned(),
        tracing_elastic_apm::config::Config::new(url.clone()).with_authorization(
            tracing_elastic_apm::config::Authorization::ApiKey(
                tracing_elastic_apm::config::ApiKey::new(id, key),
            ),
        ),
    )
    .expect("Bad APM Config");

    tracing_subscriber::registry()
        .with(ApmAugmentLayer::new(apm_layer)