mirror of
https://github.com/openobserve/openobserve.git
synced 2026-11-05 16:10:43 +00:00
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user