Add power rails energy data to Perfetto traces

Bug: 122584217
Test: Manual
Change-Id: I81aaa68e221e6e32e07d32d647cd7e79e047acc7
diff --git a/Android.bp b/Android.bp
index edd36be..d22e64c 100644
--- a/Android.bp
+++ b/Android.bp
@@ -1591,6 +1591,7 @@
   name: "perfetto_protos_perfetto_trace_power_lite_gen",
   srcs: [
     "protos/perfetto/trace/power/battery_counters.proto",
+    "protos/perfetto/trace/power/power_rails.proto",
   ],
   tools: [
     "aprotoc",
@@ -1598,6 +1599,7 @@
   cmd: "mkdir -p $(genDir)/external/perfetto/protos && $(location aprotoc) --cpp_out=$(genDir)/external/perfetto/protos --proto_path=external/perfetto/protos $(in)",
   out: [
     "external/perfetto/protos/perfetto/trace/power/battery_counters.pb.cc",
+    "external/perfetto/protos/perfetto/trace/power/power_rails.pb.cc",
   ],
 }
 
@@ -1606,6 +1608,7 @@
   name: "perfetto_protos_perfetto_trace_power_lite_gen_headers",
   srcs: [
     "protos/perfetto/trace/power/battery_counters.proto",
+    "protos/perfetto/trace/power/power_rails.proto",
   ],
   tools: [
     "aprotoc",
@@ -1613,6 +1616,7 @@
   cmd: "mkdir -p $(genDir)/external/perfetto/protos && $(location aprotoc) --cpp_out=$(genDir)/external/perfetto/protos --proto_path=external/perfetto/protos $(in)",
   out: [
     "external/perfetto/protos/perfetto/trace/power/battery_counters.pb.h",
+    "external/perfetto/protos/perfetto/trace/power/power_rails.pb.h",
   ],
   export_include_dirs: [
     "protos",
@@ -1624,6 +1628,7 @@
   name: "perfetto_protos_perfetto_trace_power_zero_gen",
   srcs: [
     "protos/perfetto/trace/power/battery_counters.proto",
+    "protos/perfetto/trace/power/power_rails.proto",
   ],
   tools: [
     "aprotoc",
@@ -1632,6 +1637,7 @@
   cmd: "mkdir -p $(genDir)/external/perfetto/protos && $(location aprotoc) --cpp_out=$(genDir)/external/perfetto/protos --proto_path=external/perfetto/protos --plugin=protoc-gen-plugin=$(location perfetto_src_protozero_protoc_plugin_protoc_plugin___gn_standalone_toolchain_gcc_like_host_) --plugin_out=wrapper_namespace=pbzero:$(genDir)/external/perfetto/protos $(in)",
   out: [
     "external/perfetto/protos/perfetto/trace/power/battery_counters.pbzero.cc",
+    "external/perfetto/protos/perfetto/trace/power/power_rails.pbzero.cc",
   ],
 }
 
@@ -1640,6 +1646,7 @@
   name: "perfetto_protos_perfetto_trace_power_zero_gen_headers",
   srcs: [
     "protos/perfetto/trace/power/battery_counters.proto",
+    "protos/perfetto/trace/power/power_rails.proto",
   ],
   tools: [
     "aprotoc",
@@ -1648,6 +1655,7 @@
   cmd: "mkdir -p $(genDir)/external/perfetto/protos && $(location aprotoc) --cpp_out=$(genDir)/external/perfetto/protos --proto_path=external/perfetto/protos --plugin=protoc-gen-plugin=$(location perfetto_src_protozero_protoc_plugin_protoc_plugin___gn_standalone_toolchain_gcc_like_host_) --plugin_out=wrapper_namespace=pbzero:$(genDir)/external/perfetto/protos $(in)",
   out: [
     "external/perfetto/protos/perfetto/trace/power/battery_counters.pbzero.h",
+    "external/perfetto/protos/perfetto/trace/power/power_rails.pbzero.h",
   ],
   export_include_dirs: [
     "protos",
diff --git a/include/perfetto/tracing/core/android_power_config.h b/include/perfetto/tracing/core/android_power_config.h
index 51242ba..9b9220e 100644
--- a/include/perfetto/tracing/core/android_power_config.h
+++ b/include/perfetto/tracing/core/android_power_config.h
@@ -78,9 +78,13 @@
     return &battery_counters_.back();
   }
 
+  bool collect_power_rails() const { return collect_power_rails_; }
+  void set_collect_power_rails(bool value) { collect_power_rails_ = value; }
+
  private:
   uint32_t battery_poll_ms_ = {};
   std::vector<BatteryCounters> battery_counters_;
+  bool collect_power_rails_ = {};
 
   // Allows to preserve unknown protobuf fields for compatibility
   // with future versions of .proto files.
diff --git a/protos/perfetto/config/perfetto_config.proto b/protos/perfetto/config/perfetto_config.proto
index 1a8d7fd..f777530 100644
--- a/protos/perfetto/config/perfetto_config.proto
+++ b/protos/perfetto/config/perfetto_config.proto
@@ -329,6 +329,9 @@
   }
   optional uint32 battery_poll_ms = 1;
   repeated BatteryCounters battery_counters = 2;
+
+  // Where available enables per-power-rail measurements.
+  optional bool collect_power_rails = 3;
 }
 
 // End of protos/perfetto/config/power/android_power_config.proto
diff --git a/protos/perfetto/config/power/android_power_config.proto b/protos/perfetto/config/power/android_power_config.proto
index 50cb4b2..c17371d 100644
--- a/protos/perfetto/config/power/android_power_config.proto
+++ b/protos/perfetto/config/power/android_power_config.proto
@@ -29,4 +29,7 @@
   }
   optional uint32 battery_poll_ms = 1;
   repeated BatteryCounters battery_counters = 2;
