Add android.power_stats data source to poll energy breakdown data

AndroidPowerStatsDataSource reports data from the
android.hardware.power.IPowerStats/default service.

It does not report power entity info as PowerEntityDescriptor yet.

Bug: 178219601

Change-Id: Id01b2dc1160d123251fef968defee5156238a900
diff --git a/Android.bp b/Android.bp
index 2e1741c..6d2115d 100644
--- a/Android.bp
+++ b/Android.bp
@@ -620,6 +620,7 @@
     "src/android_internal/atrace_hal.cc",
     "src/android_internal/health_hal.cc",
     "src/android_internal/incident_service.cc",
+    "src/android_internal/power_stats_aidl.cc",
     "src/android_internal/power_stats_hal.cc",
     "src/android_internal/statsd_logging.cc",
     "src/android_internal/tracing_service_proxy.cc",
@@ -627,6 +628,7 @@
   shared_libs: [
     "android.hardware.atrace@1.0",
     "android.hardware.health@2.0",
+    "android.hardware.power.stats-V1-cpp",
     "android.hardware.power.stats@1.0",
     "libbase",
     "libbinder",
@@ -8314,6 +8316,7 @@
   name: "perfetto_src_traced_probes_power_power",
   srcs: [
     "src/traced/probes/power/android_power_data_source.cc",
+    "src/traced/probes/power/android_power_stats_data_source.cc",
   ],
 }
 
diff --git a/BUILD b/BUILD
index 737f499..7d25fc3 100644
--- a/BUILD
+++ b/BUILD
@@ -551,6 +551,7 @@
         "src/android_internal/atrace_hal.h",
         "src/android_internal/health_hal.h",
         "src/android_internal/incident_service.h",
+        "src/android_internal/power_stats_aidl.h",
         "src/android_internal/power_stats_hal.h",
         "src/android_internal/statsd_logging.h",
         "src/android_internal/tracing_service_proxy.h",
@@ -1452,6 +1453,8 @@
     srcs = [
         "src/traced/probes/power/android_power_data_source.cc",
         "src/traced/probes/power/android_power_data_source.h",
+        "src/traced/probes/power/android_power_stats_data_source.cc",
+        "src/traced/probes/power/android_power_stats_data_source.h",
     ],
 )
 
diff --git a/protos/perfetto/trace/perfetto_trace.proto b/protos/perfetto/trace/perfetto_trace.proto
index 6afd0fb..cbeaed0 100644
--- a/protos/perfetto/trace/perfetto_trace.proto
+++ b/protos/perfetto/trace/perfetto_trace.proto
@@ -7281,6 +7281,67 @@
 
 // End of protos/perfetto/trace/perfetto/tracing_service_event.proto
 
