async_graphql/extensions/
apollo_tracing.rs

1use std::sync::Arc;
2
3use chrono::{DateTime, Utc};
4use futures_util::lock::Mutex;
5use serde::{Serialize, Serializer, ser::SerializeMap};
6
7use crate::{
8    Response, ServerResult, Value,
9    extensions::{
10        Extension, ExtensionContext, ExtensionFactory, NextExecute, NextResolve, ResolveInfo,
11    },
12    value,
13};
14
15struct ResolveState {
16    path: Vec<String>,
17    field_name: String,
18    parent_type: String,
19    return_type: String,
20    start_time: DateTime<Utc>,
21    end_time: DateTime<Utc>,
22    start_offset: i64,
23}
24
25impl Serialize for ResolveState {
26    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
27        let mut map = serializer.serialize_map(None)?;
28        map.serialize_entry("path", &self.path)?;
29        map.serialize_entry("fieldName", &self.field_name)?;
30        map.serialize_entry("parentType", &self.parent_type)?;
31        map.serialize_entry("returnType", &self.return_type)?;
32        map.serialize_entry("startOffset", &self.start_offset)?;
33        map.serialize_entry(
34            "duration",
35            &(self.end_time - self.start_time).num_nanoseconds(),
36        )?;
37        map.end()
38    }
39}
40
41/// Apollo tracing extension for performance tracing
42///
43/// Apollo Tracing works by including data in the extensions field of the
44/// GraphQL response, which is reserved by the GraphQL spec for extra
45/// information that a server wants to return. That way, you have access to
46/// performance traces alongside the data returned by your query. It's already
47/// supported by `Apollo Engine`, and we're excited to see what other kinds of
48/// integrations people can build on top of this format.
49#[cfg_attr(docsrs, doc(cfg(feature = "apollo_tracing")))]
50pub struct ApolloTracing;
51
52impl ExtensionFactory for ApolloTracing {
53    fn create(&self) -> Arc<dyn Extension> {
54        Arc::new(ApolloTracingExtension {
55            inner: Mutex::new(Inner {
56                start_time: Utc::now(),
57                end_time: Utc::now(),
58                resolves: Default::default(),
59            }),
60        })
61    }
62}
63
64struct Inner {
65    start_time: DateTime<Utc>,
66    end_time: DateTime<Utc>,
67    resolves: Vec<ResolveState>,
68}
69
70struct ApolloTracingExtension {
71    inner: Mutex<Inner>,
72}
73
74#[async_trait::async_trait]
75impl Extension for ApolloTracingExtension {
76    async fn execute(
77        &self,
78        ctx: &ExtensionContext<'_>,
79        operation_name: Option<&str>,
80        next: NextExecute<'_>,
81    ) -> Response {
82        self.inner.lock().await.start_time = Utc::now();
83        let resp = next.run(ctx, operation_name).await;
84
85        let mut inner = self.inner.lock().await;
86        inner.end_time = Utc::now();
87        inner
88            .resolves
89            .sort_by(|a, b| a.start_offset.cmp(&b.start_offset));
90        resp.extension(
91            "tracing",
92            value!({
93                "version": 1,
94                "startTime": inner.start_time.to_rfc3339(),
95                "endTime": inner.end_time.to_rfc3339(),
96                "duration": (inner.end_time - inner.start_time).num_nanoseconds(),
97                "execution": {
98                    "resolvers": inner.resolves
99                }
100            }),
101        )
102    }
103
104    async fn resolve(
105        &self,
106        ctx: &ExtensionContext<'_>,
107        info: ResolveInfo<'_>,
108        next: NextResolve<'_>,
109    ) -> ServerResult<Option<Value>> {
110        let path = info.path_node.to_string_vec();
111        let field_name = info.path_node.field_name().to_string();
112        let parent_type = info.parent_type.to_string();
113        let return_type = info.return_type.to_string();
114        let start_time = Utc::now();
115        let start_offset = (start_time - self.inner.lock().await.start_time)
116            .num_nanoseconds()
117            .unwrap();
118
119        let res = next.run(ctx, info).await;
120        let end_time = Utc::now();
121
122        self.inner.lock().await.resolves.push(ResolveState {
123            path,
124            field_name,
125            parent_type,
126            return_type,
127            start_time,
128            end_time,
129            start_offset,
130        });
131        res
132    }
133}
async_graphql/extensions/apollo_tracing.rs

async_graphql/extensions/
apollo_tracing.rs
