Merge "Add WorkSource capability to AlarmManager."
diff --git a/core/java/android/app/AlarmManager.java b/core/java/android/app/AlarmManager.java
index b73fbff..1a56826 100644
--- a/core/java/android/app/AlarmManager.java
+++ b/core/java/android/app/AlarmManager.java
@@ -20,6 +20,7 @@
 import android.content.Intent;
 import android.os.Build;
 import android.os.RemoteException;
+import android.os.WorkSource;
 
 /**
  * This class provides access to the system alarm services.  These allow you
@@ -155,7 +156,7 @@
      * @see #RTC_WAKEUP
      */
     public void set(int type, long triggerAtMillis, PendingIntent operation) {
-        setImpl(type, triggerAtMillis, legacyExactLength(), 0, operation);
+        setImpl(type, triggerAtMillis, legacyExactLength(), 0, operation, null);
     }
 
     /**
@@ -209,7 +210,7 @@
      */
     public void setRepeating(int type, long triggerAtMillis,
             long intervalMillis, PendingIntent operation) {
-        setImpl(type, triggerAtMillis, legacyExactLength(), intervalMillis, operation);
+        setImpl(type, triggerAtMillis, legacyExactLength(), intervalMillis, operation, null);
     }
 
     /**
@@ -244,7 +245,7 @@
      */
     public void setWindow(int type, long windowStartMillis, long windowLengthMillis,
             PendingIntent operation) {
-        setImpl(type, windowStartMillis, windowLengthMillis, 0, operation);
+        setImpl(type, windowStartMillis, windowLengthMillis, 0, operation, null);
     }
 
     /**
@@ -252,11 +253,17 @@
      * to the precise time specified.
      */
     public void setExact(int type, long triggerAtMillis, PendingIntent operation) {
-        setImpl(type, triggerAtMillis, WINDOW_EXACT, 0, operation);
+        setImpl(type, triggerAtMillis, WINDOW_EXACT, 0, operation, null);
+    }
+
+    /** @hide */
+    public void set(int type, long triggerAtMillis, long windowMillis, long intervalMillis,
+            PendingIntent operation, WorkSource workSource) {
+        setImpl(type, triggerAtMillis, windowMillis, intervalMillis, operation, workSource);
     }
 
     private void setImpl(int type, long triggerAtMillis, long windowMillis, long intervalMillis,
-            PendingIntent operation) {
+            PendingIntent operation, WorkSource workSource) {
         if (triggerAtMillis < 0) {
             /* NOTYET
             if (mAlwaysExact) {
@@ -267,8 +274,10 @@
             */
             triggerAtMillis = 0;
         }
+
         try {
-            mService.set(type, triggerAtMillis, windowMillis, intervalMillis, operation);
+            mService.set(type, triggerAtMillis, windowMillis, intervalMillis, operation,
+                    workSource);
         } catch (RemoteException ex) {
         }
     }
@@ -361,7 +370,7 @@
     @Deprecated
     public void setInexactRepeating(int type, long triggerAtMillis,
             long intervalMillis, PendingIntent operation) {
-        setImpl(type, triggerAtMillis, WINDOW_HEURISTIC, intervalMillis, operation);
+        setImpl(type, triggerAtMillis, WINDOW_HEURISTIC, intervalMillis, operation, null);
     }
     
     /**
diff --git a/core/java/android/app/IAlarmManager.aidl b/core/java/android/app/IAlarmManager.aidl
index 0a49ddf..8476609 100644
--- a/core/java/android/app/IAlarmManager.aidl
+++ b/core/java/android/app/IAlarmManager.aidl
@@ -17,6 +17,7 @@
 package android.app;
 
 import android.app.PendingIntent;
+import android.os.WorkSource;
 
 /**
  * System private API for talking with the alarm manager service.
@@ -26,7 +27,7 @@
 interface IAlarmManager {
 	/** windowLength == 0 means exact; windowLength < 0 means the let the OS decide */
     void set(int type, long triggerAtTime, long windowLength,
-            long interval, in PendingIntent operation);
+            long interval, in PendingIntent operation, in WorkSource workSource);
     void setTime(long millis);
     void setTimeZone(String zone);
     void remove(in PendingIntent operation);
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");
diff --git a/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java b/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java
index aca77b8..5fa046aa 100644
--- a/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java
@@ -878,7 +878,8 @@
         mAlarmManager.remove(isA(PendingIntent.class));
         expectLastCall().anyTimes();
 
-        mAlarmManager.set(eq(AlarmManager.ELAPSED_REALTIME), anyLong(), anyLong(), anyLong(), isA(PendingIntent.class));
+        mAlarmManager.set(eq(AlarmManager.ELAPSED_REALTIME), anyLong(), anyLong(), anyLong(),
+                isA(PendingIntent.class), isA(WorkSource.class));
         expectLastCall().atLeastOnce();
 
         mNetManager.setGlobalAlert(anyLong());