fix(idpe-17943): spans in futures stream need parent (#8364)

ingester/src/server/grpc/query.rs

@@ -359,6 +359,8 @@ fn encode_response(
     span: Option<Span>,
     frame_encoding_duration_metric: Arc<DurationHistogram>,
 ) -> impl Stream<Item = Result<FlightData, FlightError>> {
+    let span = SpanRecorder::new(span.clone()).span().cloned();
+
     response.into_partition_stream().flat_map(move |partition| {
         let partition_id = partition.id().clone();
         let completed_persistence_count = partition.completed_persistence_count();
@@ -413,6 +415,7 @@ mod tests {
     use data_types::PartitionId;
     use proto::ingester_query_response_metadata::PartitionIdentifier;
     use tonic::Code;
+    use trace::{ctx::SpanContext, RingBufferTraceCollector, TraceCollector};
 
     #[tokio::test]
     async fn sends_only_partition_hash_id_if_present() {
@@ -531,6 +534,44 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_encoded_spans_attached_to_collector() {
+        let ingester_id = IngesterId::new();
+
+        // A dummy batch of data.
+        let (batch, _) = make_batch!(
+            Int32Array("int" => vec![1, 2, 3]),
+        );
+
+        let query_response = QueryResponse::new(PartitionStream::new(futures::stream::iter([
+            PartitionResponse::new(vec![batch], ARBITRARY_TRANSITION_PARTITION_ID.clone(), 42),
+        ])));
+
+        let histogram = Arc::new(
+            metric::Registry::default()
+                .register_metric::<DurationHistogram>("test", "")
+                .recorder([]),
+        );
+
+        // Initialise a tracing backend to capture the emitted traces.
+        let trace_collector = Arc::new(RingBufferTraceCollector::new(5));
+        let trace_observer: Arc<dyn TraceCollector> = Arc::new(Arc::clone(&trace_collector));
+        let span_ctx = SpanContext::new(Arc::clone(&trace_observer));
+        let query_span = span_ctx.child("query span");
+
+        // test with encode_response
+        let call_chain = encode_response(query_response, ingester_id, Some(query_span), histogram);
+        call_chain.collect::<Vec<_>>().await;
+
+        let spans = trace_collector.spans();
+        assert_matches!(spans.as_slice(), [parent_span, frame_encoding_span_1, frame_encoding_span_2, frame_encoding_span_3] => {
+            assert_eq!(parent_span.name, "query span");
+            assert_eq!(frame_encoding_span_1.name, "frame encoding");
+            assert_eq!(frame_encoding_span_2.name, "frame encoding");
+            assert_eq!(frame_encoding_span_3.name, "frame encoding");
+        });
+    }
+
     /// Regression test for https://github.com/influxdata/idpe/issues/17408
     #[tokio::test]
     async fn test_chunks_with_different_schemas() {