Add proto output of battery usage stats

Bug: 245832387
Test: adb shell dumpsys batterystats --usage --proto | \
      gqui from rawproto:- proto $ANDROID_BUILD_TOP/frameworks/base/core/proto/android/os/batteryusagestats.proto:android.os.BatteryUsageStatsAtomsProto --protoprint_annotations=no
Change-Id: I297f09892accf0d2adef41d4cbd384a505c9c189
diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index 7105b1a..0c5f778 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -33,6 +33,7 @@
 import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.Closeable;
+import java.io.FileDescriptor;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
@@ -448,6 +449,16 @@
         return proto.getBytes();
     }
 
+    /**
+     * Writes contents in a binary protobuffer format, using
+     * the android.os.BatteryUsageStatsAtomsProto proto.
+     */
+    public void dumpToProto(FileDescriptor fd) {
+        final ProtoOutputStream proto = new ProtoOutputStream(fd);
+        writeStatsProto(proto, /* max size */ Integer.MAX_VALUE);
+        proto.flush();
+    }
+
     @NonNull
     private void writeStatsProto(ProtoOutputStream proto, int maxRawSize) {
         final AggregateBatteryConsumer deviceBatteryConsumer = getAggregateBatteryConsumer(
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index cc2b693..a41a311 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -44,6 +44,7 @@
 import android.net.INetworkManagementEventObserver;
 import android.net.Network;
 import android.net.NetworkCapabilities;
+import android.os.BatteryConsumer;
 import android.os.BatteryManagerInternal;
 import android.os.BatteryStats;
 import android.os.BatteryStatsInternal;
@@ -2292,6 +2293,10 @@
         pw.println("  --settings: dump the settings key/values related to batterystats");
         pw.println("  --cpu: dump cpu stats for debugging purpose");
         pw.println("  --power-profile: dump the power profile constants");
+        pw.println("  --usage: write battery usage stats. Optional arguments:");
+        pw.println("     --proto: output as a binary protobuffer");
+        pw.println("     --model power-profile: use the power profile model"
+                + " even if measured energy is available");
         pw.println("  <package.name>: optional name of package to filter output by.");
         pw.println("  -h: print this help text.");
         pw.println("Battery stats (batterystats) commands:");
@@ -2335,6 +2340,31 @@
         }
     }
 
+    private void dumpUsageStatsToProto(FileDescriptor fd, PrintWriter pw, int model,
+            boolean proto) {
+        awaitCompletion();
+        syncStats("dump", BatteryExternalStatsWorker.UPDATE_ALL);
+
+        BatteryUsageStatsQuery.Builder builder = new BatteryUsageStatsQuery.Builder()
+                .setMaxStatsAgeMs(0)
+                .includeProcessStateData()
+                .includePowerModels();
+        if (model == BatteryConsumer.POWER_MODEL_POWER_PROFILE) {
+            builder.powerProfileModeledOnly();
+        }
+        BatteryUsageStatsQuery query = builder.build();
+        synchronized (mStats) {
+            mStats.prepareForDumpLocked();
+            BatteryUsageStats batteryUsageStats =
+                    mBatteryUsageStatsProvider.getBatteryUsageStats(query);
+            if (proto) {
+                batteryUsageStats.dumpToProto(fd);
+            } else {
+                batteryUsageStats.dump(pw, "");
+            }
+        }
+    }
+
     private int doEnableOrDisable(PrintWriter pw, int i, String[] args, boolean enable) {
         i++;
         if (i >= args.length) {
@@ -2488,6 +2518,35 @@
                 } else if ("--power-profile".equals(arg)) {
                     dumpPowerProfile(pw);
                     return;
+                } else if ("--usage".equals(arg)) {
+                    int model = BatteryConsumer.POWER_MODEL_UNDEFINED;
+                    boolean proto = false;
+                    for (int j = i + 1; j < args.length; j++) {
+                        switch (args[j]) {
+                            case "--proto":
+                                proto = true;
+                                break;
+                            case "--model": {
+                                if (j + 1 < args.length) {
+                                    j++;
+                                    if ("power-profile".equals(args[j])) {
+                                        model = BatteryConsumer.POWER_MODEL_POWER_PROFILE;
+                                    } else {
+                                        pw.println("Unknown power model: " + args[j]);
+                                        dumpHelp(pw);
+                                        return;
+                                    }
+                                } else {
+                                    pw.println("--model without a value");
+                                    dumpHelp(pw);
+                                    return;
+                                }
+                                break;
+                            }
+                        }
+                    }
+                    dumpUsageStatsToProto(fd, pw, model, proto);
+                    return;
                 } else if ("-a".equals(arg)) {
                     flags |= BatteryStats.DUMP_VERBOSE;
                 } else if (arg.length() > 0 && arg.charAt(0) == '-'){