Support chained attribution for alarm related WorkSources.

.. those passed down via the AlarmManager.set() variant that takes
a WorkSource. This required a minor re-arrangement of code in
the ActivityManager. We now treat WorkSources as opaque in the
AlarmManager and simply push them down to the AM (and eventually
to BatteryStats) where they are picked apart.

Test: BatteryStatsNoteTest, AlarmManagerTest
Bug: 62390666
Change-Id: I118f1a1d16aafa41b4f401f1a6a3ba4d2d5eca8f
diff --git a/services/core/java/com/android/server/AlarmManagerService.java b/services/core/java/com/android/server/AlarmManagerService.java
index 06cf982..472723d 100644
--- a/services/core/java/com/android/server/AlarmManagerService.java
+++ b/services/core/java/com/android/server/AlarmManagerService.java
@@ -3035,15 +3035,8 @@
                     Slog.v(TAG, "sending alarm " + alarm);
                 }
                 if (RECORD_ALARMS_IN_HISTORY) {
-                    if (alarm.workSource != null && alarm.workSource.size() > 0) {
-                        for (int wi=0; wi<alarm.workSource.size(); wi++) {
-                            ActivityManager.noteAlarmStart(
-                                    alarm.operation, alarm.workSource.get(wi), alarm.statsTag);
-                        }
-                    } else {
-                        ActivityManager.noteAlarmStart(
-                                alarm.operation, alarm.uid, alarm.statsTag);
-                    }
+                    ActivityManager.noteAlarmStart(alarm.operation, alarm.workSource, alarm.uid,
+                            alarm.statsTag);
                 }
                 mDeliveryTracker.deliverLocked(alarm, nowELAPSED, allowWhileIdle);
             } catch (RuntimeException e) {
@@ -3553,15 +3546,8 @@
                 fs.aggregateTime += nowELAPSED - fs.startTime;
             }
             if (RECORD_ALARMS_IN_HISTORY) {
-                if (inflight.mWorkSource != null && inflight.mWorkSource.size() > 0) {
-                    for (int wi=0; wi<inflight.mWorkSource.size(); wi++) {
-                        ActivityManager.noteAlarmFinish(
-                                inflight.mPendingIntent, inflight.mWorkSource.get(wi), inflight.mTag);
-                    }
-                } else {
-                    ActivityManager.noteAlarmFinish(
-                            inflight.mPendingIntent, inflight.mUid, inflight.mTag);
-                }
+                ActivityManager.noteAlarmFinish(inflight.mPendingIntent, inflight.mWorkSource,
+                        inflight.mUid, inflight.mTag);
             }
         }
 
@@ -3771,18 +3757,9 @@
                     || alarm.type == RTC_WAKEUP) {
                 bs.numWakeup++;
                 fs.numWakeup++;
-                if (alarm.workSource != null && alarm.workSource.size() > 0) {
-                    for (int wi=0; wi<alarm.workSource.size(); wi++) {
-                        final String wsName = alarm.workSource.getName(wi);
-                        ActivityManager.noteWakeupAlarm(
-                                alarm.operation, alarm.workSource.get(wi),
-                                (wsName != null) ? wsName : alarm.packageName,
-                                alarm.statsTag);
-                    }
-                } else {
-                    ActivityManager.noteWakeupAlarm(
-                            alarm.operation, alarm.uid, alarm.packageName, alarm.statsTag);
-                }
+                ActivityManager.noteWakeupAlarm(
+                        alarm.operation, alarm.workSource, alarm.uid, alarm.packageName,
+                        alarm.statsTag);
             }
         }
     }
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index d92b3b8..5936ce1 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -13858,68 +13858,100 @@
                 Context.WINDOW_SERVICE)).addView(v, lp);
     }
 
