From af56985d70ce20c58a9d1e16a343e870afe3c2b1 Mon Sep 17 00:00:00 2001
From: Dom Dwyer <dom@itsallbroken.com>
Date: Wed, 12 Jul 2023 14:42:43 +0200
Subject: [PATCH] refactor(ingester): emit span for query handler

Emit a span that covers the entire flight query handler.
---
 ingester/src/server/grpc/query.rs | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/ingester/src/server/grpc/query.rs b/ingester/src/server/grpc/query.rs
index 763ad89688..b4dde8853c 100644
--- a/ingester/src/server/grpc/query.rs
+++ b/ingester/src/server/grpc/query.rs
@@ -19,7 +19,7 @@ use tokio::sync::{Semaphore, TryAcquireError};
 use tonic::{Request, Response, Streaming};
 use trace::{
     ctx::SpanContext,
-    span::{Span, SpanExt},
+    span::{Span, SpanExt, SpanRecorder},
 };
 
 mod instrumentation;
@@ -175,7 +175,7 @@ where
         request: Request<Ticket>,
     ) -> Result<Response<Self::DoGetStream>, tonic::Status> {
         let span_ctx: Option<SpanContext> = request.extensions().get().cloned();
-        let span = span_ctx.child_span("ingester query");
+        let mut query_recorder = SpanRecorder::new(span_ctx.child_span("ingester query"));
 
         // Acquire and hold a permit for the duration of this request, or return
         // an error if the existing requests have already exhausted the
@@ -211,16 +211,23 @@ where
 
         let response = match self
             .query_handler
-            .query_exec(namespace_id, table_id, projection, span.clone(), predicate)
+            .query_exec(
+                namespace_id,
+                table_id,
+                projection,
+                query_recorder.child_span("query exec"),
+                predicate,
+            )
             .await
         {
             Ok(v) => v,
             Err(e @ (QueryError::TableNotFound(_, _) | QueryError::NamespaceNotFound(_))) => {
                 debug!(
-                        error=%e,
-                        %namespace_id,
-                        %table_id,
-                        "query error, no buffered data found");
+                    error=%e,
+                    %namespace_id,
+                    %table_id,
+                    "no buffered data found for query"
+                );
 
                 return Err(e)?;
             }
@@ -229,11 +236,12 @@ where
         let output = encode_response(
             response,
             self.ingester_id,
-            span,
+            query_recorder.child_span("serialise response"),
             Arc::clone(&self.query_request_frame_encoding_duration),
         )
         .map_err(tonic::Status::from);
 
+        query_recorder.ok("query exec complete - streaming results");
         Ok(Response::new(Box::pin(output) as Self::DoGetStream))
     }