From 462172b9bbf54c22707a4bf1c9ff37e12151dfb5 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 07:46:40 +0000 Subject: [PATCH] perf(servers): reduce temporary memory in protocol handling (#9249) * perf(prom): release the remote write v1 decode buffer before writing remote_write_v1 kept the decoded builder alive until the handler returned, so the decompressed request payload stayed resident across the downstream write or pipeline await. With 8 concurrent large requests that is one extra copy of every payload held for the whole write. Rows and pipeline values own their data, so the builder can be dropped as soon as the conversion is done. Sustained-write A/B, 8 runs per side, 50M samples each: jemalloc allocated median drops 7.8% (155.0-162.7 MiB -> 141.4-158.6 MiB); samples per CPU second is unchanged (-0.6%, fully overlapping ranges). Signed-off-by: Dennis Zhuang * perf(servers): move row values into SQL JSON responses The format=json renderer cloned every serde_json::Value and kept the whole row set alive while building the response. Move the values out instead, so each row is released as soon as it is converted. Slicing the row to the schema width keeps the panic on rows narrower than the schema; a plain zip would silently truncate them. Duplicate column names still resolve to the last value and extra row values are still ignored. Isolated conversion measurements: live peak drops 33% on a 4096-row 16 KiB string fixture and 36% on a nested-JSON fixture, with no measured slowdown. Signed-off-by: Dennis Zhuang * perf(prom): release the compressed remote write body after decompression Both the v1 and v2 decoders held the compressed Bytes until they returned, which spans the whole protobuf decode and row conversion. Decompression copies the payload into an independent buffer, so the body can go as soon as it succeeds. The compression fallback, decode errors and request counting are unchanged. Signed-off-by: Dennis Zhuang * revert(prom): keep the remote write v1 decode buffer until the write finishes This reverts commit 8232c5b3bb00a656031620d1b30929e604d42717. The v1 decoder fabricates `&'static [u8]` pointing into its own decode buffer (prom_remote_write/types.rs), so the compiler checks nothing about that buffer's lifetime. Holding the builder until the handler returns is what keeps the decoder safe by construction; releasing it early made that safety depend on every consumer copying out of the buffer, which holds today but nothing enforces. Document the requirement at the binding instead. Signed-off-by: Dennis Zhuang * perf(prom): release the remote write v1 decode buffer before writing This reverts commit ca3af722d00ede48e259a798ee55c9320391c13e, restoring 8232c5b3bb00a656031620d1b30929e604d42717. The decoder borrows the decompressed buffer while parsing, but copies everything out when it builds rows: tag values through `PromValidationMode::decode_string`, column names through `to_owned`, and the only live borrows (`TableBuilder::col_indexes`) are dropped inside `as_insert_requests`. The resulting `ContextReq` holds prost types with no lifetime parameters, so it cannot reference the buffer. Record that at the binding so the next reader does not have to re-derive it from three files. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/servers/src/http/prom_store.rs | 12 ++++- src/servers/src/http/result/json_result.rs | 57 ++++++++++++++++++++-- src/servers/src/prom_remote_write/mod.rs | 2 + src/servers/src/prom_remote_write/v2.rs | 2 + 4 files changed, 66 insertions(+), 7 deletions(-) diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index cb7b6307c4b..60c000403fd 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -175,12 +175,20 @@ async fn remote_write_v1( processor.set_pipeline(pipeline_handler, query_ctx.clone(), pipeline_def); } - let mut req = decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + let mut decoded = + decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + // Parsing borrows the decode buffer, but row building copies out of it: tag + // values through `decode_string`, column names through `to_owned`, and the + // borrowing `col_indexes` dies inside `as_insert_requests`. Nothing below + // references the buffer, so it need not span the write. let req = if processor.use_pipeline { + drop(decoded); processor.exec_pipeline().await? } else { - req.as_insert_requests() + let req = decoded.as_insert_requests(); + drop(decoded); + req }; let batches = into_prom_write_batches(req, query_ctx); diff --git a/src/servers/src/http/result/json_result.rs b/src/servers/src/http/result/json_result.rs index 537ea66cd02..74a75f26309 100644 --- a/src/servers/src/http/result/json_result.rs +++ b/src/servers/src/http/result/json_result.rs @@ -90,17 +90,20 @@ impl IntoResponse for JsonResponse { .to_string(), Some(GreptimeQueryOutput::Records(records)) => { - let schema = records.schema(); + // Borrow the schema field directly so the rows can be moved out. + let schema = &records.schema; let data: Vec> = records .rows - .iter() - .map(|row| { + .into_iter() + .map(|mut row| { + // Slicing keeps the out-of-bounds panic for short rows; + // a plain zip would silently truncate them. schema .column_schemas .iter() - .enumerate() - .map(|(i, col)| (col.name.clone(), row[i].clone())) + .zip(row[..schema.column_schemas.len()].iter_mut()) + .map(|(col, value)| (col.name.clone(), value.take())) .collect::>() }) .collect(); @@ -133,3 +136,47 @@ impl IntoResponse for JsonResponse { .into_response() } } + +#[cfg(test)] +mod tests { + use axum::body::to_bytes; + + use super::*; + + #[tokio::test] + async fn test_records_response_preserves_values_and_duplicate_columns() { + let response: JsonResponse = serde_json::from_value(json!({ + "output": [{"records": { + "schema": {"column_schemas": [ + {"name": "duplicate", "data_type": "String"}, + {"name": "nested", "data_type": "Json"}, + {"name": "duplicate", "data_type": "String"}, + {"name": "escaped\"column", "data_type": "String"} + ]}, + "rows": [ + ["discarded", {"array": [null, true, "中文"]}, "last", "line\n\\\""], + ["discarded", [1, {"key": "value"}], null, ""] + ] + }}], + "execution_time_ms": 7 + })) + .unwrap(); + + let response = response.into_response(); + assert_eq!(response.status(), axum::http::StatusCode::OK); + assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json"); + assert_eq!(response.headers()[&GREPTIME_DB_HEADER_FORMAT], "json"); + assert_eq!(response.headers()[&GREPTIME_DB_HEADER_EXECUTION_TIME], "7"); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert_eq!( + serde_json::from_slice::(&body).unwrap(), + json!({ + "data": [ + {"duplicate": "last", "nested": {"array": [null, true, "中文"]}, "escaped\"column": "line\n\\\""}, + {"duplicate": null, "nested": [1, {"key": "value"}], "escaped\"column": ""} + ], + "execution_time_ms": 7 + }) + ); + } +} diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index 73f968b6084..c9aabfb2745 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -77,6 +77,8 @@ pub fn decode_remote_write_request( // fallback to the other compression method try_decompress(!is_zstd, &body[..])? }; + // Decompression copied the payload out, so the compressed body is no longer needed. + drop(body); let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default); diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index b388071106f..2f1f4d410c4 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -143,6 +143,8 @@ pub(crate) fn decode_remote_write_v2( } else { try_decompress(!is_zstd, &body[..])? }; + // Decompression copied the payload out, so the compressed body is no longer needed. + drop(body); let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; drop(decode_timer);