async_graphql/extensions/
apollo_tracing.rs

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

async_graphql/extensions/
apollo_tracing.rs