+
+  // Where available enables per-power-rail measurements.
+  optional bool collect_power_rails = 3;
 }
diff --git a/protos/perfetto/trace/perfetto_trace.proto b/protos/perfetto/trace/perfetto_trace.proto
index fc066e8..7b7888e 100644
--- a/protos/perfetto/trace/perfetto_trace.proto
+++ b/protos/perfetto/trace/perfetto_trace.proto
@@ -2325,6 +2325,38 @@
 
 // End of protos/perfetto/trace/power/battery_counters.proto
 
+// Begin of protos/perfetto/trace/power/power_rails.proto
+
+message PowerRails {
+
+  message RailDescriptor {
+    // Index corresponding to the rail
+    optional uint32 index = 1;
+    // Name of the rail
+    optional string rail_name = 2;
+    // Name of the subsystem to which this rail belongs
+    optional string subsys_name = 3;
+    // Hardware sampling rate
+    optional uint32 sampling_rate = 4;
+  }
+
+  // This is only emitted at the beginning of the trace.
+  repeated RailDescriptor rail_descriptor = 1;
+
+  message EnergyData {
+    // Index corresponding to RailDescriptor.index
+    optional uint32 index = 1;
+    // Time since device boot(CLOCK_BOOTTIME) in milli-seconds
+    optional uint64 timestamp_ms = 2;
+    // Accumulated energy since device boot in microwatt-seconds (uWs)
+    optional uint64 energy = 3;
+  }
+
+  repeated EnergyData energy_data = 2;
+}
+
+// End of protos/perfetto/trace/power/power_rails.proto
+
 // Begin of protos/perfetto/trace/ps/process_stats.proto
 
 // Per-process periodically sampled stats. These samples are wrapped in a
@@ -2473,7 +2505,7 @@
 // The root object emitted by Perfetto. A perfetto trace is just a stream of
 // TracePacket(s).
 //