-    public void noteWakeupAlarm(IIntentSender sender, int sourceUid, String sourcePkg, String tag) {
-        if (sender != null && !(sender instanceof PendingIntentRecord)) {
-            return;
+    @Override
+    public void noteWakeupAlarm(IIntentSender sender, WorkSource workSource, int sourceUid,
+            String sourcePkg, String tag) {
+        if (workSource != null && workSource.isEmpty()) {
+            workSource = null;
         }
-        final PendingIntentRecord rec = (PendingIntentRecord)sender;
-        final BatteryStatsImpl stats = mBatteryStatsService.getActiveStatistics();
-        synchronized (stats) {
-            if (mBatteryStatsService.isOnBattery()) {
-                mBatteryStatsService.enforceCallingPermission();
-                int MY_UID = Binder.getCallingUid();
-                final int uid;
-                if (sender == null) {
-                    uid = sourceUid;
-                } else {
-                    uid = rec.uid == MY_UID ? SYSTEM_UID : rec.uid;
+
+        if (sourceUid <= 0 && workSource == null) {
+            // Try and derive a UID to attribute things to based on the caller.
+            if (sender != null) {
+                if (!(sender instanceof PendingIntentRecord)) {
+                    return;
                 }
-                BatteryStatsImpl.Uid.Pkg pkg =
-                    stats.getPackageStatsLocked(sourceUid >= 0 ? sourceUid : uid,
-                            sourcePkg != null ? sourcePkg : rec.key.packageName);
-                pkg.noteWakeupAlarmLocked(tag);
-                StatsLog.write(StatsLog.WAKEUP_ALARM_OCCURRED, sourceUid >= 0 ? sourceUid : uid,
-                        tag);
+
+                final PendingIntentRecord rec = (PendingIntentRecord) sender;
+                final int callerUid = Binder.getCallingUid();
+                sourceUid = rec.uid == callerUid ? SYSTEM_UID : rec.uid;
+            } else {
+                // TODO(narayan): Should we throw an exception in this case ? It means that we
+                // haven't been able to derive a UID to attribute things to.
+                return;
             }
         }
+
+        if (DEBUG_POWER) {
+            Slog.w(TAG, "noteWakupAlarm[ sourcePkg=" + sourcePkg + ", sourceUid=" + sourceUid
+                    + ", workSource=" + workSource + ", tag=" + tag + "]");
+        }
+
+        mBatteryStatsService.noteWakupAlarm(sourcePkg, sourceUid, workSource, tag);
     }
 
-    public void noteAlarmStart(IIntentSender sender, int sourceUid, String tag) {
-        if (sender != null && !(sender instanceof PendingIntentRecord)) {
-            return;
+    @Override
+    public void noteAlarmStart(IIntentSender sender, WorkSource workSource, int sourceUid,
+            String tag) {
+        if (workSource != null && workSource.isEmpty()) {
+            workSource = null;
         }
-        final PendingIntentRecord rec = (PendingIntentRecord)sender;
-        final BatteryStatsImpl stats = mBatteryStatsService.getActiveStatistics();
-        synchronized (stats) {
-            mBatteryStatsService.enforceCallingPermission();
-            int MY_UID = Binder.getCallingUid();
-            final int uid;
-            if (sender == null) {
-                uid = sourceUid;
+
+        if (sourceUid <= 0 && workSource == null) {
+            // Try and derive a UID to attribute things to based on the caller.
+            if (sender != null) {
+                if (!(sender instanceof PendingIntentRecord)) {
+                    return;
+                }
+
+                final PendingIntentRecord rec = (PendingIntentRecord) sender;
+                final int callerUid = Binder.getCallingUid();
+                sourceUid = rec.uid == callerUid ? SYSTEM_UID : rec.uid;
             } else {
-                uid = rec.uid == MY_UID ? SYSTEM_UID : rec.uid;
+                // TODO(narayan): Should we throw an exception in this case ? It means that we
+                // haven't been able to derive a UID to attribute things to.
+                return;
             }
-            mBatteryStatsService.noteAlarmStart(tag, sourceUid >= 0 ? sourceUid : uid);
         }
+
+        if (DEBUG_POWER) {
+            Slog.w(TAG, "noteAlarmStart[sourceUid=" + sourceUid + ", workSource=" + workSource +
+                    ", tag=" + tag + "]");
+        }
+
+        mBatteryStatsService.noteAlarmStart(tag, workSource, sourceUid);
     }
 
-    public void noteAlarmFinish(IIntentSender sender, int sourceUid, String tag) {
-        if (sender != null && !(sender instanceof PendingIntentRecord)) {
-            return;
+    @Override
+    public void noteAlarmFinish(IIntentSender sender, WorkSource workSource, int sourceUid,
+            String tag) {
+        if (workSource != null && workSource.isEmpty()) {
+            workSource = null;
         }
-        final PendingIntentRecord rec = (PendingIntentRecord)sender;
-        final BatteryStatsImpl stats = mBatteryStatsService.getActiveStatistics();
-        synchronized (stats) {
-            mBatteryStatsService.enforceCallingPermission();
-            int MY_UID = Binder.getCallingUid();
-            final int uid;
-            if (sender == null) {
-                uid = sourceUid;
+
+        if (sourceUid <= 0 && workSource == null) {
+            // Try and derive a UID to attribute things to based on the caller.
+            if (sender != null) {
+                if (!(sender instanceof PendingIntentRecord)) {
+                    return;
+                }
+
+                final PendingIntentRecord rec = (PendingIntentRecord) sender;
+                final int callerUid = Binder.getCallingUid();
+                sourceUid = rec.uid == callerUid ? SYSTEM_UID : rec.uid;
             } else {
-                uid = rec.uid == MY_UID ? SYSTEM_UID : rec.uid;
+                // TODO(narayan): Should we throw an exception in this case ? It means that we
+                // haven't been able to derive a UID to attribute things to.
+                return;
             }
-            mBatteryStatsService.noteAlarmFinish(tag, sourceUid >= 0 ? sourceUid : uid);
         }
+
+        if (DEBUG_POWER) {
+            Slog.w(TAG, "noteAlarmFinish[sourceUid=" + sourceUid + ", workSource=" + workSource +
+                    ", tag=" + tag + "]");
+        }
+
+        mBatteryStatsService.noteAlarmFinish(tag, workSource, sourceUid);
     }
 
     public boolean killPids(int[] pids, String pReason, boolean secure) {
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 35318f6..430320a 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -38,6 +38,7 @@
 import android.os.UserHandle;
 import android.os.UserManagerInternal;
 import android.os.WorkSource;
+import android.os.WorkSource.WorkChain;
 import android.os.connectivity.CellularBatteryStats;
 import android.os.health.HealthStatsParceler;
 import android.os.health.HealthStatsWriter;
@@ -66,6 +67,7 @@
 import java.nio.charset.CharsetDecoder;
 import java.nio.charset.CodingErrorAction;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
@@ -446,17 +448,24 @@
         }
     }
 
-    public void noteAlarmStart(String name, int uid) {
+    public void noteWakupAlarm(String name, int uid, WorkSource workSource, String tag) {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteAlarmStartLocked(name, uid);
+            mStats.noteWakupAlarmLocked(name, uid, workSource, tag);
         }
     }
 
-    public void noteAlarmFinish(String name, int uid) {
+    public void noteAlarmStart(String name, WorkSource workSource, int uid) {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteAlarmFinishLocked(name, uid);
+            mStats.noteAlarmStartLocked(name, workSource, uid);
+        }
+    }
+
+    public void noteAlarmFinish(String name, WorkSource workSource, int uid) {
+        enforceCallingPermission();
+        synchronized (mStats) {
+            mStats.noteAlarmFinishLocked(name, workSource, uid);
         }
     }