From 4652dbec27b23dd1f9c49fcffc0deb20ae2e2e2d Mon Sep 17 00:00:00 2001 From: Loakesh Indiran <66478092+Loaki07@users.noreply.github.com> Date: Mon, 25 May 2026 19:45:43 +0530 Subject: [PATCH] fix: TopKHeap for non-ts ORDER BY in `latest_stream` (#12134) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem `sort_by=duration` (and any non-timestamp sort) in the `/traces/latest_stream` endpoint had three bugs: 1. All partitions collapsed into one full-range scan 2. Q1 results re-sorted by `trace_start_time`, discarding SQL engine order 3. `fetch_size` capped at `query_default_limit` (1000), so top-k ranked on a truncated sample ## Fix Mirrors the approach from #11654 (_search_stream): - Detect non-ts ORDER BY generically via \`sql_order_expr\` prefix — no hardcoded field names - Run Q1 across **all** partitions; feed deduped hits into a \`TopKHeap\` bounded to \`k = from + size\` - Q2 (span details) runs **once** after all partitions complete, over the narrowed time window of the deliverable traces - Final hits emitted in heap order — correct global rank by ORDER BY column ## Refactor Extracted `run_q2_for_traces` async helper to eliminate ~300 lines of duplicated code between the ts-path and non-ts-path. Both paths now call the same helper for the `TraceResponseItem` build + Q2a (per-trace aggregates) + Q2b (per-service breakdown). Also removed the dead `"duration"` sort arm from the ts-path and the now-unused `sort_by` parameter from `process_latest_traces_stream`. ## Known Limitation `zo_sql_duration` is computed per-partition, so traces whose spans straddle a partition boundary get an underestimated duration. Rare in practice; a 2-pass global aggregation would be needed to fix. ## Related Companion to #11654 — correct pagination for ORDER BY non-timestamp column in `_search_stream` --- src/handler/http/request/traces/mod.rs | 820 ++++++++++++++----------- 1 file changed, 467 insertions(+), 353 deletions(-) diff --git a/src/handler/http/request/traces/mod.rs b/src/handler/http/request/traces/mod.rs index 85b35db31d..38641959e2 100644 --- a/src/handler/http/request/traces/mod.rs +++ b/src/handler/http/request/traces/mod.rs @@ -52,7 +52,10 @@ use crate::{ CONTENT_TYPE_JSON, CONTENT_TYPE_PROTO, search::error_utils::map_error_to_http_response, }, }, - service::{search as SearchService, traces}, + service::{ + search::{self as SearchService, streaming::sorting::TopKHeap}, + traces, + }, }; pub mod dag; @@ -1172,7 +1175,6 @@ pub async fn get_latest_traces_stream( from, size, timeout, - sort_by, sort_order.to_string(), sql_order_expr, validated_schema, @@ -1235,7 +1237,6 @@ async fn process_latest_traces_stream( from: i64, size: i64, timeout: i64, - sort_by: String, sort_order: String, sql_order_expr: String, validated_schema: Option, @@ -1353,29 +1354,41 @@ async fn process_latest_traces_stream( } let total_partitions = partitions.len(); - // Partitions from search_partition are already in DESC order (newest-first) by default, - // since the partition SQL has no ORDER BY and the default is OrderBy::Desc. - let partitions_desc = if sql_order_expr == "zo_sql_timestamp DESC" { - log::info!( - "[TRACES_STREAM trace_id {trace_id}] using {} time partitions, order=DESC", - partitions.len() - ); - partitions - } else if sql_order_expr == "zo_sql_timestamp ASC" { + // Non-ts ORDER BY uses TopKHeap across all partitions — no single-partition collapse. + // ts ASC reverses partition order; everything else (ts DESC and all non-ts) is newest-first. + let is_non_ts_order_by = !sql_order_expr.starts_with("zo_sql_timestamp"); + let partitions_desc: Vec<[i64; 2]> = if sql_order_expr == "zo_sql_timestamp ASC" { log::info!( "[TRACES_STREAM trace_id {trace_id}] using {} time partitions, order=ASC", partitions.len() ); - partitions.into_iter().rev().collect::>() + partitions.into_iter().rev().collect() } else { - // order by other fields can't be multiple partitions log::info!( - "[TRACES_STREAM trace_id {trace_id}] sort_by non-timestamp ({sql_order_expr}), \ - forcing single partition [{}, {}]", - partition_req.start_time, - partition_req.end_time, + "[TRACES_STREAM trace_id {trace_id}] using {} time partitions, order={}", + partitions.len(), + if is_non_ts_order_by { + format!("non-ts ({sql_order_expr}), TopKHeap") + } else { + "DESC".to_string() + } ); - vec![[partition_req.start_time, partition_req.end_time]] + partitions + }; + + // For non-ts ORDER BY: TopKHeap merges Q1 results across all partitions. + // k = from + size so the heap is bounded and O(k) memory at all times. + // Q2 runs once after all partitions are drained — no per-partition streaming. + let heap_k = if from + size > 0 { + (from + size) as usize + } else { + get_config().limit.query_default_limit as usize + }; + let (non_ts_col, non_ts_is_desc) = parse_order_expr(&sql_order_expr); + let mut topk_heap: Option = if is_non_ts_order_by { + Some(TopKHeap::new(heap_k, &[(non_ts_col, non_ts_is_desc)])) + } else { + None }; // `from` is a global offset: skip the first `from` hits across all partitions, @@ -1406,12 +1419,16 @@ async fn process_latest_traces_stream( let p_start = partition[0]; let p_end = partition[1]; - // Ask for enough rows to cover remaining offset + remaining needed, - // capped at query_default_limit (default 1000) to prevent oversized Q1 scans. let remaining_to_skip = (hits_to_skip - hits_seen).max(0); let remaining_needed = hits_to_deliver - hits_delivered; - let fetch_size = - (remaining_to_skip + remaining_needed).min(get_config().limit.query_default_limit); + // Non-ts ORDER BY: each partition fetches its local top-k = heap_k rows so the heap can + // select the global top-k after all partitions complete. No query_default_limit cap here. + // ts ORDER BY: cap per partition to avoid oversized Q1 scans. + let fetch_size = if is_non_ts_order_by { + heap_k as i64 + } else { + (remaining_to_skip + remaining_needed).min(get_config().limit.query_default_limit) + }; // ---------- Query 1: aggregate over this partition ---------- let mut req1 = base_req.clone(); @@ -1471,8 +1488,29 @@ async fn process_latest_traces_stream( continue; } - // Sort Q1 hits by trace_start_time descending (same order as the final result). - // Q1 returns traces ordered by zo_sql_timestamp DESC, but we need start_time order. + // Non-ts ORDER BY: push this partition's Q1 hits into the global heap and move on. + // Seen trace IDs are registered here so later partitions skip duplicates. + // Q2 and final streaming happen after all partitions complete. + if is_non_ts_order_by { + for item in &deduped_hits { + if let Some(tid) = item.get("trace_id").and_then(|v| v.as_str()) + && !tid.is_empty() + { + seen_trace_ids.insert(tid.to_string()); + } + } + if let Some(ref mut heap) = topk_heap { + heap.push_hits(deduped_hits); + } + let progress = ((idx + 1) * 100) / total_partitions; + let _ = sender + .send(Ok(StreamResponses::Progress { percent: progress })) + .await; + continue; + } + + // ts ORDER BY: sort Q1 hits by trace_start_time (Q1 orders by zo_sql_timestamp; UI expects + // trace_start_time order). Register seen IDs before slicing so dedup covers skipped traces. let mut sorted_hits = deduped_hits; sorted_hits.sort_by(|a, b| { let a_t = json::get_int_value(a.get("trace_start_time").unwrap_or_default()); @@ -1514,342 +1552,29 @@ async fn process_latest_traces_stream( } // ---------- Query 2: fetch span details only for the traces we will deliver ---------- - let mut traces_data: HashMap = - HashMap::with_capacity(deliverable_q1.len()); - let mut p_start_actual = p_start; - let mut p_end_actual = p_end; - - for item in &deliverable_q1 { - let tid = item - .get("trace_id") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let trace_start_time = json::get_int_value( - item.get("trace_start_time") - .unwrap_or(&serde_json::Value::Null), - ); - let trace_end_time = json::get_int_value( - item.get("trace_end_time") - .unwrap_or(&serde_json::Value::Null), - ); - // Expand window to cover actual span times (nanoseconds → microseconds) - if trace_start_time > 0 && trace_start_time / 1000 < p_start_actual { - p_start_actual = trace_start_time / 1000; - } - if trace_end_time > 0 && trace_end_time / 1000 > p_end_actual { - // becuase we search with < end_time, so we need to add 1 to the end_time - p_end_actual = trace_end_time / 1000 + 1; - } - traces_data.insert( - tid.clone(), - TraceResponseItem { - trace_id: tid, - start_time: trace_start_time, - end_time: trace_end_time, - duration: 0, - spans: [0, 0], - service_name: Vec::new(), - first_event: serde_json::Value::Null, - gen_ai_usage_input_tokens: json::get_int_value( - item.get("gen_ai_usage_details_input").unwrap_or_default(), - ), - gen_ai_usage_output_tokens: json::get_int_value( - item.get("gen_ai_usage_details_output").unwrap_or_default(), - ), - gen_ai_usage_total_tokens: json::get_int_value( - item.get("gen_ai_usage_details_total").unwrap_or_default(), - ), - gen_ai_usage_cost: json::get_float_value( - item.get("gen_ai_usage_cost_details").unwrap_or_default(), - ), - gen_ai_input_messages: item.get("gen_ai_input_messages").cloned(), - models: item - .get("gen_ai_response_models") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(String::from)) - .collect() - }) - .unwrap_or_default(), - }, - ); - } - - // Sanitize trace IDs before interpolating into SQL: allow only hex chars and hyphens. - // Trace IDs originate from ingested data and could contain injected SQL if not validated. - let sanitized_ids: Vec = traces_data - .keys() - .map(|tid| { - tid.chars() - .filter(|c| c.is_ascii_hexdigit() || *c == '-') - .collect::() - }) - .filter(|tid| !tid.is_empty()) - .collect(); - // Check whether reference_parent_span_id exists in the stream schema. - let has_ref_parent_id = infra::schema::get_stream_schema_from_cache( - org_id.as_str(), - stream_name.as_str(), - StreamType::Traces, - ) - .await - .map(|s| s.field_with_name("reference_parent_span_id").is_ok()) - .unwrap_or(true); - - let stream_root_service_name_expr = if has_ref_parent_id { - "max(CASE WHEN reference_parent_span_id IS NULL OR reference_parent_span_id = '' THEN service_name END)" - } else { - "null" - }; - let stream_root_operation_name_expr = if has_ref_parent_id { - "max(CASE WHEN reference_parent_span_id IS NULL OR reference_parent_span_id = '' THEN operation_name END)" - } else { - "null" - }; - - let trace_ids_str = sanitized_ids.join("','"); - - // Q2a: per-trace aggregates. One row per trace_id (bounded by N traces) so this - // never trips the search-service row cap, which the previous SELECT * approach - // could hit when traces had tens of thousands of spans. - let detail_sql = format!( - "SELECT trace_id, \ - count(*) AS span_count, \ - sum(CASE WHEN span_status = 'ERROR' THEN 1 ELSE 0 END) AS error_count, \ - min(start_time) AS min_start_time, \ - max(end_time) AS max_end_time, \ - max(duration) AS max_duration, \ - count(DISTINCT service_name) AS service_count, \ - {stream_root_service_name_expr} AS root_service_name, \ - {stream_root_operation_name_expr} AS root_operation_name, \ - first_value(service_name ORDER BY {TIMESTAMP_COL_NAME} ASC) AS first_service_name, \ - first_value(operation_name ORDER BY {TIMESTAMP_COL_NAME} ASC) AS first_operation_name \ - FROM \"{stream_name}\" WHERE trace_id IN ('{trace_ids_str}') GROUP BY trace_id" - ); - - let mut req2 = base_req.clone(); - req2.query.sql = detail_sql; - req2.query.from = 0; - // Q2a returns one row per trace_id, so size = number of traces. - req2.query.size = traces_data.len() as i64; - req2.query.start_time = p_start_actual; - req2.query.end_time = p_end_actual; - - // Trace IDs that have more than one service — these need a follow-up Q2b. - // Single-service traces are fully populated from Q2a alone. - // Track the total expected (trace_id, service_name) row count so Q2b can - // size its request exactly and skip pagination. - let mut multi_service_tids: Vec = Vec::new(); - let mut multi_service_total: i64 = 0; - - if sender.is_closed() { - return; - } - let detail_res = match SearchService::cache::search( + let Some(traces_data) = run_q2_for_traces( + &deliverable_q1, + p_start, + p_end, + &base_req, &trace_id, &org_id, stream_type, - Some(user_id.clone()), - &req2, - "".to_string(), - false, - None, - false, + &user_id, + &stream_name, + &sender, + &format!("partition [{p_start},{p_end}]"), ) .await - { - Ok(res) => res, - Err(e) => { - log::error!( - "[TRACES_STREAM trace_id {trace_id}] Query2 error on partition [{p_start},{p_end}]: {e}" - ); - let _ = sender.send(Err(e)).await; - return; - } + else { + return; }; - for item in detail_res.hits { - let tid = item - .get("trace_id") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - if tid.is_empty() { - continue; - } - let span_count = json::get_int_value(item.get("span_count").unwrap_or_default()); - let error_count = json::get_int_value(item.get("error_count").unwrap_or_default()); - let min_start = json::get_int_value(item.get("min_start_time").unwrap_or_default()); - let max_end = json::get_int_value(item.get("max_end_time").unwrap_or_default()); - let max_duration = json::get_int_value(item.get("max_duration").unwrap_or_default()); - let service_count = json::get_int_value(item.get("service_count").unwrap_or_default()); - let root_service_name = item - .get("root_service_name") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let root_operation_name = item - .get("root_operation_name") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let first_service_name = item - .get("first_service_name") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let first_operation_name = item - .get("first_operation_name") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - - // Prefer the true root span's labels; fall back to earliest span. - let event_service_name = if !root_service_name.is_empty() { - root_service_name - } else { - first_service_name - }; - let event_operation_name = if !root_operation_name.is_empty() { - root_operation_name - } else { - first_operation_name - }; - - if let Some(trace) = traces_data.get_mut(&tid) { - trace.spans[0] = span_count.try_into().unwrap_or_default(); - trace.spans[1] = error_count.try_into().unwrap_or_default(); - if trace.start_time == 0 || trace.start_time > min_start { - trace.start_time = min_start; - } - if trace.end_time < max_end { - trace.end_time = max_end; - } - trace.duration = max_duration; - // update the duration if the duration is not correct - if trace.end_time - trace.start_time > trace.duration * 1000 { - trace.duration = (trace.end_time - trace.start_time) / 1000; - } - trace.first_event = serde_json::json!({ - "service_name": event_service_name, - "operation_name": event_operation_name, - }); - - if service_count <= 1 { - // Single service: populate the services list directly from Q2a; - // skip Q2b for this trace. - if !event_service_name.is_empty() { - trace.service_name.push(TraceServiceNameItem { - service_name: event_service_name, - count: span_count.try_into().unwrap_or_default(), - duration: max_duration, - }); - } - } else { - multi_service_tids.push(tid); - multi_service_total += service_count; - } - } - } - - // Q2b: per-(trace_id, service_name) breakdown, only for multi-service traces. - if !multi_service_tids.is_empty() { - // Same sanitization as Q2a above — trace IDs come from DB rows and must be - // validated before interpolating into SQL. - let multi_ids_str = multi_service_tids - .iter() - .map(|tid| { - tid.chars() - .filter(|c| c.is_ascii_hexdigit() || *c == '-') - .collect::() - }) - .filter(|tid| !tid.is_empty()) - .collect::>() - .join("','"); - let svc_sql = format!( - "SELECT trace_id, service_name, count(*) AS svc_count, max(duration) AS svc_duration \ - FROM \"{stream_name}\" WHERE trace_id IN ('{multi_ids_str}') \ - GROUP BY trace_id, service_name" - ); - - let mut req3 = base_req.clone(); - req3.query.sql = svc_sql; - req3.query.from = 0; - // Output is bounded to the total service count summed from Q2a, so a - // single request fetches everything. - req3.query.size = multi_service_total; - req3.query.start_time = p_start_actual; - req3.query.end_time = p_end_actual; - - if sender.is_closed() { - return; - } - let svc_res = match SearchService::cache::search( - &trace_id, - &org_id, - stream_type, - Some(user_id.clone()), - &req3, - "".to_string(), - false, - None, - false, - ) - .await - { - Ok(res) => res, - Err(e) => { - log::error!( - "[TRACES_STREAM trace_id {trace_id}] Query2b error on partition [{p_start},{p_end}]: {e}" - ); - let _ = sender.send(Err(e)).await; - return; - } - }; - - for item in svc_res.hits { - let tid = item - .get("trace_id") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let svc_name = item - .get("service_name") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let svc_count = json::get_int_value(item.get("svc_count").unwrap_or_default()); - let svc_duration = - json::get_int_value(item.get("svc_duration").unwrap_or_default()); - if let Some(trace) = traces_data.get_mut(&tid) { - trace.service_name.push(TraceServiceNameItem { - service_name: svc_name, - count: svc_count.try_into().unwrap_or_default(), - duration: svc_duration, - }); - } - } - } - - // Sort by start_time descending (Q2 may have refined start_time from actual span data). let mut partition_hits: Vec<&TraceResponseItem> = traces_data.values().collect(); - match sort_by.as_str() { - "duration" => { - if sort_order == "ASC" { - partition_hits.sort_by(|a, b| a.duration.cmp(&b.duration)); - } else { - partition_hits.sort_by(|a, b| b.duration.cmp(&a.duration)); - } - } - _ => { - if sort_order == "ASC" { - partition_hits.sort_by(|a, b| a.start_time.cmp(&b.start_time)); - } else { - partition_hits.sort_by(|a, b| b.start_time.cmp(&a.start_time)); - } - } + if sort_order == "ASC" { + partition_hits.sort_by(|a, b| a.start_time.cmp(&b.start_time)); + } else { + partition_hits.sort_by(|a, b| b.start_time.cmp(&a.start_time)); } let deliverable: Vec = partition_hits @@ -1897,9 +1622,398 @@ async fn process_latest_traces_stream( } } + // ── Non-ts ORDER BY: drain heap → Q2 once → stream final result ───────────────────────── + if is_non_ts_order_by { + let deliverable_q1 = topk_heap + .take() + .map(|h| h.into_sorted_vec(from as usize)) + .unwrap_or_default(); + + if deliverable_q1.is_empty() { + let _ = sender.send(Ok(StreamResponses::Done)).await; + return; + } + + let Some(traces_data) = run_q2_for_traces( + &deliverable_q1, + start_time, + end_time, + &base_req, + &trace_id, + &org_id, + stream_type, + &user_id, + &stream_name, + &sender, + "non-ts", + ) + .await + else { + return; + }; + + // Preserve heap ordering: walk deliverable_q1 in heap-drain order (already globally + // sorted by the non-ts column). Q2 may have refined start_time/duration but the + // relative ORDER BY column rank from Q1 is authoritative. + let partition_hits: Vec<&TraceResponseItem> = deliverable_q1 + .iter() + .filter_map(|item| { + let tid = item.get("trace_id").and_then(|v| v.as_str())?; + traces_data.get(tid) + }) + .collect(); + + let deliverable: Vec = partition_hits + .iter() + .map(|t| serde_json::to_value(t).unwrap_or(serde_json::Value::Null)) + .collect(); + + let mut results = config::meta::search::Response::default(); + for h in deliverable { + results.add_hit(&h); + } + results.set_total(total_across_partitions); + if !range_error.is_empty() { + results.function_error.push(range_error.clone()); + } + + let _ = sender + .send(Ok(StreamResponses::SearchResponse { + results, + streaming_aggs: false, + streaming_id: None, + time_offset: TimeOffset { + start_time, + end_time, + }, + })) + .await; + let _ = sender.send(Ok(StreamResponses::Done)).await; + return; + } + let _ = sender.send(Ok(StreamResponses::Done)).await; } +/// Extract ORDER BY column name and direction from a `" DESC|ASC"` expression. +fn parse_order_expr(sql_order_expr: &str) -> (String, bool) { + let mut parts = sql_order_expr.split_whitespace(); + let col = parts.next().unwrap_or("zo_sql_duration").to_string(); + let is_desc = parts + .next() + .map(|d| d.to_uppercase() != "ASC") + .unwrap_or(true); + (col, is_desc) +} + +/// Runs Q2a (per-trace aggregates) and Q2b (per-service breakdown) for a slice of Q1 hits, +/// building the `TraceResponseItem` map. +/// +/// `window_start` / `window_end` (microseconds) are the initial search window; they expand +/// to cover actual span times reported by Q1. +/// +/// Returns `Some(traces_data)` on success. On error the error is sent on `sender` and +/// `None` is returned — callers should `return` immediately. +#[allow(clippy::too_many_arguments)] +async fn run_q2_for_traces( + q1_hits: &[serde_json::Value], + window_start: i64, + window_end: i64, + base_req: &config::meta::search::Request, + req_trace_id: &str, + org_id: &str, + stream_type: StreamType, + user_id: &str, + stream_name: &str, + sender: &mpsc::Sender>, + log_ctx: &str, +) -> Option> { + let mut q2_start = window_start; + let mut q2_end = window_end; + let mut traces_data: HashMap = HashMap::with_capacity(q1_hits.len()); + + for item in q1_hits { + let tid = item + .get("trace_id") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let trace_start_time = + json::get_int_value(item.get("trace_start_time").unwrap_or_default()); + let trace_end_time = json::get_int_value(item.get("trace_end_time").unwrap_or_default()); + if trace_start_time > 0 && trace_start_time / 1000 < q2_start { + q2_start = trace_start_time / 1000; + } + if trace_end_time > 0 && trace_end_time / 1000 + 1 > q2_end { + q2_end = trace_end_time / 1000 + 1; + } + if tid.is_empty() { + continue; + } + traces_data.insert( + tid.clone(), + TraceResponseItem { + trace_id: tid, + start_time: trace_start_time, + end_time: trace_end_time, + duration: 0, + spans: [0, 0], + service_name: Vec::new(), + first_event: serde_json::Value::Null, + gen_ai_usage_input_tokens: json::get_int_value( + item.get("gen_ai_usage_details_input").unwrap_or_default(), + ), + gen_ai_usage_output_tokens: json::get_int_value( + item.get("gen_ai_usage_details_output").unwrap_or_default(), + ), + gen_ai_usage_total_tokens: json::get_int_value( + item.get("gen_ai_usage_details_total").unwrap_or_default(), + ), + gen_ai_usage_cost: json::get_float_value( + item.get("gen_ai_usage_cost_details").unwrap_or_default(), + ), + gen_ai_input_messages: item.get("gen_ai_input_messages").cloned(), + models: item + .get("gen_ai_response_models") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(), + }, + ); + } + + let sanitized_ids: Vec = traces_data + .keys() + .map(|tid| { + tid.chars() + .filter(|c| c.is_ascii_hexdigit() || *c == '-') + .collect::() + }) + .filter(|tid| !tid.is_empty()) + .collect(); + + let has_ref_parent_id = + infra::schema::get_stream_schema_from_cache(org_id, stream_name, StreamType::Traces) + .await + .map(|s| s.field_with_name("reference_parent_span_id").is_ok()) + .unwrap_or(true); + + let root_svc_expr = if has_ref_parent_id { + "max(CASE WHEN reference_parent_span_id IS NULL OR reference_parent_span_id = '' THEN service_name END)" + } else { + "null" + }; + let root_op_expr = if has_ref_parent_id { + "max(CASE WHEN reference_parent_span_id IS NULL OR reference_parent_span_id = '' THEN operation_name END)" + } else { + "null" + }; + + let trace_ids_str = sanitized_ids.join("','"); + let detail_sql = format!( + "SELECT trace_id, \ + count(*) AS span_count, \ + sum(CASE WHEN span_status = 'ERROR' THEN 1 ELSE 0 END) AS error_count, \ + min(start_time) AS min_start_time, \ + max(end_time) AS max_end_time, \ + max(duration) AS max_duration, \ + count(DISTINCT service_name) AS service_count, \ + {root_svc_expr} AS root_service_name, \ + {root_op_expr} AS root_operation_name, \ + first_value(service_name ORDER BY {TIMESTAMP_COL_NAME} ASC) AS first_service_name, \ + first_value(operation_name ORDER BY {TIMESTAMP_COL_NAME} ASC) AS first_operation_name \ + FROM \"{stream_name}\" WHERE trace_id IN ('{trace_ids_str}') GROUP BY trace_id" + ); + + let mut req2 = base_req.clone(); + req2.query.sql = detail_sql; + req2.query.from = 0; + req2.query.size = traces_data.len() as i64; + req2.query.start_time = q2_start; + req2.query.end_time = q2_end; + + let mut multi_service_tids: Vec = Vec::new(); + let mut multi_service_total: i64 = 0; + + if sender.is_closed() { + return None; + } + let detail_res = match SearchService::cache::search( + req_trace_id, + org_id, + stream_type, + Some(user_id.to_string()), + &req2, + "".to_string(), + false, + None, + false, + ) + .await + { + Ok(res) => res, + Err(e) => { + log::error!("[TRACES_STREAM trace_id {req_trace_id}] Q2a error ({log_ctx}): {e}"); + let _ = sender.send(Err(e)).await; + return None; + } + }; + + for item in detail_res.hits { + let tid = item + .get("trace_id") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + if tid.is_empty() { + continue; + } + let span_count = json::get_int_value(item.get("span_count").unwrap_or_default()); + let error_count = json::get_int_value(item.get("error_count").unwrap_or_default()); + let min_start = json::get_int_value(item.get("min_start_time").unwrap_or_default()); + let max_end = json::get_int_value(item.get("max_end_time").unwrap_or_default()); + let max_duration = json::get_int_value(item.get("max_duration").unwrap_or_default()); + let service_count = json::get_int_value(item.get("service_count").unwrap_or_default()); + let root_service_name = item + .get("root_service_name") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let root_operation_name = item + .get("root_operation_name") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let first_service_name = item + .get("first_service_name") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let first_operation_name = item + .get("first_operation_name") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + + let event_service_name = if !root_service_name.is_empty() { + root_service_name + } else { + first_service_name + }; + let event_operation_name = if !root_operation_name.is_empty() { + root_operation_name + } else { + first_operation_name + }; + + if let Some(trace) = traces_data.get_mut(&tid) { + trace.spans[0] = span_count.try_into().unwrap_or_default(); + trace.spans[1] = error_count.try_into().unwrap_or_default(); + if trace.start_time == 0 || trace.start_time > min_start { + trace.start_time = min_start; + } + if trace.end_time < max_end { + trace.end_time = max_end; + } + trace.duration = max_duration; + if trace.end_time - trace.start_time > trace.duration * 1000 { + trace.duration = (trace.end_time - trace.start_time) / 1000; + } + trace.first_event = serde_json::json!({ + "service_name": event_service_name, + "operation_name": event_operation_name, + }); + if service_count <= 1 { + if !event_service_name.is_empty() { + trace.service_name.push(TraceServiceNameItem { + service_name: event_service_name, + count: span_count.try_into().unwrap_or_default(), + duration: max_duration, + }); + } + } else { + multi_service_tids.push(tid); + multi_service_total += service_count; + } + } + } + + if !multi_service_tids.is_empty() { + let multi_ids_str = multi_service_tids + .iter() + .map(|tid| { + tid.chars() + .filter(|c| c.is_ascii_hexdigit() || *c == '-') + .collect::() + }) + .filter(|tid| !tid.is_empty()) + .collect::>() + .join("','"); + let svc_sql = format!( + "SELECT trace_id, service_name, count(*) AS svc_count, max(duration) AS svc_duration \ + FROM \"{stream_name}\" WHERE trace_id IN ('{multi_ids_str}') \ + GROUP BY trace_id, service_name" + ); + let mut req3 = base_req.clone(); + req3.query.sql = svc_sql; + req3.query.from = 0; + req3.query.size = multi_service_total; + req3.query.start_time = q2_start; + req3.query.end_time = q2_end; + + if sender.is_closed() { + return None; + } + let svc_res = match SearchService::cache::search( + req_trace_id, + org_id, + stream_type, + Some(user_id.to_string()), + &req3, + "".to_string(), + false, + None, + false, + ) + .await + { + Ok(res) => res, + Err(e) => { + log::error!("[TRACES_STREAM trace_id {req_trace_id}] Q2b error ({log_ctx}): {e}"); + let _ = sender.send(Err(e)).await; + return None; + } + }; + for item in svc_res.hits { + let tid = item + .get("trace_id") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let svc_name = item + .get("service_name") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let svc_count = json::get_int_value(item.get("svc_count").unwrap_or_default()); + let svc_duration = json::get_int_value(item.get("svc_duration").unwrap_or_default()); + if let Some(trace) = traces_data.get_mut(&tid) { + trace.service_name.push(TraceServiceNameItem { + service_name: svc_name, + count: svc_count.try_into().unwrap_or_default(), + duration: svc_duration, + }); + } + } + } + + Some(traces_data) +} + #[derive(Debug, Serialize)] struct TraceResponseItem { trace_id: String,