-// Next id: 40.
+// Next id: 41.
 message TracePacket {
   // TODO(primiano): in future we should add a timestamp_clock_domain field to
   // allow mixing timestamps from different clock domains.
@@ -2496,6 +2528,7 @@
     // removed field with id 35
     ProfilePacket profile_packet = 37;
     BatteryCounters battery = 38;
+    PowerRails power_rails = 40;
     AndroidLogPacket android_log = 39;
 
     // This field is emitted at periodic intervals (~10s) and
diff --git a/protos/perfetto/trace/power/BUILD.gn b/protos/perfetto/trace/power/BUILD.gn
index 34d131b..6c6e88a 100644
--- a/protos/perfetto/trace/power/BUILD.gn
+++ b/protos/perfetto/trace/power/BUILD.gn
@@ -16,7 +16,7 @@
 import("../../../../gn/proto_library.gni")
 import("../../../../gn/protozero_library.gni")
 
-profiling_proto_names = [ "battery_counters.proto" ]
+profiling_proto_names = [ "battery_counters.proto", "power_rails.proto" ]
 
 proto_library("lite") {
   generate_python = false
diff --git a/protos/perfetto/trace/power/power_rails.proto b/protos/perfetto/trace/power/power_rails.proto
new file mode 100644
index 0000000..39486fc
--- /dev/null
+++ b/protos/perfetto/trace/power/power_rails.proto
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * 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.
+ */
+
+syntax = "proto2";
+option optimize_for = LITE_RUNTIME;
+package perfetto.protos;
+
+message PowerRails {
+
+  message RailDescriptor {
+    // Index corresponding to the rail
+    optional uint32 index = 1;
+    // Name of the rail
+    optional string rail_name = 2;
+    // Name of the subsystem to which this rail belongs
+    optional string subsys_name = 3;
+    // Hardware sampling rate
+    optional uint32 sampling_rate = 4;
+  }
+
+  // This is only emitted at the beginning of the trace.
+  repeated RailDescriptor rail_descriptor = 1;
+
+  message EnergyData {
+    // Index corresponding to RailDescriptor.index
+    optional uint32 index = 1;
+    // Time since device boot(CLOCK_BOOTTIME) in milli-seconds
+    optional uint64 timestamp_ms = 2;
+    // Accumulated energy since device boot in microwatt-seconds (uWs)
+    optional uint64 energy = 3;
+  }
+
+  repeated EnergyData energy_data = 2;
+}
diff --git a/protos/perfetto/trace/trace_packet.proto b/protos/perfetto/trace/trace_packet.proto
index 9a25982..1080211 100644
--- a/protos/perfetto/trace/trace_packet.proto
+++ b/protos/perfetto/trace/trace_packet.proto
@@ -26,6 +26,7 @@
 import "perfetto/trace/ftrace/ftrace_event_bundle.proto";
 import "perfetto/trace/ftrace/ftrace_stats.proto";
 import "perfetto/trace/power/battery_counters.proto";
+import "perfetto/trace/power/power_rails.proto";
 import "perfetto/trace/profiling/profile_packet.proto";
 import "perfetto/trace/ps/process_stats.proto";
 import "perfetto/trace/ps/process_tree.proto";
@@ -37,7 +38,7 @@
 // The root object emitted by Perfetto. A perfetto trace is just a stream of
 // TracePacket(s).
 //
-// Next id: 40.
+// Next id: 41.
 message TracePacket {
   // TODO(primiano): in future we should add a timestamp_clock_domain field to
   // allow mixing timestamps from different clock domains.
@@ -60,6 +61,7 @@
     TraceStats trace_stats = 35;
     ProfilePacket profile_packet = 37;
     BatteryCounters battery = 38;
+    PowerRails power_rails = 40;
     AndroidLogPacket android_log = 39;
 
     // This field is emitted at periodic intervals (~10s) and
diff --git a/src/perfetto_cmd/perfetto_config.descriptor.h b/src/perfetto_cmd/perfetto_config.descriptor.h
index 92d1c51..2fb0272 100644
--- a/src/perfetto_cmd/perfetto_config.descriptor.h
+++ b/src/perfetto_cmd/perfetto_config.descriptor.h
@@ -12,15 +12,15 @@
 // SHA1(tools/gen_binary_descriptors)
 // e329b1e1e964417db57f83d8ecf081e041923e78
 // SHA1(protos/perfetto/config/perfetto_config.proto)
-// 0932f50f6ec35592fef800ef2096680025a5ccd9
+// 8b62198e2e55b3dc1e6cdb31807d883cb455bcc5
 
 // This is the proto PerfettoConfig encoded as a ProtoFileDescriptor to allow
 // for reflection without libprotobuf full/non-lite protos.
 
 namespace perfetto {
 
-constexpr std::array<uint8_t, 9726> kPerfettoConfigDescriptor{
-    {0x0a, 0xfb, 0x4b, 0x0a, 0x25, 0x70, 0x65, 0x72, 0x66, 0x65, 0x74, 0x74,
+constexpr std::array<uint8_t, 9774> kPerfettoConfigDescriptor{
+    {0x0a, 0xab, 0x4c, 0x0a, 0x25, 0x70, 0x65, 0x72, 0x66, 0x65, 0x74, 0x74,
      0x6f, 0x2f, 0x63, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x2f, 0x70, 0x65, 0x72,
      0x66, 0x65, 0x74, 0x74, 0x6f, 0x5f, 0x63, 0x6f, 0x6e, 0x66, 0x69, 0x67,
      0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x0f, 0x70, 0x65, 0x72, 0x66,
@@ -170,7 +170,7 @@
      0x70, 0x6f, 0x69, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x63, 0x61,
      0x6e, 0x5f, 0x72, 0x6f, 0x6f, 0x74, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28,
      0x09, 0x52, 0x09, 0x73, 0x63, 0x61, 0x6e, 0x52, 0x6f, 0x6f, 0x74, 0x73,
-     0x22, 0xd1, 0x02, 0x0a, 0x12, 0x41, 0x6e, 0x64, 0x72, 0x6f, 0x69, 0x64,
+     0x22, 0x81, 0x03, 0x0a, 0x12, 0x41, 0x6e, 0x64, 0x72, 0x6f, 0x69, 0x64,
      0x50, 0x6f, 0x77, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12,
      0x26, 0x0a, 0x0f, 0x62, 0x61, 0x74, 0x74, 0x65, 0x72, 0x79, 0x5f, 0x70,
      0x6f, 0x6c, 0x6c, 0x5f, 0x6d, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0d,
@@ -183,7 +183,11 @@
      0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x2e, 0x42, 0x61, 0x74, 0x74, 0x65,
      0x72, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x73, 0x52, 0x0f,
      0x62, 0x61, 0x74, 0x74, 0x65, 0x72, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74,
-     0x65, 0x72, 0x73, 0x22, 0xb2, 0x01, 0x0a, 0x0f, 0x42, 0x61, 0x74, 0x74,
+     0x65, 0x72, 0x73, 0x12, 0x2e, 0x0a, 0x13, 0x63, 0x6f, 0x6c, 0x6c, 0x65,
+     0x63, 0x74, 0x5f, 0x70, 0x6f, 0x77, 0x65, 0x72, 0x5f, 0x72, 0x61, 0x69,
+     0x6c, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x11, 0x63, 0x6f,
+     0x6c, 0x6c, 0x65, 0x63, 0x74, 0x50, 0x6f, 0x77, 0x65, 0x72, 0x52, 0x61,
+     0x69, 0x6c, 0x73, 0x22, 0xb2, 0x01, 0x0a, 0x0f, 0x42, 0x61, 0x74, 0x74,
      0x65, 0x72, 0x79, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x65, 0x72, 0x73, 0x12,
      0x1f, 0x0a, 0x1b, 0x42, 0x41, 0x54, 0x54, 0x45, 0x52, 0x59, 0x5f, 0x43,
      0x4f, 0x55, 0x4e, 0x54, 0x45, 0x52, 0x5f, 0x55, 0x4e, 0x53, 0x50, 0x45,
diff --git a/src/traced/probes/power/android_power_data_source.cc b/src/traced/probes/power/android_power_data_source.cc
index 5400fe8..43b2422 100644
--- a/src/traced/probes/power/android_power_data_source.cc
+++ b/src/traced/probes/power/android_power_data_source.cc
@@ -18,6 +18,8 @@
 
 #include <dlfcn.h>
 
+#include <vector>
+
 #include "perfetto/base/logging.h"
 #include "perfetto/base/optional.h"
 #include "perfetto/base/scoped_file.h"
@@ -27,15 +29,18 @@
 #include "perfetto/tracing/core/trace_packet.h"
 #include "perfetto/tracing/core/trace_writer.h"
 #include "src/android_internal/health_hal.h"
+#include "src/android_internal/power_stats_hal.h"
 
 #include "perfetto/trace/power/battery_counters.pbzero.h"
+#include "perfetto/trace/power/power_rails.pbzero.h"
 #include "perfetto/trace/trace_packet.pbzero.h"
 
 namespace perfetto {
 
 namespace {
 constexpr uint32_t kMinPollRateMs = 250;
-}
+constexpr size_t kMaxNumRails = 32;
+}  // namespace
 
 // Dynamically loads / unloads the libperfetto_android_internal.so library which
 // allows to proxy calls to android hwbinder in in-tree builds.
@@ -55,6 +60,21 @@
       return;
     }
     get_battery_counter_ = reinterpret_cast<decltype(get_battery_counter_)>(fn);
+
+    fn = dlsym(*handle_, "GetAvailableRails");
+    if (!fn) {
+      PERFETTO_PLOG("dlsym(GetAvailableRails) failed");
+      return;
+    }
+    get_available_rails_ = reinterpret_cast<decltype(get_available_rails_)>(fn);
+
+    fn = dlsym(*handle_, "GetRailEnergyData");
+    if (!fn) {
+      PERFETTO_PLOG("dlsym(GetRailEnergyData) failed");
+      return;
+    }
+    get_rail_energy_data_ =
+        reinterpret_cast<decltype(get_rail_energy_data_)>(fn);
   }
 
   base::Optional<int64_t> GetCounter(android_internal::BatteryCounter counter) {
@@ -66,10 +86,36 @@
     return base::nullopt;
   }
 
+  std::vector<android_internal::RailDescriptor> GetRailDescriptors() {
+    if (!get_available_rails_)
+      return std::vector<android_internal::RailDescriptor>();
+
+    std::vector<android_internal::RailDescriptor> rail_descriptors(
+        kMaxNumRails);
+    size_t num_rails = rail_descriptors.size();
+    get_available_rails_(&rail_descriptors[0], &num_rails);
+    rail_descriptors.resize(num_rails);
+    return rail_descriptors;
+  }
+
+  std::vector<android_internal::RailEnergyData> GetRailEnergyData() {
+    if (!get_rail_energy_data_)
+      return std::vector<android_internal::RailEnergyData>();
+
+    std::vector<android_internal::RailEnergyData> energy_data(kMaxNumRails);
+    size_t num_rails = energy_data.size();
+    get_rail_energy_data_(&energy_data[0], &num_rails);
+    energy_data.resize(num_rails);
+    return energy_data;
+  }
+
   bool is_loaded() const { return !!handle_; }
 
  private:
   decltype(&android_internal::GetBatteryCounter) get_battery_counter_ = nullptr;
+  decltype(&android_internal::GetAvailableRails) get_available_rails_ = nullptr;
+  decltype(&android_internal::GetRailEnergyData) get_rail_energy_data_ =
+      nullptr;
   ScopedDlHandle handle_;
 };
 
@@ -81,6 +127,9 @@
     : ProbesDataSource(session_id, kTypeId),
       task_runner_(task_runner),
       poll_rate_ms_(cfg.android_power_config().battery_poll_ms()),
+      rails_collection_enabled_(
+          cfg.android_power_config().collect_power_rails()),
+      rail_descriptors_logged_(false),
       writer_(std::move(writer)),
       weak_factory_(this) {
   if (poll_rate_ms_ < kMinPollRateMs) {
@@ -132,6 +181,14 @@
       },
       poll_rate_ms_ - (now_ms % poll_rate_ms_));
 
+  WriteBatteryCounters();
+  WritePowerRailsData();
+}
+
+void AndroidPowerDataSource::WriteBatteryCounters() {
+  if (counters_enabled_.none())
+    return;
+
   auto packet = writer_->NewTracePacket();
   packet->set_timestamp(static_cast<uint64_t>(base::GetBootTimeNs().count()));
   auto* counters_proto = packet->set_battery();
@@ -168,6 +225,42 @@
   }
 }
 
+void AndroidPowerDataSource::WritePowerRailsData() {
+  if (!rails_collection_enabled_)
+    return;
+
+  auto packet = writer_->NewTracePacket();
+  packet->set_timestamp(static_cast<uint64_t>(base::GetBootTimeNs().count()));
+  auto* rails_proto = packet->set_power_rails();
+
+  if (!rail_descriptors_logged_) {
+    // We only add the rail descriptors to the first package, to avoid logging
+    // all rail names etc. on each one.
+    rail_descriptors_logged_ = true;
+    auto rail_descriptors = lib_->GetRailDescriptors();
+    if (rail_descriptors.size() == 0) {
+      // No rails to collect data for. Don't try again in the next iteration.
+      rails_collection_enabled_ = false;
+      return;
+    }
+
+    for (const auto& rail_descriptor : rail_descriptors) {
+      auto* descriptor = rails_proto->add_rail_descriptor();
+      descriptor->set_index(rail_descriptor.index);
+      descriptor->set_rail_name(rail_descriptor.rail_name);
+      descriptor->set_subsys_name(rail_descriptor.subsys_name);
+      descriptor->set_sampling_rate(rail_descriptor.sampling_rate);
+    }
+  }
+
+  for (const auto& energy_data : lib_->GetRailEnergyData()) {
+    auto* data = rails_proto->add_energy_data();
+    data->set_index(energy_data.index);
+    data->set_timestamp_ms(energy_data.timestamp);
+    data->set_energy(energy_data.energy);
+  }
+}
+
 void AndroidPowerDataSource::Flush(FlushRequestID,
                                    std::function<void()> callback) {
   writer_->Flush(callback);
diff --git a/src/traced/probes/power/android_power_data_source.h b/src/traced/probes/power/android_power_data_source.h
index 1685664..b3b45e8 100644
--- a/src/traced/probes/power/android_power_data_source.h
+++ b/src/traced/probes/power/android_power_data_source.h
@@ -54,10 +54,14 @@
   struct DynamicLibLoader;
 
   void Tick();
+  void WriteBatteryCounters();
+  void WritePowerRailsData();
 
   base::TaskRunner* const task_runner_;
   uint32_t poll_rate_ms_ = 0;
   std::bitset<8> counters_enabled_;
+  bool rails_collection_enabled_;
+  bool rail_descriptors_logged_;
   std::unique_ptr<TraceWriter> writer_;
   std::unique_ptr<DynamicLibLoader> lib_;
   base::WeakPtrFactory<AndroidPowerDataSource> weak_factory_;  // Keep last.
diff --git a/src/tracing/core/android_power_config.cc b/src/tracing/core/android_power_config.cc
index 0ce4fa4..7368469 100644
--- a/src/tracing/core/android_power_config.cc
+++ b/src/tracing/core/android_power_config.cc
@@ -56,6 +56,12 @@
     battery_counters_.back() =
         static_cast<decltype(battery_counters_)::value_type>(field);
   }
+
+  static_assert(
+      sizeof(collect_power_rails_) == sizeof(proto.collect_power_rails()),
+      "size mismatch");
+  collect_power_rails_ =
+      static_cast<decltype(collect_power_rails_)>(proto.collect_power_rails());
   unknown_fields_ = proto.unknown_fields();
 }
 
@@ -74,6 +80,13 @@
     static_assert(sizeof(it) == sizeof(proto->battery_counters(0)),
                   "size mismatch");
   }
+
+  static_assert(
+      sizeof(collect_power_rails_) == sizeof(proto->collect_power_rails()),
+      "size mismatch");
+  proto->set_collect_power_rails(
+      static_cast<decltype(proto->collect_power_rails())>(
+          collect_power_rails_));
   *(proto->mutable_unknown_fields()) = unknown_fields_;
 }
 
diff --git a/tools/gen_merged_protos b/tools/gen_merged_protos
index 05ac277..a88ba6d 100755
--- a/tools/gen_merged_protos
+++ b/tools/gen_merged_protos
@@ -65,6 +65,7 @@
   'protos/perfetto/trace/ftrace/task.proto',
   'protos/perfetto/trace/ftrace/vmscan.proto',
   'protos/perfetto/trace/power/battery_counters.proto',
+  'protos/perfetto/trace/power/power_rails.proto',
   'protos/perfetto/trace/ps/process_stats.proto',
   'protos/perfetto/trace/ps/process_tree.proto',
   'protos/perfetto/trace/sys_stats/sys_stats.proto',