fix: TopKHeap for non-ts ORDER BY in latest_stream (#12134)

## 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`
This commit is contained in:
Loakesh Indiran
2026-05-25 14:15:43 +00:00
committed by GitHub
parent 31ddecd23b
commit 4652dbec27
+467 -353
View File
@@ -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<schema_compat::ValidatedLlmSchema>,
@@ -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::<Vec<_>>()
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<TopKHeap> = 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<String, TraceResponseItem> =
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<String> = traces_data
.keys()
.map(|tid| {
tid.chars()
.filter(|c| c.is_ascii_hexdigit() || *c == '-')
.collect::<String>()
})
.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<String> = 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::<String>()
})
.filter(|tid| !tid.is_empty())
.collect::<Vec<String>>()
.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<serde_json::Value> = 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<serde_json::Value> = 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 `"<col> 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<Result<StreamResponses, infra::errors::Error>>,
log_ctx: &str,
) -> Option<HashMap<String, TraceResponseItem>> {
let mut q2_start = window_start;
let mut q2_end = window_end;
let mut traces_data: HashMap<String, TraceResponseItem> = 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<String> = traces_data
.keys()
.map(|tid| {
tid.chars()
.filter(|c| c.is_ascii_hexdigit() || *c == '-')
.collect::<String>()
})
.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<String> = 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::<String>()
})
.filter(|tid| !tid.is_empty())
.collect::<Vec<String>>()
.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,