+// Begin of protos/perfetto/common/android_energy_consumer_descriptor.proto
+
+// Energy consumer based on aidl class:
+// android.hardware.power.stats.EnergyConsumer.
+message AndroidEnergyConsumer {
+  // Unique ID of this energy consumer.  Matches the ID in a
+  // AndroidEnergyEstimationBreakdown.
+  optional int32 energy_consumer_id = 1;
+
+  // For a group of energy consumers of the same logical type, sorting by
+  // ordinal gives their physical order. Ordinals must be consecutive integers
+  // starting from 0.
+  optional int32 ordinal = 2;
+
+  // Type of this energy consumer.
+  optional string type = 3;
+
+  // Unique name of this energy consumer. Vendor/device specific. Opaque to
+  // framework.
+  optional string name = 4;
+}
+
+message AndroidEnergyConsumerDescriptor {
+  repeated AndroidEnergyConsumer energy_consumers = 1;
+}
+
+// End of protos/perfetto/common/android_energy_consumer_descriptor.proto
+
+// Begin of protos/perfetto/trace/power/android_energy_estimation_breakdown.proto
+
+// Energy data retrieve using the ODPM(On Device Power Monitor) API.
+// This proto represents the aidl class:
+// android.hardware.power.stats.EnergyConsumerResult.
+message AndroidEnergyEstimationBreakdown {
+  // The first trace packet of each session should include a energy consumer
+  // descriptor.
+  optional AndroidEnergyConsumerDescriptor energy_consumer_descriptor = 1;
+
+  // ID of the AndroidEnergyConsumer associated with this result.  Matches
+  // the energy_consumer_id in the AndroidEnergyConsumerDescriptor that
+  // should be sent at the beginning of a trace.
+  optional int32 energy_consumer_id = 2;
+
+  // Total accumulated energy since boot in microwatt-seconds (uWs)
+  optional int64 energy_uws = 3;
+
+  message EnergyUidBreakdown {
+    // Android ID/Linux UID, the accumulated energy is attributed to.
+    optional int32 uid = 1;
+
+    // Accumulated energy since boot in microwatt-seconds (uWs).
+    optional int64 energy_uws = 2;
+  }
+  // Optional attributed energy per Android ID / Linux UID for this
+  // EnergyConsumer. Sum total of attributed energy must be less than or equal
+  // to total accumulated energy.
+  repeated EnergyUidBreakdown per_uid_breakdown = 4;
+}
+
+// End of protos/perfetto/trace/power/android_energy_estimation_breakdown.proto
+
 // Begin of protos/perfetto/trace/power/battery_counters.proto
 
 message BatteryCounters {
@@ -8465,7 +8526,7 @@
 // See the [Buffers and Dataflow](/docs/concepts/buffers.md) doc for details.
 //
 // Next reserved id: 14 (up to 15).
-// Next id: 77.
+// Next id: 78.
 message TracePacket {
   // The timestamp of the TracePacket.
   // By default this timestamps refers to the trace clock (CLOCK_BOOTTIME on
@@ -8525,6 +8586,7 @@
     GpuMemTotalEvent gpu_mem_total_event = 71;
     MemoryTrackerSnapshot memory_tracker_snapshot = 73;
     FrameTimelineEvent frame_timeline_event = 76;
+    AndroidEnergyEstimationBreakdown android_energy_estimation_breakdown = 77;
 
     // Only used in profile packets.
     ProfiledFrameSymbols profiled_frame_symbols = 55;
diff --git a/protos/perfetto/trace/trace_packet.proto b/protos/perfetto/trace/trace_packet.proto
index 3324e09..0924313 100644
--- a/protos/perfetto/trace/trace_packet.proto
+++ b/protos/perfetto/trace/trace_packet.proto
@@ -41,6 +41,7 @@
 import "protos/perfetto/trace/memory_graph.proto";
 import "protos/perfetto/trace/perfetto/perfetto_metatrace.proto";
 import "protos/perfetto/trace/perfetto/tracing_service_event.proto";
+import "protos/perfetto/trace/power/android_energy_estimation_breakdown.proto";
 import "protos/perfetto/trace/power/battery_counters.proto";
 import "protos/perfetto/trace/power/power_rails.proto";
 import "protos/perfetto/trace/profiling/deobfuscation.proto";
@@ -83,7 +84,7 @@
 // See the [Buffers and Dataflow](/docs/concepts/buffers.md) doc for details.
 //
 // Next reserved id: 14 (up to 15).
-// Next id: 77.
+// Next id: 78.
 message TracePacket {
   // The timestamp of the TracePacket.
   // By default this timestamps refers to the trace clock (CLOCK_BOOTTIME on
@@ -143,6 +144,7 @@
     GpuMemTotalEvent gpu_mem_total_event = 71;
     MemoryTrackerSnapshot memory_tracker_snapshot = 73;
     FrameTimelineEvent frame_timeline_event = 76;
+    AndroidEnergyEstimationBreakdown android_energy_estimation_breakdown = 77;
 
     // Only used in profile packets.
     ProfiledFrameSymbols profiled_frame_symbols = 55;
diff --git a/src/android_internal/BUILD.gn b/src/android_internal/BUILD.gn
index 6ef8f58..1ded343 100644
--- a/src/android_internal/BUILD.gn
+++ b/src/android_internal/BUILD.gn
@@ -31,6 +31,7 @@
       "atrace_hal.cc",
       "health_hal.cc",
       "incident_service.cc",
+      "power_stats_aidl.cc",
       "power_stats_hal.cc",
       "statsd_logging.cc",
       "tracing_service_proxy.cc",
@@ -38,6 +39,7 @@
     libs = [
       "android.hardware.health@2.0",
       "android.hardware.power.stats@1.0",
+      "android.hardware.power.stats-V1-cpp",
       "android.hardware.atrace@1.0",
       "statslog_perfetto",
       "statssocket",
@@ -71,6 +73,7 @@
     "atrace_hal.h",
     "health_hal.h",
     "incident_service.h",
+    "power_stats_aidl.h",
     "power_stats_hal.h",
     "statsd_logging.h",
     "tracing_service_proxy.h",
diff --git a/src/android_internal/power_stats_aidl.cc b/src/android_internal/power_stats_aidl.cc
new file mode 100644
index 0000000..739d3ad
--- /dev/null
+++ b/src/android_internal/power_stats_aidl.cc
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+#include "src/android_internal/power_stats_aidl.h"
+
+#include "perfetto/ext/base/utils.h"
+
+#include <vector>
+
+#include <android/hardware/power/stats/IPowerStats.h>
+#include <binder/IServiceManager.h>
+
+namespace perfetto {
+namespace android_internal {
+
+using android::hardware::power::stats::EnergyConsumerResult;
+using android::hardware::power::stats::IPowerStats;
+
+namespace {
+
+android::sp<IPowerStats> g_svc;
+
+IPowerStats* MaybeGetService() {
+  if (!g_svc) {
+    static const char kInstance[] =
+        "android.hardware.power.stats.IPowerStats/default";
+    g_svc = android::checkDeclaredService<IPowerStats>(
+        android::String16(kInstance));
+  }
+  return g_svc.get();
+}
+
+void ResetService() {
+  g_svc.clear();
+}
+
+}  // namespace
+
+bool GetEnergyConsumed(EnergyEstimationBreakdown* breakdown,
+                       size_t* size_of_arr) {
+  const size_t in_array_size = *size_of_arr;
+  *size_of_arr = 0;
+  IPowerStats* svc = MaybeGetService();
+  if (svc == nullptr) {
+    return false;
+  }
+
+  std::vector<int> ids;
+  std::vector<EnergyConsumerResult> results;
+  android::binder::Status status = svc->getEnergyConsumed(ids, &results);
+
+  if (!status.isOk()) {
+    if (status.transactionError() == android::DEAD_OBJECT) {
+      // Service has died.  Reset it to attempt to acquire a new one next time.
+      ResetService();
+    }
+    return false;
+  }
+  size_t max_size = std::min(in_array_size, results.size());
+  // Iterate through all consumer ID.
+  for (const auto& result : results) {
+    if (*size_of_arr >= max_size) {
+      break;
+    }
+    auto& cur = breakdown[(*size_of_arr)++];
+    cur.energy_consumer_id = result.id;
+    cur.uid = ALL_UIDS_FOR_CONSUMER;
+    cur.energy_uws = result.energyUWs;
+
+    // Iterate through all UIDs for this consumer.
+    for (const auto& attribution : result.attribution) {
+      if (*size_of_arr >= max_size) {
+        break;
+      }
+      auto& cur = breakdown[(*size_of_arr)++];
+      cur.energy_consumer_id = result.id;
+      cur.uid = attribution.uid;
+      cur.energy_uws = attribution.energyUWs;
+    }
+  }
+
+  return true;
+}
+
+}  // namespace android_internal
+}  // namespace perfetto
diff --git a/src/android_internal/power_stats_aidl.h b/src/android_internal/power_stats_aidl.h
new file mode 100644
index 0000000..b0d33a0
--- /dev/null
+++ b/src/android_internal/power_stats_aidl.h
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+#ifndef SRC_ANDROID_INTERNAL_POWER_STATS_AIDL_H_
+#define SRC_ANDROID_INTERNAL_POWER_STATS_AIDL_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+namespace perfetto {
+namespace android_internal {
+
+const int32_t ALL_UIDS_FOR_CONSUMER = -1;
+
+struct EnergyEstimationBreakdown {
+  // Energy consumer ID.
+  int32_t energy_consumer_id;
+
+  // Process uid.  ALL_UIDS_FOR_CONSUMER represents energy for all processes
+  // for the energy_consumer_id.
+  int32_t uid;
+
+  // Energy usage in microwatts-second(µWs).
+  int64_t energy_uws;
+};
+
+extern "C" {
+
+// These functions are not thread safe unless specified otherwise.
+
+// Retrieve the energy estimation breakdown for all energy consumer.  For each
+// consumer, there will be an entry with a uid of ALL_UIDS_FOR_CONSUMER,
+// followed by the energy breakdown for each process contributing to that
+// consumer.
+bool __attribute__((visibility("default")))
+GetEnergyConsumed(EnergyEstimationBreakdown* breakdown, size_t* size_of_arr);
+
+}  // extern "C"
+
+}  // namespace android_internal
+}  // namespace perfetto
+
+#endif  // SRC_ANDROID_INTERNAL_POWER_STATS_AIDL_H_
diff --git a/src/traced/probes/power/BUILD.gn b/src/traced/probes/power/BUILD.gn
index 7b75c6a..b02301d 100644
--- a/src/traced/probes/power/BUILD.gn
+++ b/src/traced/probes/power/BUILD.gn
@@ -27,5 +27,7 @@
   sources = [
     "android_power_data_source.cc",
     "android_power_data_source.h",
+    "android_power_stats_data_source.cc",
+    "android_power_stats_data_source.h",
   ]
 }
diff --git a/src/traced/probes/power/android_power_stats_data_source.cc b/src/traced/probes/power/android_power_stats_data_source.cc
new file mode 100644
index 0000000..a3e6284
--- /dev/null
+++ b/src/traced/probes/power/android_power_stats_data_source.cc
@@ -0,0 +1,132 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+#include "src/traced/probes/power/android_power_stats_data_source.h"
+
+#include <vector>
+
+#include "perfetto/base/logging.h"
+#include "perfetto/base/task_runner.h"
+#include "perfetto/base/time.h"
+#include "perfetto/ext/base/optional.h"
+#include "perfetto/ext/base/scoped_file.h"
+#include "perfetto/ext/tracing/core/trace_packet.h"
+#include "perfetto/ext/tracing/core/trace_writer.h"
+#include "perfetto/tracing/core/data_source_config.h"
+#include "src/android_internal/health_hal.h"
+#include "src/android_internal/lazy_library_loader.h"
+#include "src/android_internal/power_stats_aidl.h"
+
+#include "protos/perfetto/trace/power/android_energy_estimation_breakdown.pbzero.h"
+#include "protos/perfetto/trace/trace_packet.pbzero.h"
+
+namespace perfetto {
+
+namespace {
+constexpr uint32_t kPollIntervalMs = 100;
+constexpr size_t kMaxNumPowerEntities = 256;
+}  // namespace
+
+// static
+const ProbesDataSource::Descriptor AndroidPowerStatsDataSource::descriptor = {
+    /*name*/ "android.power_stats",
+    /*flags*/ Descriptor::kFlagsNone,
+};
+
+// Dynamically loads the libperfetto_android_internal.so library which
+// allows to proxy calls to android hwbinder in in-tree builds.
+struct AndroidPowerStatsDataSource::DynamicLibLoader {
+  PERFETTO_LAZY_LOAD(android_internal::GetEnergyConsumed, get_energy_consumed_);
+
+  std::vector<android_internal::EnergyEstimationBreakdown> GetEnergyConsumed() {
+    if (!get_energy_consumed_)
+      return std::vector<android_internal::EnergyEstimationBreakdown>();
+
+    std::vector<android_internal::EnergyEstimationBreakdown> energy_breakdown(
+        kMaxNumPowerEntities);
+    size_t num_power_entities = energy_breakdown.size();
+    get_energy_consumed_(&energy_breakdown[0], &num_power_entities);
+    energy_breakdown.resize(num_power_entities);
+    return energy_breakdown;
+  }
+};
+
+AndroidPowerStatsDataSource::AndroidPowerStatsDataSource(
+    DataSourceConfig,
+    base::TaskRunner* task_runner,
+    TracingSessionID session_id,
+    std::unique_ptr<TraceWriter> writer)
+    : ProbesDataSource(session_id, &descriptor),
+      task_runner_(task_runner),
+      writer_(std::move(writer)),
+      weak_factory_(this) {}
+
+AndroidPowerStatsDataSource::~AndroidPowerStatsDataSource() = default;
+
+void AndroidPowerStatsDataSource::Start() {
+  lib_.reset(new DynamicLibLoader());
+  Tick();
+}
+
+void AndroidPowerStatsDataSource::Tick() {
+  // Post next task.
+  auto now_ms = base::GetWallTimeMs().count();
+  auto weak_this = weak_factory_.GetWeakPtr();
+  task_runner_->PostDelayedTask(
+      [weak_this] {
+        if (weak_this)
+          weak_this->Tick();
+      },
+      kPollIntervalMs - static_cast<uint32_t>(now_ms % kPollIntervalMs));
+  WriteEnergyEstimationBreakdown();
+}
+
+void AndroidPowerStatsDataSource::WriteEnergyEstimationBreakdown() {
+  auto energy_breakdowns = lib_->GetEnergyConsumed();
+  auto timestamp = static_cast<uint64_t>(base::GetBootTimeNs().count());
+
+  TraceWriter::TracePacketHandle packet;
+  protos::pbzero::AndroidEnergyEstimationBreakdown* energy_estimation_proto =
+      nullptr;
+  for (const auto& breakdown : energy_breakdowns) {
+    if (breakdown.uid == android_internal::ALL_UIDS_FOR_CONSUMER) {
+      // Finalize packet before calling NewTracePacket.
+      if (packet) {
+        packet->Finalize();
+      }
+      packet = writer_->NewTracePacket();
+      packet->set_timestamp(timestamp);
+      energy_estimation_proto =
+          packet->set_android_energy_estimation_breakdown();
+      energy_estimation_proto->set_energy_consumer_id(
+          breakdown.energy_consumer_id);
+      energy_estimation_proto->set_energy_uws(breakdown.energy_uws);
+    } else {
+      PERFETTO_CHECK(energy_estimation_proto != nullptr);
+      auto* uid_breakdown_proto =
+          energy_estimation_proto->add_per_uid_breakdown();
+      uid_breakdown_proto->set_uid(breakdown.uid);
+      uid_breakdown_proto->set_energy_uws(breakdown.energy_uws);
+    }
+  }
+}
+
+void AndroidPowerStatsDataSource::Flush(FlushRequestID,
+                                        std::function<void()> callback) {
+  writer_->Flush(callback);
+}
+
+}  // namespace perfetto
diff --git a/src/traced/probes/power/android_power_stats_data_source.h b/src/traced/probes/power/android_power_stats_data_source.h
new file mode 100644
index 0000000..5f6dbf6
--- /dev/null
+++ b/src/traced/probes/power/android_power_stats_data_source.h
@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+#ifndef SRC_TRACED_PROBES_POWER_ANDROID_POWER_STATS_DATA_SOURCE_H_
+#define SRC_TRACED_PROBES_POWER_ANDROID_POWER_STATS_DATA_SOURCE_H_
+
+#include <bitset>
+#include <memory>
+
+#include "perfetto/ext/base/scoped_file.h"
+#include "perfetto/ext/base/weak_ptr.h"
+#include "perfetto/tracing/core/data_source_config.h"
+#include "src/traced/probes/probes_data_source.h"
+
+namespace perfetto {
+
+class TraceWriter;
+namespace base {
+class TaskRunner;
+}
+
+class AndroidPowerStatsDataSource : public ProbesDataSource {
+ public:
+  static const ProbesDataSource::Descriptor descriptor;
+
+  AndroidPowerStatsDataSource(DataSourceConfig,
+                              base::TaskRunner*,
+                              TracingSessionID,
+                              std::unique_ptr<TraceWriter> writer);
+
+  ~AndroidPowerStatsDataSource() override;
+
+  base::WeakPtr<AndroidPowerStatsDataSource> GetWeakPtr() const;
+
+  // ProbesDataSource implementation.
+  void Start() override;
+  void Flush(FlushRequestID, std::function<void()> callback) override;
+
+ private:
+  struct DynamicLibLoader;
+
+  void Tick();
+  void WriteEnergyEstimationBreakdown();
+
+  base::TaskRunner* const task_runner_;
+  std::unique_ptr<TraceWriter> writer_;
+  std::unique_ptr<DynamicLibLoader> lib_;
+  base::WeakPtrFactory<AndroidPowerStatsDataSource>
+      weak_factory_;  // Keep last.
+};
+
+}  // namespace perfetto
+
+#endif  // SRC_TRACED_PROBES_POWER_ANDROID_POWER_STATS_DATA_SOURCE_H_
diff --git a/src/traced/probes/probes_producer.cc b/src/traced/probes/probes_producer.cc
index cfa893d..396266a 100644
--- a/src/traced/probes/probes_producer.cc
+++ b/src/traced/probes/probes_producer.cc
@@ -41,6 +41,7 @@
 #include "src/traced/probes/metatrace/metatrace_data_source.h"
 #include "src/traced/probes/packages_list/packages_list_data_source.h"
 #include "src/traced/probes/power/android_power_data_source.h"
+#include "src/traced/probes/power/android_power_stats_data_source.h"
 #include "src/traced/probes/probes_data_source.h"
 #include "src/traced/probes/ps/process_stats_data_source.h"
 #include "src/traced/probes/sys_stats/sys_stats_data_source.h"
@@ -70,6 +71,7 @@
     &InodeFileDataSource::descriptor,            //
     &SysStatsDataSource::descriptor,             //
     &AndroidPowerDataSource::descriptor,         //
+    &AndroidPowerStatsDataSource::descriptor,    //
     &AndroidLogDataSource::descriptor,           //
     &PackagesListDataSource::descriptor,         //
     &MetatraceDataSource::descriptor,            //
@@ -170,6 +172,8 @@
     data_source = CreateSysStatsDataSource(session_id, config);
   } else if (config.name() == AndroidPowerDataSource::descriptor.name) {
     data_source = CreateAndroidPowerDataSource(session_id, config);
+  } else if (config.name() == AndroidPowerStatsDataSource::descriptor.name) {
+    data_source = CreateAndroidPowerStatsDataSource(session_id, config);
   } else if (config.name() == AndroidLogDataSource::descriptor.name) {
     data_source = CreateAndroidLogDataSource(session_id, config);
   } else if (config.name() == PackagesListDataSource::descriptor.name) {
@@ -282,6 +286,16 @@
                                  endpoint_->CreateTraceWriter(buffer_id)));
 }
 
