rules: Hook up auto-uploads for traces

Summary: On the two trace subtargets, `run` them to produce an upload Reviewed By: emersonford Differential Revision: D70435505 fbshipit-source-id: c909bb48ca1fc4e71def9f5fa46c45e88ed0cdd7
facebook · Mar 1, 2025 · 8ead5b1 · 8ead5b1
1 parent 6d2f458
commit 8ead5b1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/rust/profile.bzl b/rust/profile.bzl
@@ -8,6 +8,13 @@
 load(":context.bzl", "CompileContext")  # @unused: Used as type
 load(":outputs.bzl", "RustcOutput")
 
+def _make_trace_providers(compile_ctx: CompileContext, a: Artifact) -> list[Provider]:
+    p = [DefaultInfo(default_output = a)]
+    make_trace_upload = compile_ctx.toolchain_info.make_trace_upload
+    if make_trace_upload != None:
+        p.append(make_trace_upload(a))
+    return p
+
 def _analyze_llvm_lines(
         ctx: AnalysisContext,
         compile_ctx: CompileContext,
@@ -28,12 +35,10 @@ def _analyze_llvm_lines(
     )
     return [DefaultInfo(default_output = llvm_lines)]
 
-def _llvm_time_trace(llvm_time_trace: RustcOutput) -> list[Provider]:
-    return [
-        DefaultInfo(
-            default_output = llvm_time_trace.profile_output,
-        ),
-    ]
+def _llvm_time_trace(
+        compile_ctx: CompileContext,
+        llvm_time_trace: RustcOutput) -> list[Provider]:
+    return _make_trace_providers(compile_ctx, llvm_time_trace.profile_output)
 
 def _self_profile(
         ctx: AnalysisContext,
@@ -68,7 +73,7 @@ def _self_profile(
             ),
             category = "run_crox",
         )
-        sub_targets["trace"] = [DefaultInfo(default_output = proftrace)]
+        sub_targets["trace"] = _make_trace_providers(compile_ctx, proftrace)
     return [DefaultInfo(sub_targets = sub_targets)]
 
 def make_profile_providers(
@@ -83,7 +88,7 @@ def make_profile_providers(
     if llvm_lines != None:
         sub_targets["llvm_lines"] = llvm_lines
 
-    sub_targets["llvm_passes"] = _llvm_time_trace(llvm_time_trace)
+    sub_targets["llvm_passes"] = _llvm_time_trace(compile_ctx, llvm_time_trace)
     sub_targets["rustc_stages"] = _self_profile(ctx, compile_ctx, self_profile)
 
     return [DefaultInfo(sub_targets = sub_targets)]
diff --git a/rust/rust_toolchain.bzl b/rust/rust_toolchain.bzl
@@ -130,6 +130,8 @@ rust_toolchain_attrs = {
     # The `crox` binary from measure-me, if present, used to generate a
     # self-profile trace subtarget
     "measureme_crox": provider_field(RunInfo | None, default = None),
+    # Constructs an upload command for the given chrome trace
+    "make_trace_upload": provider_field(typing.Callable[[Artifact], RunInfo] | None, default = None),
 }
 
 RustToolchainInfo = provider(fields = rust_toolchain_attrs)