tonic icon indicating copy to clipboard operation
tonic copied to clipboard

feat(grpc): Add task-local context propagation

Open sauravzg opened this issue 1 month ago • 0 comments

This commit introduces a new task-local context propagation mechanism to gRPC, enabling context to be preserved across await points in asynchronous tasks.

The implementation is based on a combination of:

  • Otel FutureExt: Adopts the extension trait pattern from OpenTelemetry's FutureExt to allow fluent context attachment via .with_context(ctx).
  • Tokio Task Local: Implements a runtime-agnostic task-local storage mechanism similar to tokio::task_local!, ensuring context is correctly scoped and propagated.

Changes:

  • Add task_local_context module for managing context scope and restoring context to and from thread local.
  • Add FutureExt and StreamExt traits in extensions.rs.
  • Update Context to use Arc<dyn Context> for efficient sharing.

I could've chosen Context to be a struct instead of a dyn compatible trait, but chose to do so for the odd case where we may need another google internal implemenation of Context in future. Happy to revert this and move to a struct (like tokio or OTEL) instead.

The context itself currently simply contains deadline for all practical purposes for enabling deadline propagation. We can later add more fields as needed.

This change uses the modern file hierarchy recommended by https://doc.rust-lang.org/rust-by-example/mod/split.html which is a departure from the current crate which uses the legacy mod.rs based structure. I am happy to move to the legacy method in favor of consistency with rest of the code if needed.

sauravzg avatar Nov 26 '25 08:11 sauravzg