+std::unique_ptr<ProbesDataSource>
+ProbesProducer::CreateAndroidPowerStatsDataSource(
+    TracingSessionID session_id,
+    const DataSourceConfig& config) {
+  auto buffer_id = static_cast<BufferID>(config.target_buffer());
+  return std::unique_ptr<ProbesDataSource>(
+      new AndroidPowerStatsDataSource(config, task_runner_, session_id,
+                                      endpoint_->CreateTraceWriter(buffer_id)));
+}
+
 std::unique_ptr<ProbesDataSource> ProbesProducer::CreateAndroidLogDataSource(
     TracingSessionID session_id,
     const DataSourceConfig& config) {
diff --git a/src/traced/probes/probes_producer.h b/src/traced/probes/probes_producer.h
index b879a01..eb4df63 100644
--- a/src/traced/probes/probes_producer.h
+++ b/src/traced/probes/probes_producer.h
@@ -80,6 +80,9 @@
   std::unique_ptr<ProbesDataSource> CreateAndroidPowerDataSource(
       TracingSessionID session_id,
       const DataSourceConfig& config);
+  std::unique_ptr<ProbesDataSource> CreateAndroidPowerStatsDataSource(
+      TracingSessionID session_id,
+      const DataSourceConfig& config);
   std::unique_ptr<ProbesDataSource> CreateAndroidLogDataSource(
       TracingSessionID session_id,
       const DataSourceConfig& config);
diff --git a/tools/gen_android_bp b/tools/gen_android_bp
index 954e8ee..157f8f1 100755
--- a/tools/gen_android_bp
+++ b/tools/gen_android_bp
@@ -107,8 +107,10 @@
     'android.hardware.atrace@1.0',
     'android.hardware.health@2.0',
     'android.hardware.power.stats@1.0',
+    "android.hardware.power.stats-V1-cpp",
     'base',
     'binder',
+    'binder_ndk',
     'cutils',
     'hidlbase',
     'hidltransport',
@@ -872,8 +874,9 @@
     module.defaults = [defaults_module]
     for lib in target.libs:
       # Generally library names should be mangled as 'libXXX', unless they
-      # are HAL libraries (e.g., android.hardware.health@2.0).
-      android_lib = lib if '@' in lib else 'lib' + lib
+      # are HAL libraries (e.g., android.hardware.health@2.0) or AIDL c++
+      # libraries (e.g. "android.hardware.power.stats-V1-cpp")
+      android_lib = lib if '@' in lib or "-cpp" in lib else 'lib' + lib
       if lib in shared_library_allowlist:
         module.add_android_shared_lib(android_lib)
       if lib in static_library_allowlist: