hive-router 0.2.11

GraphQL router for Federation, part of the Hive platform
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

use std::future::Future;
use std::sync::Arc;

use crate::telemetry::logging::request_id::{RequestIdentifiers, REQUEST_IDENTIFIERS};
use crate::telemetry::{
    logging::summary::{RequestSummary, REQUEST_SUMMARY},
    traces::hive_trace_context::{HiveTraceDocuments, HiveTraceFuture, HIVE_TRACE_DOCUMENTS},
};

/// Captures per-request task-locals for work that outlives the request future, such as subscriptions.
pub(crate) struct RequestTaskScope {
    identifiers: Option<Arc<RequestIdentifiers>>,
    summary: Option<Arc<RequestSummary>>,
    hive_trace_documents: Option<Arc<HiveTraceDocuments>>,
}

impl RequestTaskScope {
    pub fn capture() -> Self {
        Self {
            identifiers: REQUEST_IDENTIFIERS.try_with(Arc::clone).ok(),
            summary: REQUEST_SUMMARY.try_with(Arc::clone).ok(),
            hive_trace_documents: HIVE_TRACE_DOCUMENTS.try_with(Arc::clone).ok(),
        }
    }

    pub async fn scope<F: Future>(self, fut: F) -> F::Output {
        let Self {
            identifiers,
            summary,
            hive_trace_documents,
        } = self;
        let fut = async move {
            match (identifiers, summary) {
                (Some(identifiers), Some(summary)) => {
                    REQUEST_IDENTIFIERS
                        .scope(identifiers, REQUEST_SUMMARY.scope(summary, fut))
                        .await
                }
                (Some(identifiers), None) => REQUEST_IDENTIFIERS.scope(identifiers, fut).await,
                (None, Some(summary)) => REQUEST_SUMMARY.scope(summary, fut).await,
                (None, None) => fut.await,
            }
        };
        match hive_trace_documents {
            Some(documents) => HiveTraceFuture::new(fut, documents).await,
            None => fut.await,
        }
    }
}