diff --git a/src/frontend/src/instance/jaeger.rs b/src/frontend/src/instance/jaeger.rs index ac59698678b..7a7a2c8b44d 100644 --- a/src/frontend/src/instance/jaeger.rs +++ b/src/frontend/src/instance/jaeger.rs @@ -35,6 +35,7 @@ use datafusion::dataframe::DataFrame; use datafusion::execution::SessionStateBuilder; use datafusion::execution::context::SessionContext; use datafusion::functions::core::expr_fn::coalesce; +use datafusion::functions_aggregate::expr_fn::max; use datafusion::functions_window::expr_fn::row_number; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{Expr, ExprFunctionExt, SortExpr, col, lit, lit_timestamp_nano, wildcard}; @@ -59,6 +60,7 @@ use crate::instance::Instance; const DEFAULT_LIMIT: usize = 2000; const KEY_RN: &str = "greptime_rn"; +const KEY_LATEST_MATCH: &str = "greptime_latest_match"; impl Instance { async fn check_jaeger_query_permission(&self, ctx: &QueryContextRef) -> ServerResult<()> { @@ -226,11 +228,11 @@ impl JaegerQueryHandler for Instance { filters.push(col(DURATION_NANO_COLUMN).lt_eq(lit(max_duration))); } - // Get all distinct trace ids that match the filters. + // Get the ids of the most recent traces that match the filters. // It's equivalent to the following SQL query: // // ``` - // SELECT DISTINCT trace_id + // SELECT trace_id // FROM // {db}.{trace_table} // WHERE @@ -240,17 +242,16 @@ impl JaegerQueryHandler for Instance { // timestamp <= {end_time} AND // duration >= {min_duration} AND // duration <= {max_duration} + // GROUP BY trace_id + // ORDER BY max(timestamp) DESC, trace_id // LIMIT {limit} // ```. - let output = query_trace_table( + let output = find_latest_trace_ids( ctx.clone(), self, - vec![wildcard()], filters, - vec![], - Some(query_params.limit.unwrap_or(DEFAULT_LIMIT)), + query_params.limit.unwrap_or(DEFAULT_LIMIT), query_params.tags, - vec![col(TRACE_ID_COLUMN)], ) .await?; @@ -338,6 +339,79 @@ async fn query_trace_table( tags: Option>, distincts: Vec, ) -> ServerResult { + let dataframe = + build_trace_dataframe(&ctx, instance, selects, filters, tags, &distincts).await?; + + // Apply the distinct if needed. + let dataframe = if !distincts.is_empty() { + dataframe + .distinct_on(distincts.clone(), distincts, None) + .context(DataFusionSnafu)? + } else { + dataframe + }; + + // Apply the sorts if needed. + let dataframe = if !sorts.is_empty() { + dataframe.sort(sorts).context(DataFusionSnafu)? + } else { + dataframe + }; + + // Apply the limit if needed. + let dataframe = if let Some(limit) = limit { + dataframe.limit(0, Some(limit)).context(DataFusionSnafu)? + } else { + dataframe + }; + + execute_trace_dataframe(dataframe).await +} + +/// Returns the ids of the `limit` traces whose latest span matching `filters` +/// and `tags` is the most recent, newest first. +/// +/// Ordering before the limit keeps a search that matches more traces than the +/// limit deterministic. +async fn find_latest_trace_ids( + ctx: QueryContextRef, + instance: &Instance, + filters: Vec, + limit: usize, + tags: Option>, +) -> ServerResult { + let dataframe = + build_trace_dataframe(&ctx, instance, vec![wildcard()], filters, tags, &[]).await?; + let dataframe = dataframe + .aggregate( + vec![col(TRACE_ID_COLUMN)], + vec![max(col(TIMESTAMP_COLUMN)).alias(KEY_LATEST_MATCH)], + ) + .context(DataFusionSnafu)? + .sort(vec![ + col(KEY_LATEST_MATCH).sort(false, false), + col(TRACE_ID_COLUMN).sort(true, false), + ]) + .context(DataFusionSnafu)? + .limit(0, Some(limit)) + .context(DataFusionSnafu)? + .select(vec![col(TRACE_ID_COLUMN)]) + .context(DataFusionSnafu)?; + + execute_trace_dataframe(dataframe).await +} + +/// Reads the trace table, or the services or operations lookup table that +/// answers `selects` and `distincts`, then projects `selects` and applies +/// `filters` and the Jaeger tag filters. +async fn build_trace_dataframe( + ctx: &QueryContextRef, + instance: &Instance, + selects: Vec, + filters: Vec, + tags: Option>, + distincts: &[Expr], +) -> ServerResult { let trace_table_name = ctx .extension(JAEGER_QUERY_TABLE_NAME_KEY) .unwrap_or(TRACE_TABLE_NAME); @@ -366,7 +440,7 @@ async fn query_trace_table( ctx.current_catalog(), &ctx.current_schema(), table_name, - Some(&ctx), + Some(ctx), ) .await? .with_context(|| TableNotFoundSnafu { @@ -403,30 +477,12 @@ async fn query_trace_table( df.filter(expr).context(DataFusionSnafu) })?; - // Apply the distinct if needed. - let dataframe = if !distincts.is_empty() { - dataframe - .distinct_on(distincts.clone(), distincts, None) - .context(DataFusionSnafu)? - } else { - dataframe - }; + Ok(dataframe) +} - // Apply the sorts if needed. - let dataframe = if !sorts.is_empty() { - dataframe.sort(sorts).context(DataFusionSnafu)? - } else { - dataframe - }; - - // Apply the limit if needed. - let dataframe = if let Some(limit) = limit { - dataframe.limit(0, Some(limit)).context(DataFusionSnafu)? - } else { - dataframe - }; - - // Execute the query and collect the result. +/// Executes `dataframe` and wraps its stream into an [`Output`] with dictionary +/// columns decoded to values. +async fn execute_trace_dataframe(dataframe: DataFrame) -> ServerResult { let stream = dataframe.execute_stream().await.context(DataFusionSnafu)?; let output = Output::new_with_stream(Box::pin( diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f3b439e19bc..1bb4af7839c 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -230,6 +230,7 @@ macro_rules! http_tests { test_jaeger_v3_query_api, test_jaeger_query_api_for_trace_v1, test_jaeger_query_api_for_trace_v2, + test_jaeger_search_limit_returns_latest_traces, test_influxdb_write, test_influxdb_write_with_hints, @@ -12480,6 +12481,88 @@ pub async fn test_jaeger_query_api_for_trace_v2( Ok(()) } +pub async fn test_jaeger_search_limit_returns_latest_traces(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_jaeger_search_limit").await; + let client = TestClient::new(app).await; + + // Twelve single-span traces spread over all trace partitions. The oldest + // traces get the smallest ids and are written first, so neither id nor + // write order points at the newest ones. + let now = Utc::now().timestamp_nanos_opt().unwrap() as u64; + let traces = (0..12u64) + .map(|i| { + let start = now - (60 - i) * 1_000_000_000; + (format!("{:x}{:031x}", i, i + 1), start) + }) + .collect::>(); + let spans = traces + .iter() + .map(|(trace_id, start)| { + json!({ + "traceId": trace_id, + "spanId": &trace_id[16..], + "name": "op", + "kind": 2, + "startTimeUnixNano": start.to_string(), + "endTimeUnixNano": (start + 1_000).to_string(), + }) + }) + .collect::>(); + let content = json!({ + "resourceSpans": [{ + "resource": {"attributes": [{"key": "service.name", "value": {"stringValue": "svc"}}]}, + "scopeSpans": [{"scope": {"name": "s"}, "spans": spans}], + }] + }); + let req: ExportTraceServiceRequest = serde_json::from_value(content).unwrap(); + let res = send_req( + &client, + vec![ + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ), + ( + HeaderName::from_static("x-greptime-pipeline-name"), + HeaderValue::from_static(GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME), + ), + ], + "/v1/otlp/v1/traces", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + + let mut expected = traces[9..] + .iter() + .map(|(trace_id, _)| trace_id.clone()) + .collect::>(); + expected.sort(); + let path = format!( + "/v1/jaeger/api/traces?service=svc&operation=op&start={}&end={}&limit=3", + (now - 120_000_000_000) / 1_000, + now / 1_000 + ); + for _ in 0..5 { + let res = client.get(&path).send().await; + assert_eq!(StatusCode::OK, res.status()); + let resp: Value = serde_json::from_str(&res.text().await).unwrap(); + let mut found = resp["data"] + .as_array() + .unwrap() + .iter() + .map(|trace| trace["traceID"].as_str().unwrap().to_string()) + .collect::>(); + found.sort(); + assert_eq!(found, expected); + } + + guard.remove_all().await; +} + pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) =