From 973eb69cd3913cef2b225ab31a1bb7130f1afef8 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 1 Oct 2024 13:17:17 +0100 Subject: [PATCH] try reduce memory usage of returned data --- proxy/src/serverless/sql_over_http.rs | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 55bf75cf76..af9e170847 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -678,6 +678,9 @@ async fn handle_db_inner( None => return Err(SqlOverHttpError::Cancelled(SqlOverHttpCancel::Connect)), }; + arena.params_arena.shrink_to_fit(); + arena.str_arena.shrink_to_fit(); + let mut response = Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, "application/json"); @@ -955,7 +958,8 @@ async fn query_batch( queries: BatchQueryData, parsed_headers: HttpHeaders, ) -> Result { - let mut results = Vec::with_capacity(queries.queries.len()); + let mut result_bytes = br#"{"results":["#.to_vec(); + let mut current_size = 0; for stmt in queries.queries { let query = pin!(query_to_json( @@ -971,7 +975,12 @@ async fn query_batch( match res { // TODO: maybe we should check that the transaction bit is set here Either::Left((Ok((_, values)), _cancelled)) => { - results.push(values); + if !result_bytes.is_empty() { + result_bytes.push(b','); + } + values + .serialize(&mut serde_json::Serializer::new(&mut result_bytes)) + .expect("serializing results to memory should always succeed"); } Either::Left((Err(e), _cancelled)) => { return Err(e); @@ -982,8 +991,10 @@ async fn query_batch( } } - let results = json!({ "results": results }); - let json_output = serde_json::to_string(&results).expect("json serialization should not fail"); + result_bytes.extend_from_slice(b"]}"); + + let json_output = + String::from_utf8(result_bytes).expect("JSON serialization should always be utf8"); Ok(json_output) }