1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use crate::extensions::{Extension, ResolveInfo};
use crate::Variables;
use async_graphql_parser::types::ExecutableDocument;
use std::collections::BTreeMap;
use tracing::{span, Level, Span};

/// Tracing extension
///
/// # References
///
/// <https://crates.io/crates/tracing>
#[derive(Default)]
#[cfg_attr(feature = "nightly", doc(cfg(feature = "tracing")))]
pub struct Tracing {
    root: Option<Span>,
    parse: Option<Span>,
    validation: Option<Span>,
    execute: Option<Span>,
    fields: BTreeMap<usize, Span>,
}

impl Extension for Tracing {
    #[allow(clippy::deref_addrof)]
    fn parse_start(&mut self, query_source: &str, _variables: &Variables) {
        let root_span = span!(
            target: "async_graphql::graphql",
            parent: None,
            Level::INFO,
            "query",
            source = %query_source
        );

        let parse_span = span!(
            target: "async_graphql::graphql",
            parent: &root_span,
            Level::INFO,
            "parse"
        );

        root_span.with_subscriber(|(id, d)| d.enter(id));
        self.root.replace(root_span);

        parse_span.with_subscriber(|(id, d)| d.enter(id));
        self.parse.replace(parse_span);
    }

    fn parse_end(&mut self, _document: &ExecutableDocument) {
        self.parse
            .take()
            .unwrap()
            .with_subscriber(|(id, d)| d.exit(id));
    }

    fn validation_start(&mut self) {
        let parent = self.root.as_ref().unwrap();
        let validation_span = span!(
            target: "async_graphql::graphql",
            parent: parent,
            Level::INFO,
            "validation"
        );
        validation_span.with_subscriber(|(id, d)| d.enter(id));
        self.validation.replace(validation_span);
    }

    fn validation_end(&mut self) {
        self.validation
            .take()
            .unwrap()
            .with_subscriber(|(id, d)| d.exit(id));
    }

    fn execution_start(&mut self) {
        let parent = self.root.as_ref().unwrap();
        let execute_span = span!(
            target: "async_graphql::graphql",
            parent: parent,
            Level::INFO,
            "execute"
        );
        execute_span.with_subscriber(|(id, d)| d.enter(id));
        self.execute.replace(execute_span);
    }

    fn execution_end(&mut self) {
        self.execute
            .take()
            .unwrap()
            .with_subscriber(|(id, d)| d.exit(id));

        self.root
            .take()
            .unwrap()
            .with_subscriber(|(id, d)| d.exit(id));
    }

    fn resolve_start(&mut self, info: &ResolveInfo<'_>) {
        let parent_span = match info.resolve_id.parent {
            Some(parent_id) if parent_id > 0 => self.fields.get(&parent_id).unwrap(),
            _ => self.execute.as_ref().unwrap(),
        };

        let span = span!(
            target: "async_graphql::graphql",
            parent: parent_span,
            Level::INFO,
            "field",
            id = %info.resolve_id.current,
            path = %info.path_node
        );
        span.with_subscriber(|(id, d)| d.enter(id));
        self.fields.insert(info.resolve_id.current, span);
    }

    fn resolve_end(&mut self, info: &ResolveInfo<'_>) {
        if let Some(span) = self.fields.remove(&info.resolve_id.current) {
            span.with_subscriber(|(id, d)| d.exit(id));
        }
    }
}