rust-sdk: Add proto extension support to Rust SDK codegen (#5444)

Adds extension code generation to the protozero Rust plugin. The plugin
already collected proto extension declarations into an extensions_ map
but never emitted code for them. Now it generates pb_msg_ext!() macro
invocations that create extension traits with setter methods, matching
the existing pb_msg_ext! runtime macro.

The plugin also gains two new options for cross-crate imports:
- external_crate: the crate name to use for types not found in the
current proto set (e.g. "perfetto_sdk")
- local_files: pipe-separated list of proto files belonging to the
current crate, used to distinguish local vs external imports

This enables gpu_track_event.proto to be fully auto-generated (not
hand-written) in the GPU protos crate, with the correct cross-crate
import for TrackEvent from perfetto_sdk.

Also regenerates all Rust proto bindings, picking up:
- event_wait_ids field on GpuRenderStageEvent
- GpuCorrelation message and TrackEvent extension from
gpu_track_event.proto
- disable_extension_descriptors field on TraceConfig

Author: dreveman

contrib/rust-sdk/Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 edition = "2024"
 name = "perfetto-sdk-protos-gpu"
-version = "0.3.3"
+version = "0.3.4"
 authors = ["David Reveman <reveman@meta.com>"]
 description = "Extra protobuf bindings for GPU events"
 readme = "README.md"

contrib/rust-sdk/perfetto-protos-gpu/Cargo.toml

@@ -62,6 +62,7 @@ pb_msg!(GpuRenderStageEvent {
     render_subpass_index_mask: u64, primitive, 15,
     command_buffer_handle: u64, primitive, 12,
     name: String, primitive, 17,
+    event_wait_ids: u64, primitive, 18,
     specifications: GpuRenderStageEventSpecifications, msg, 7,
     hw_queue_id: i32, primitive, 3,
     stage_id: i32, primitive, 4,

contrib/rust-sdk/perfetto-protos-gpu/src/protos/trace/gpu/gpu_render_stage_event.pz.rs

@@ -0,0 +1,30 @@
+// Copyright (C) 2025 Rivos Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Autogenerated by the ProtoZero Rust compiler plugin.
+// Invoked by contrib/rust-sdk/tools/gen_rust_protos
+// DO NOT EDIT.
+
+use crate::pb_msg;
+use crate::pb_msg_ext;
+use perfetto_sdk::protos::trace::track_event::track_event::*;
+
+pb_msg!(GpuCorrelation {
+    render_stage_submission_event_ids: u64, primitive, 1,
+    render_stage_wait_event_ids: u64, primitive, 2,
+});
+
+pb_msg_ext!(TrackEvent {
+    gpu_correlation: GpuCorrelation, msg, 3000,
+});

contrib/rust-sdk/perfetto-protos-gpu/src/protos/trace/gpu/gpu_track_event.pz.rs

@@ -27,6 +27,10 @@ pub mod gpu_log;
 #[path = "gpu_render_stage_event.pz.rs"]
 pub mod gpu_render_stage_event;
 
+/// `gpu_track_event` protos.
+#[path = "gpu_track_event.pz.rs"]
+pub mod gpu_track_event;
+
 /// `vulkan_api_event` protos.
 #[path = "vulkan_api_event.pz.rs"]
 pub mod vulkan_api_event;

contrib/rust-sdk/perfetto-protos-gpu/src/protos/trace/gpu/mod.rs

@@ -1,7 +1,7 @@
 [package]
 edition = "2024"
 name = "perfetto-sdk"
-version = "0.3.0"
+version = "0.3.1"
 authors = ["David Reveman <reveman@meta.com>"]
 description = "Bindings for the Perfetto tracing framework"
 readme = "README.md"

contrib/rust-sdk/perfetto/Cargo.toml

@@ -221,6 +221,7 @@ pb_msg!(TraceConfigBuiltinDataSource {
     snapshot_interval_ms: u32, primitive, 6,
     prefer_suspend_clock_for_snapshot: bool, primitive, 7,
     disable_chunk_usage_histograms: bool, primitive, 8,
+    disable_extension_descriptors: bool, primitive, 9,
 });
 
 pb_msg!(TraceConfigDataSource {

contrib/rust-sdk/perfetto/src/protos/config/trace_config.pz.rs

@@ -82,6 +82,7 @@ SOURCE_FILES = [
             "protos/perfetto/trace/gpu/gpu_counter_event.proto",
             "protos/perfetto/trace/gpu/gpu_log.proto",
             "protos/perfetto/trace/gpu/gpu_render_stage_event.proto",
+            "protos/perfetto/trace/gpu/gpu_track_event.proto",
             "protos/perfetto/trace/gpu/vulkan_api_event.proto",
             "protos/perfetto/trace/gpu/vulkan_memory_event.proto",
             "protos/perfetto/trace/system_info/gpu_info.proto",
@@ -92,6 +93,7 @@ SOURCE_FILES = [
             "protos/perfetto/trace/interned_data/interned_data.proto",
             "protos/perfetto/trace/trace_packet.proto",
         ],
+        "external_crate": "perfetto_sdk",
         "path_strip_prefix": "protos/perfetto",
         "path_add_prefix": "contrib/rust-sdk/perfetto-protos-gpu/src/protos",
     },
@@ -159,12 +161,18 @@ def generate(
     protozero_rust_plugin_path,
     path_strip_prefix,
     path_add_prefix,
+    external_crate=None,
+    local_files=None,
 ):
   options = {
       "path_strip_prefix": path_strip_prefix,
       "path_add_prefix": path_add_prefix,
       "invoker": SCRIPT_PATH,
   }
+  if external_crate:
+    options["external_crate"] = external_crate
+  if local_files:
+    options["local_files"] = "|".join(local_files)
   serialized_options = ",".join(
       ["{}={}".format(name, value) for name, value in options.items()])
   subprocess.check_call(
@@ -245,6 +253,7 @@ def main():
   try:
     with tempfile.TemporaryDirectory() as tmpdirname:
       for sources in SOURCE_FILES:
+        all_files = sources["files"] + sources["custom_files"]
         for source in sources["files"]:
           generate(
               source,
@@ -253,6 +262,8 @@ def main():
               protozero_rust_plugin_path(out),
               path_strip_prefix=sources["path_strip_prefix"],
               path_add_prefix=sources["path_add_prefix"],
+              external_crate=sources.get("external_crate"),
+              local_files=all_files,
           )
 
           tmpfilename = os.path.join(tmpdirname, transform_extension(source))

contrib/rust-sdk/tools/gen_rust_protos

@@ -107,6 +107,8 @@ class GeneratorJob {
       GenerateEnumDescriptor(enumeration);
     for (const Descriptor* message : messages_)
       GenerateMessageDescriptor(message);
+    for (const auto& key_value : extensions_)
+      GenerateExtensionDescriptor(key_value.first, key_value.second);
     return error_.empty();
   }
 
@@ -119,6 +121,12 @@ class GeneratorJob {
       path_add_prefix_ = value;
     } else if (name == "invoker") {
       invoker_ = value;
+    } else if (name == "external_crate") {
+      external_crate_ = value;
+    } else if (name == "local_files") {
+      for (const auto& f : SplitString(value, "|")) {
+        local_files_.insert(std::string(f));
+      }
     } else {
       Abort(std::string() + "Unknown plugin option '" + name + "'.");
     }
@@ -299,6 +307,27 @@ class GeneratorJob {
         }
       }
     }
+
+    // Collect dependencies for extension fields (base message and field types).
+    for (const auto& key_value : extensions_) {
+      for (const FieldDescriptor* field : key_value.second) {
+        // The extended message type (e.g. TrackEvent).
+        const Descriptor* extendee = field->containing_type();
+        if (public_imports_.count(extendee->file()) == 0) {
+          referenced_messages_.insert(extendee);
+        }
+        // Message/enum types used in extension fields.
+        if (field->type() == FieldDescriptor::TYPE_MESSAGE) {
+          if (public_imports_.count(field->message_type()->file()) == 0) {
+            referenced_messages_.insert(field->message_type());
+          }
+        } else if (field->type() == FieldDescriptor::TYPE_ENUM) {
+          if (public_imports_.count(field->enum_type()->file()) == 0) {
+            referenced_enums_.insert(field->enum_type());
+          }
+        }
+      }
+    }
   }
 
   void Preprocess() {
@@ -350,6 +379,9 @@ class GeneratorJob {
     if (!messages_.empty()) {
       stub_rs_->Print("use crate::pb_msg;\n");
     }
+    if (!extensions_.empty()) {
+      stub_rs_->Print("use crate::pb_msg_ext;\n");
+    }
 
     // Print use statements for public imports, enums and messages.
     std::vector<std::string> imports;
@@ -378,8 +410,13 @@ class GeneratorJob {
       if (!path_strip_prefix_.empty()) {
         mod_path = StripPrefix(imp, path_strip_prefix_);
       }
-      stub_rs_->Print("use crate::protos$mod$::*;\n", "mod",
-                      ReplaceAll(mod_path, "/", "::"));
+      // When external_crate is set and this import is not a local file,
+      // use the external crate path instead of crate::.
+      bool is_external =
+          !external_crate_.empty() && local_files_.count(imp + ".proto") == 0;
+      std::string crate_prefix = is_external ? external_crate_ : "crate";
+      stub_rs_->Print("use $crate$::protos$mod$::*;\n", "crate", crate_prefix,
+                      "mod", ReplaceAll(mod_path, "/", "::"));
     }
   }
 
@@ -473,6 +510,21 @@ class GeneratorJob {
     }
   }
 
+  void GenerateExtensionDescriptor(
+      const std::string& /*extension_name*/,
+      const std::vector<const FieldDescriptor*>& fields) {
+    // The extended message name (e.g. "TrackEvent").
+    std::string base_name =
+        GetFullRustMessageName(fields[0]->containing_type());
+
+    stub_rs_->Print("\npb_msg_ext!($base$ {\n", "base", base_name);
+    for (const FieldDescriptor* field : fields) {
+      stub_rs_->Print("    $field$\n", "field",
+                      GetFieldDescriptorContent(field));
+    }
+    stub_rs_->Print("});\n");
+  }
+
   const FileDescriptor* const source_;
   Printer* const stub_rs_;
   std::string error_;
@@ -481,6 +533,8 @@ class GeneratorJob {
   std::string wrapper_namespace_;
   std::string path_strip_prefix_;
   std::string path_add_prefix_;
+  std::string external_crate_;
+  std::set<std::string> local_files_;
   std::string invoker_;
   std::vector<std::string> namespaces_;
   std::string full_namespace_prefix_;