Implement optelemetry tracing for FVR (#2979)

* Implement optelemetry tracing for FVR * implement eran's suggestions
mobilecoinfoundation · Jan 10, 2023 · fd52c2d · fd52c2d
1 parent 199bcac
commit fd52c2d
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 14 deletions.
diff --git a/fog/view/server/src/fog_view_router_service.rs b/fog/view/server/src/fog_view_router_service.rs
@@ -13,6 +13,7 @@ use mc_fog_uri::FogViewStoreUri;
 use mc_fog_view_enclave_api::ViewEnclaveProxy;
 use mc_util_grpc::{check_request_chain_id, rpc_logger, send_result, Authenticator};
 use mc_util_metrics::{ServiceMetrics, SVC_COUNTERS};
+use mc_util_telemetry::tracer;
 use std::{
     collections::HashMap,
     sync::{Arc, RwLock},
@@ -134,11 +135,13 @@ where
 
             // This will block the async API. We should use some sort of differentiator...
             let shard_clients = self.shard_clients.read().expect("RwLock poisoned");
+            let tracer = tracer!();
             let result = block_on(router_request_handler::handle_query_request(
                 request,
                 self.enclave.clone(),
                 shard_clients.values().cloned().collect(),
                 self.logger.clone(),
+                &tracer,
             ))
             .map(|mut response| response.take_query());
 

diff --git a/fog/view/server/src/router_request_handler.rs b/fog/view/server/src/router_request_handler.rs
@@ -18,6 +18,7 @@ use mc_fog_uri::FogViewStoreUri;
 use mc_fog_view_enclave_api::ViewEnclaveProxy;
 use mc_util_grpc::{rpc_invalid_arg_error, ConnectionUriGrpcioChannel, ResponseStatus};
 use mc_util_metrics::{GrpcMethodName, SVC_COUNTERS};
+use mc_util_telemetry::{create_context, tracer, BoxedTracer, FutureExt, Tracer};
 use mc_util_uri::ConnectionUri;
 use std::sync::Arc;
 
@@ -71,10 +72,21 @@ pub async fn handle_request<E>(
 where
     E: ViewEnclaveProxy,
 {
+    let tracer = tracer!();
     if request.has_auth() {
-        handle_auth_request(enclave, request.take_auth(), logger)
+        tracer.in_span("router_auth", |_cx| {
+            handle_auth_request(enclave, request.take_auth(), logger)
+        })
     } else if request.has_query() {
-        handle_query_request(request.take_query(), enclave, shard_clients, logger).await
+        handle_query_request(
+            request.take_query(),
+            enclave,
+            shard_clients,
+            logger,
+            &tracer,
+        )
+        .with_context(create_context(&tracer, "router_query"))
+        .await
     } else {
         let rpc_status = rpc_invalid_arg_error(
             "Inavlid FogViewRouterRequest request",
@@ -109,6 +121,7 @@ pub async fn handle_query_request<E>(
     enclave: E,
     shard_clients: Vec<Arc<FogViewStoreApiClient>>,
     logger: Logger,
+    tracer: &BoxedTracer,
 ) -> Result<FogViewRouterResponse, RpcStatus>
 where
     E: ViewEnclaveProxy,
@@ -129,17 +142,20 @@ where
         shard_clients.clone(),
         logger.clone(),
     )
+    .with_context(create_context(tracer, "router_get_query_responses"))
     .await?;
 
-    let query_response = enclave
-        .collate_shard_query_responses(sealed_query, query_responses)
-        .map_err(|err| {
-            router_server_err_to_rpc_status(
-                "Query: shard response collation",
-                RouterServerError::Enclave(err),
-                logger.clone(),
-            )
-        })?;
+    let query_response = tracer.in_span("router_collate_query_responses", |_cx| {
+        enclave
+            .collate_shard_query_responses(sealed_query, query_responses)
+            .map_err(|err| {
+                router_server_err_to_rpc_status(
+                    "Query: shard response collation",
+                    RouterServerError::Enclave(err),
+                    logger.clone(),
+                )
+            })
+    })?;
 
     let mut response = FogViewRouterResponse::new();
     response.set_query(query_response.into());

diff --git a/util/telemetry/src/lib.rs b/util/telemetry/src/lib.rs
@@ -3,7 +3,8 @@
 //! OpenTelemetry wrappers and helper utilities.
 
 pub use opentelemetry::{
-    trace::{mark_span_as_active, Span, SpanKind, TraceContextExt, Tracer},
+    global::BoxedTracer,
+    trace::{mark_span_as_active, FutureExt, Span, SpanKind, TraceContextExt, Tracer},
     Context, Key,
 };
 
@@ -13,8 +14,6 @@ use opentelemetry::{
 };
 use std::borrow::Cow;
 
-pub use opentelemetry::global::BoxedTracer;
-
 #[macro_export]
 macro_rules! tracer {
     () => {
@@ -52,6 +51,16 @@ pub fn versioned_tracer(
     tracer_provider().versioned_tracer(name, version, schema_url)
 }
 
+/// Creates a context when an explicit context is required. Useful when tracing
+/// inside an `async` method.
+pub fn create_context<T>(tracer: &BoxedTracer, name: T) -> Context
+where
+    T: Into<Cow<'static, str>>,
+{
+    let span = tracer.start(name);
+    Context::current_with_span(span)
+}
+
 /// A utility method to create a predictable trace ID out of a block index.
 /// This is used to group traces by block index.
 pub fn block_index_to_trace_id(block_index: u64) -> TraceId {