Add OriginContext to track data across modules (#875)

Adds OriginContext that can have data set on it about the origin and read out in any module at any time. This is designed to allow us to track "requests" (usually network requests) as it traverses through the system. The most immediate use is to track what hashing algorithm was requested for VerifyStore and verify the data using it. In the future we'll likely use this for other things, like tracking a "user" and emit what resources that user used over time and emit it to another system based on a config.
TraceMachina · May 7, 2024 · 829904e · 829904e
1 parent 185eab3
commit 829904e
Show file tree

Hide file tree

Showing 11 changed files with 554 additions and 86 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,5 +1,6 @@
 bazel-*
 target/
+.vscode/
 .cache
 .terraform*
 .config

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/nativelink-macro/src/lib.rs b/nativelink-macro/src/lib.rs
@@ -35,9 +35,12 @@ pub fn nativelink_test(attr: TokenStream, item: TokenStream) -> TokenStream {
         #[tokio::test(#attr)]
         async fn #fn_name(#fn_inputs) #fn_output {
             #[warn(clippy::disallowed_methods)]
-            {
-                #fn_block
-            }
+            ::std::sync::Arc::new(::nativelink_util::origin_context::OriginContext::new()).wrap_async(
+                ::nativelink_util::__tracing::trace_span!("test"), async move {
+                    #fn_block
+                }
+            )
+            .await
         }
     };
 

diff --git a/nativelink-service/src/health_server.rs b/nativelink-service/src/health_server.rs
@@ -14,6 +14,7 @@
 
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use futures::StreamExt;
@@ -22,7 +23,7 @@ use hyper::{Body, Request, Response, StatusCode};
 use nativelink_util::health_utils::{
     HealthRegistry, HealthStatus, HealthStatusDescription, HealthStatusReporter,
 };
-use nativelink_util::task::instrument_future;
+use nativelink_util::origin_context::OriginContext;
 use tower::Service;
 use tracing::error_span;
 
@@ -51,7 +52,8 @@ impl Service<Request<hyper::Body>> for HealthServer {
 
     fn call(&mut self, _req: Request<Body>) -> Self::Future {
         let health_registry = self.health_registry.clone();
-        Box::pin(instrument_future(
+        Box::pin(Arc::new(OriginContext::new()).wrap_async(
+            error_span!("health_server_call"),
             async move {
                 let health_status_descriptions: Vec<HealthStatusDescription> =
                     health_registry.health_status_report().collect().await;
@@ -81,7 +83,6 @@ impl Service<Request<hyper::Body>> for HealthServer {
                         .unwrap()),
                 }
             },
-            error_span!("health_server_call"),
         ))
     }
 }
diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs
@@ -40,8 +40,8 @@ use nativelink_store::filesystem_store::{
 use nativelink_util::buf_channel::make_buf_channel_pair;
 use nativelink_util::common::{fs, DigestInfo};
 use nativelink_util::evicting_map::LenEntry;
+use nativelink_util::origin_context::ContextAwareFuture;
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
-use nativelink_util::task::instrument_future;
 use nativelink_util::{background_spawn, spawn};
 use once_cell::sync::Lazy;
 use rand::{thread_rng, Rng};
@@ -51,6 +51,7 @@ use tokio::sync::Barrier;
 use tokio::time::sleep;
 use tokio_stream::wrappers::ReadDirStream;
 use tokio_stream::StreamExt;
+use tracing::Instrument;
 
 trait FileEntryHooks {
     fn on_unref<Fe: FileEntry>(_entry: &Fe) {}
@@ -167,24 +168,24 @@ impl<Hooks: FileEntryHooks + 'static + Sync + Send> Drop for TestFileEntry<Hooks
         // command that will wait for all tasks and sub spawns to complete.
         // Sadly we need to rely on `active_drop_spawns` to hit zero to ensure that
         // all tasks have completed.
-        let fut = instrument_future(
+        let fut = ContextAwareFuture::new_from_active(
             async move {
                 // Drop the FileEntryImpl in a controlled setting then wait for the
                 // `active_drop_spawns` to hit zero.
                 drop(inner);
                 while shared_context.active_drop_spawns.load(Ordering::Acquire) > 0 {
                     tokio::task::yield_now().await;
                 }
-            },
-            tracing::error_span!("test_file_entry_drop"),
+            }
+            .instrument(tracing::error_span!("test_file_entry_drop")),
         );
         #[allow(clippy::disallowed_methods)]
         let thread_handle = {
             std::thread::spawn(move || {
                 let rt = tokio::runtime::Builder::new_current_thread()
                     .build()
                     .unwrap();
-                rt.block_on(fut)
+                rt.block_on(fut);
             })
         };
         thread_handle.join().unwrap();

diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel
@@ -20,6 +20,7 @@ rust_library(
         "src/health_utils.rs",
         "src/lib.rs",
         "src/metrics_utils.rs",
+        "src/origin_context.rs",
         "src/platform_properties.rs",
         "src/proto_stream_utils.rs",
         "src/resource_info.rs",
@@ -42,6 +43,7 @@ rust_library(
         "@crates//:bytes",
         "@crates//:futures",
         "@crates//:hex",
+        "@crates//:hyper",
         "@crates//:lru",
         "@crates//:parking_lot",
         "@crates//:pin-project-lite",

diff --git a/nativelink-util/Cargo.toml b/nativelink-util/Cargo.toml
@@ -14,6 +14,7 @@ blake3 = { version = "1.5.1", features = ["mmap"] }
 bytes = "1.6.0"
 futures = "0.3.30"
 hex = "0.4.3"
+hyper = "0.14.28"
 lru = "0.12.3"
 parking_lot = "0.12.2"
 pin-project-lite = "0.2.14"

diff --git a/nativelink-util/src/lib.rs b/nativelink-util/src/lib.rs
@@ -22,6 +22,7 @@ pub mod fastcdc;
 pub mod fs;
 pub mod health_utils;
 pub mod metrics_utils;
+pub mod origin_context;
 pub mod platform_properties;
 pub mod proto_stream_utils;
 pub mod resource_info;
@@ -30,3 +31,6 @@ pub mod store_trait;
 pub mod task;
 pub mod tls_utils;
 pub mod write_counter;
+
+// Re-export tracing mostly for use in macros.
+pub use tracing as __tracing;