feat(tvix/castore): instrument tokio task with current span

By default tokio::spawn does not instrument the spawned task with the
current spawn (https://github.com/tokio-rs/tokio/discussions/6008), do
this manually for all tokio::spawn functions in functions that are
instrumented.

Change-Id: I83dd8145b3a62421454aff57d34180cebbee8304
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11864
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Autosubmit: Simon Hauser <simon.hauser@helsinki-systems.de>
This commit is contained in:
Simon Hauser 2024-06-20 16:10:12 +02:00
parent bd8d74a3ee
commit 2b20d8d82d
3 changed files with 27 additions and 18 deletions

View file

@ -12,7 +12,7 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
use tonic::async_trait;
use tonic::Code;
use tonic::{transport::Channel, Status};
use tracing::{instrument, warn};
use tracing::{instrument, warn, Instrument as _};
/// Connects to a (remote) tvix-store DirectoryService over gRPC.
#[derive(Clone)]
@ -194,14 +194,17 @@ impl DirectoryService for GRPCDirectoryService {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let task: JoinHandle<Result<proto::PutDirectoryResponse, Status>> = spawn(async move {
let s = grpc_client
.put(UnboundedReceiverStream::new(rx))
.await?
.into_inner();
let task: JoinHandle<Result<proto::PutDirectoryResponse, Status>> = spawn(
async move {
let s = grpc_client
.put(UnboundedReceiverStream::new(rx))
.await?
.into_inner();
Ok(s)
});
Ok(s)
} // instrument the task with the current span, this is not done by default
.in_current_span(),
);
Box::new(GRPCPutter {
rq: Some((task, tx)),