rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-31 12:20:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: scope spawned task with trace id (#2419)

Browse Source 
* feat: scope spawned task with trace id

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix test

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2023-09-17 04:05:28 -05:00
committed by GitHub
parent 542e863ecc
commit 693e8de83a
6 changed files with 44 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/datanode/src/region_server.rs
View File

@@ -135,23 +135,19 @@ impl RegionServerHandler for RegionServer {
.context(ExecuteGrpcRequestSnafu)?;
let join_tasks = requests.into_iter().map(|(region_id, req)| {
let self_to_move = self.clone();
self.inner
.runtime
.spawn(async move { self_to_move.handle_request(region_id, req).await })
async move { self_to_move.handle_request(region_id, req).await }
});
let results = try_join_all(join_tasks)
.await
.context(servers_error::JoinTaskSnafu)?;
.map_err(BoxedError::new)
.context(ExecuteGrpcRequestSnafu)?;
// merge results by simply sum up affected rows.
// only insert/delete will have multiple results.
let mut affected_rows = 0;
for result in results {
match result
.map_err(BoxedError::new)
.context(ExecuteGrpcRequestSnafu)?
{
match result {
Output::AffectedRows(rows) => affected_rows += rows,
Output::Stream(_) | Output::RecordBatches(_) => {
// TODO: change the output type to only contains `affected_rows`
@@ -181,10 +177,15 @@ impl FlightCraft for RegionServer {
let ticket = request.into_inner().ticket;
let request = QueryRequest::decode(ticket.as_ref())
.context(servers_error::InvalidFlightTicketSnafu)?;
let trace_id = request
.header
.as_ref()
.map(|h| h.trace_id)
.unwrap_or_default();
let result = self.handle_read(request).await?;
let stream = Box::pin(FlightRecordBatchStream::new(result));
let stream = Box::pin(FlightRecordBatchStream::new(result, trace_id));
Ok(Response::new(stream))
}
}