Add WorkSource capability to AlarmManager.

Change-Id: I663ea3078d405f0fa667a04bdaa376ae6652e994
diff --git a/services/java/com/android/server/AlarmManagerService.java b/services/java/com/android/server/AlarmManagerService.java
index dc8fab6..212796c 100644
--- a/services/java/com/android/server/AlarmManagerService.java
+++ b/services/java/com/android/server/AlarmManagerService.java
@@ -389,19 +389,21 @@
                         ? whenElapsed
                                 : maxTriggerTime(nowElapsed, whenElapsed, a.repeatInterval);
                 setImplLocked(a.type, a.when, whenElapsed, maxElapsed,
-                        a.repeatInterval, a.operation, batch.standalone, doValidate);
+                        a.repeatInterval, a.operation, batch.standalone, doValidate, a.workSource);
             }
         }
     }
 
     private static final class InFlight extends Intent {
         final PendingIntent mPendingIntent;
+        final WorkSource mWorkSource;
         final Pair<String, ComponentName> mTarget;
         final BroadcastStats mBroadcastStats;
         final FilterStats mFilterStats;
 
-        InFlight(AlarmManagerService service, PendingIntent pendingIntent) {
+        InFlight(AlarmManagerService service, PendingIntent pendingIntent, WorkSource workSource) {
             mPendingIntent = pendingIntent;
+            mWorkSource = workSource;
             Intent intent = pendingIntent.getIntent();
             mTarget = intent != null
                     ? new Pair<String, ComponentName>(intent.getAction(), intent.getComponent())
@@ -498,12 +500,18 @@
 
     @Override
     public void set(int type, long triggerAtTime, long windowLength, long interval,
-            PendingIntent operation) {
-        set(type, triggerAtTime, windowLength, interval, operation, false);
+            PendingIntent operation, WorkSource workSource) {
+        if (workSource != null) {
+            mContext.enforceCallingPermission(
+                    android.Manifest.permission.UPDATE_DEVICE_STATS,
+                    "AlarmManager.set");
+        }
+
+        set(type, triggerAtTime, windowLength, interval, operation, false, workSource);
     }
 
     public void set(int type, long triggerAtTime, long windowLength, long interval,
-            PendingIntent operation, boolean isStandalone) {
+            PendingIntent operation, boolean isStandalone, WorkSource workSource) {
         if (operation == null) {
             Slog.w(TAG, "set/setRepeating ignored because there is no intent");
             return;
@@ -548,13 +556,14 @@
                         + " interval=" + interval + " standalone=" + isStandalone);
             }
             setImplLocked(type, triggerAtTime, triggerElapsed, maxElapsed,
-                    interval, operation, isStandalone, true);
+                    interval, operation, isStandalone, true, workSource);
         }
     }
 
     private void setImplLocked(int type, long when, long whenElapsed, long maxWhen, long interval,
-            PendingIntent operation, boolean isStandalone, boolean doValidate) {
-        Alarm a = new Alarm(type, when, whenElapsed, maxWhen, interval, operation);
+            PendingIntent operation, boolean isStandalone, boolean doValidate,
+            WorkSource workSource) {
+        Alarm a = new Alarm(type, when, whenElapsed, maxWhen, interval, operation, workSource);
         removeLocked(operation);
 
         boolean reschedule;
@@ -1039,7 +1048,8 @@
                     final long nextElapsed = alarm.whenElapsed + delta;
                     setImplLocked(alarm.type, alarm.when + delta, nextElapsed,
                             maxTriggerTime(nowELAPSED, nextElapsed, alarm.repeatInterval),
-                            alarm.repeatInterval, alarm.operation, batch.standalone, true);
+                            alarm.repeatInterval, alarm.operation, batch.standalone, true,
+                            alarm.workSource);
                 }
 
             }
@@ -1071,15 +1081,17 @@
         public long maxWhen;        // also in the elapsed time base
         public long repeatInterval;
         public PendingIntent operation;
+        public WorkSource workSource;
         
         public Alarm(int _type, long _when, long _whenElapsed, long _maxWhen,
-                long _interval, PendingIntent _op) {
+                long _interval, PendingIntent _op, WorkSource _ws) {
             type = _type;
             when = _when;
             whenElapsed = _whenElapsed;
             maxWhen = _maxWhen;
             repeatInterval = _interval;
             operation = _op;
+            workSource = _ws;
         }
 
         @Override
@@ -1193,11 +1205,11 @@
                             
                             // we have an active broadcast so stay awake.
                             if (mBroadcastRefCount == 0) {
-                                setWakelockWorkSource(alarm.operation);
+                                setWakelockWorkSource(alarm.operation, alarm.workSource);
                                 mWakeLock.acquire();
                             }
                             final InFlight inflight = new InFlight(AlarmManagerService.this,
-                                    alarm.operation);
+                                    alarm.operation, alarm.workSource);
                             mInFlight.add(inflight);
                             mBroadcastRefCount++;
 
@@ -1239,8 +1251,18 @@
         }
     }
 
-    void setWakelockWorkSource(PendingIntent pi) {
+    /**
+     * Attribute blame for a WakeLock.
+     * @param pi PendingIntent to attribute blame to if ws is null.
+     * @param ws WorkSource to attribute blame.
+     */
+    void setWakelockWorkSource(PendingIntent pi, WorkSource ws) {
         try {
+            if (ws != null) {
+                mWakeLock.setWorkSource(ws);
+                return;
+            }
+
             final int uid = ActivityManagerNative.getDefault()
                     .getUidForIntentSender(pi.getTarget());
             if (uid >= 0) {
@@ -1323,8 +1345,9 @@
             // the top of the next minute.
             final long tickEventDelay = nextTime - currentTime;
 
+            final WorkSource workSource = null; // Let system take blame for time tick events.
             set(ELAPSED_REALTIME, SystemClock.elapsedRealtime() + tickEventDelay, 0,
-                    0, mTimeTickSender, true);
+                    0, mTimeTickSender, true, workSource);
         }
 
         public void scheduleDateChangedEvent() {
@@ -1335,8 +1358,9 @@
             calendar.set(Calendar.SECOND, 0);
             calendar.set(Calendar.MILLISECOND, 0);
             calendar.add(Calendar.DAY_OF_MONTH, 1);
-      
-            set(RTC, calendar.getTimeInMillis(), 0, 0, mDateChangeSender, true);
+
+            final WorkSource workSource = null; // Let system take blame for date change events.
+            set(RTC, calendar.getTimeInMillis(), 0, 0, mDateChangeSender, true, workSource);
         }
     }
     
@@ -1452,7 +1476,8 @@
                 } else {
                     // the next of our alarms is now in flight.  reattribute the wakelock.
                     if (mInFlight.size() > 0) {
-                        setWakelockWorkSource(mInFlight.get(0).mPendingIntent);
+                        InFlight inFlight = mInFlight.get(0);
+                        setWakelockWorkSource(inFlight.mPendingIntent, inFlight.mWorkSource);
                     } else {
                         // should never happen
                         mLog.w("Alarm wakelock still held but sent queue empty");