Merge "Chained attribution for GnssLocationProvider."
diff --git a/core/java/android/os/WorkSource.java b/core/java/android/os/WorkSource.java
index 81d7506..d592000 100644
--- a/core/java/android/os/WorkSource.java
+++ b/core/java/android/os/WorkSource.java
@@ -493,6 +493,29 @@
         return mChains;
     }
 
+    /**
+     * DO NOT USE: Hacky API provided solely for {@code GnssLocationProvider}. See
+     * {@code setReturningDiffs} as well.
+     *
+     * @hide
+     */
+    public void transferWorkChains(WorkSource other) {
+        if (mChains != null) {
+            mChains.clear();
+        }
+
+        if (other.mChains == null || other.mChains.isEmpty()) {
+            return;
+        }
+
+        if (mChains == null) {
+            mChains = new ArrayList<>(4);
+        }
+
+        mChains.addAll(other.mChains);
+        other.mChains.clear();
+    }
+
     private boolean removeUids(WorkSource other) {
         int N1 = mNum;
         final int[] uids1 = mUids;
@@ -880,6 +903,13 @@
             return mUids[0];
         }
 
+        /**
+         * Return the tag associated with the attribution UID. See (@link #getAttributionUid}.
+         */
+        public String getAttributionTag() {
+            return mTags[0];
+        }
+
         // TODO: The following three trivial getters are purely for testing and will be removed
         // once we have higher level logic in place, e.g for serializing this WorkChain to a proto,
         // diffing it etc.
diff --git a/core/java/com/android/internal/app/IBatteryStats.aidl b/core/java/com/android/internal/app/IBatteryStats.aidl
index d3e807d..514ff76 100644
--- a/core/java/com/android/internal/app/IBatteryStats.aidl
+++ b/core/java/com/android/internal/app/IBatteryStats.aidl
@@ -91,8 +91,7 @@
 
     void noteVibratorOn(int uid, long durationMillis);
     void noteVibratorOff(int uid);
-    void noteStartGps(int uid);
-    void noteStopGps(int uid);
+    void noteGpsChanged(in WorkSource oldSource, in WorkSource newSource);
     void noteGpsSignalQuality(int signalLevel);
     void noteScreenState(int state);
     void noteScreenBrightness(int brightness);
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 51f51c2..ee3bec8 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -4587,8 +4587,35 @@
 
     int mGpsNesting;
 
-    public void noteStartGpsLocked(int uid) {
-        uid = mapUid(uid);
+    public void noteGpsChangedLocked(WorkSource oldWs, WorkSource newWs) {
+        for (int i = 0; i < newWs.size(); ++i) {
+            noteStartGpsLocked(newWs.get(i), null);
+        }
+
+        for (int i = 0; i < oldWs.size(); ++i) {
+            noteStopGpsLocked((oldWs.get(i)), null);
+        }
+
+        List<WorkChain>[] wcs = WorkSource.diffChains(oldWs, newWs);
+        if (wcs != null) {
+            if (wcs[0] != null) {
+                final List<WorkChain> newChains = wcs[0];
+                for (int i = 0; i < newChains.size(); ++i) {
+                    noteStartGpsLocked(-1, newChains.get(i));
+                }
+            }
+
+            if (wcs[1] != null) {
+                final List<WorkChain> goneChains = wcs[1];
+                for (int i = 0; i < goneChains.size(); ++i) {
+                    noteStopGpsLocked(-1, goneChains.get(i));
+                }
+            }
+        }
+    }
+
+    private void noteStartGpsLocked(int uid, WorkChain workChain) {
+        uid = getAttributionUid(uid, workChain);
         final long elapsedRealtime = mClocks.elapsedRealtime();
         final long uptime = mClocks.uptimeMillis();
         if (mGpsNesting == 0) {
@@ -4598,11 +4625,19 @@
             addHistoryRecordLocked(elapsedRealtime, uptime);
         }
         mGpsNesting++;
+
+        if (workChain == null) {
+            StatsLog.write_non_chained(StatsLog.GPS_SCAN_STATE_CHANGED, uid, null, 1);
+        } else {
+            StatsLog.write(StatsLog.GPS_SCAN_STATE_CHANGED,
+                    workChain.getUids(), workChain.getTags(), 1);
+        }
+
         getUidStatsLocked(uid).noteStartGps(elapsedRealtime);
     }
 
-    public void noteStopGpsLocked(int uid) {
-        uid = mapUid(uid);
+    private void noteStopGpsLocked(int uid, WorkChain workChain) {
+        uid = getAttributionUid(uid, workChain);
         final long elapsedRealtime = mClocks.elapsedRealtime();
         final long uptime = mClocks.uptimeMillis();
         mGpsNesting--;
@@ -4614,6 +4649,14 @@
             stopAllGpsSignalQualityTimersLocked(-1);
             mGpsSignalQualityBin = -1;
         }
+
+        if (workChain == null) {
+            StatsLog.write_non_chained(StatsLog.GPS_SCAN_STATE_CHANGED, uid, null, 0);
+        } else {
+            StatsLog.write(StatsLog.GPS_SCAN_STATE_CHANGED, workChain.getUids(),
+                    workChain.getTags(), 0);
+        }
+
         getUidStatsLocked(uid).noteStopGps(elapsedRealtime);
     }
 
@@ -9842,9 +9885,7 @@
         public void noteStartSensor(int sensor, long elapsedRealtimeMs) {
             DualTimer t = getSensorTimerLocked(sensor, /* create= */ true);
             t.startRunningLocked(elapsedRealtimeMs);
-            if (sensor == Sensor.GPS) {
-                StatsLog.write_non_chained(StatsLog.GPS_SCAN_STATE_CHANGED, getUid(), null, 1);
-            } else {
+            if (sensor != Sensor.GPS) {
                 StatsLog.write_non_chained(StatsLog.SENSOR_STATE_CHANGED, getUid(), null, sensor,
                         1);
             }
@@ -9855,14 +9896,9 @@
             DualTimer t = getSensorTimerLocked(sensor, false);
             if (t != null) {
                 t.stopRunningLocked(elapsedRealtimeMs);
-                if (!t.isRunningLocked()) { // only tell statsd if truly stopped
-                    if (sensor == Sensor.GPS) {
-                        StatsLog.write_non_chained(StatsLog.GPS_SCAN_STATE_CHANGED, getUid(), null,
-                                0);
-                    } else {
-                        StatsLog.write_non_chained(StatsLog.SENSOR_STATE_CHANGED, getUid(), null,
-                                sensor, 0);
-                    }
+                if (sensor != Sensor.GPS) {
+                    StatsLog.write_non_chained(StatsLog.SENSOR_STATE_CHANGED, getUid(), null,
+                             sensor, 0);
                 }
             }
         }
diff --git a/core/tests/coretests/src/android/os/WorkSourceTest.java b/core/tests/coretests/src/android/os/WorkSourceTest.java
index a427a2f..952a64d 100644
--- a/core/tests/coretests/src/android/os/WorkSourceTest.java
+++ b/core/tests/coretests/src/android/os/WorkSourceTest.java
@@ -349,6 +349,15 @@
         assertEquals(100, wc.getAttributionUid());
     }
 
+    public void testGetAttributionTag() {
+        WorkSource ws1 = new WorkSource();
+        WorkChain wc = ws1.createWorkChain();
+        wc.addNode(100, "tag");
+        assertEquals("tag", wc.getAttributionTag());
+        wc.addNode(200, "tag2");
+        assertEquals("tag", wc.getAttributionTag());
+    }
+
     public void testRemove_fromChainedWorkSource() {
         WorkSource ws1 = new WorkSource();
         ws1.createWorkChain().addNode(50, "foo");
@@ -368,4 +377,25 @@
         assertEquals(1, ws1.getWorkChains().size());
         assertEquals(75, ws1.getWorkChains().get(0).getAttributionUid());
     }
+
+    public void testTransferWorkChains() {
+        WorkSource ws1 = new WorkSource();
+        WorkChain wc1 = ws1.createWorkChain().addNode(100, "tag");
+        WorkChain wc2 = ws1.createWorkChain().addNode(200, "tag2");
+
+        WorkSource ws2 = new WorkSource();
+        ws2.transferWorkChains(ws1);
+
+        assertEquals(0, ws1.getWorkChains().size());
+        assertEquals(2, ws2.getWorkChains().size());
+        assertSame(wc1, ws2.getWorkChains().get(0));
+        assertSame(wc2, ws2.getWorkChains().get(1));
+
+        ws1.clear();
+        ws1.createWorkChain().addNode(300, "tag3");
+        ws1.transferWorkChains(ws2);
+        assertEquals(0, ws2.getWorkChains().size());
+        assertSame(wc1, ws1.getWorkChains().get(0));
+        assertSame(wc2, ws1.getWorkChains().get(1));
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
index a55563a..82ac9da 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -23,10 +23,12 @@
 import android.os.BatteryManager;
 import android.os.BatteryStats;
 import android.os.BatteryStats.HistoryItem;
+import android.os.BatteryStats.Uid.Sensor;
 import android.os.WorkSource;
 import android.support.test.filters.SmallTest;
 import android.view.Display;
 
+import com.android.internal.os.BatteryStatsImpl.DualTimer;
 import com.android.internal.os.BatteryStatsImpl.Uid;
 import junit.framework.TestCase;
 
@@ -446,4 +448,52 @@
         pkg = bi.getPackageStatsLocked(100, "com.foo.baz_alternate");
         assertEquals(0, pkg.getWakeupAlarmStats().size());
     }
+
+    @SmallTest
+    public void testNoteGpsChanged() {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        bi.setRecordAllHistoryLocked(true);
+        bi.forceRecordAllHistory();
+        bi.mForceOnBattery = true;
+
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
+
+        WorkSource ws = new WorkSource();
+        ws.add(UID);
+
+        bi.noteGpsChangedLocked(new WorkSource(), ws);
+        DualTimer t = bi.getUidStatsLocked(UID).getSensorTimerLocked(Sensor.GPS, false);
+        assertNotNull(t);
+        assertTrue(t.isRunningLocked());
+
+        bi.noteGpsChangedLocked(ws, new WorkSource());
+        t = bi.getUidStatsLocked(UID).getSensorTimerLocked(Sensor.GPS, false);
+        assertFalse(t.isRunningLocked());
+    }
+
+    @SmallTest
+    public void testNoteGpsChanged_workSource() {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        bi.setRecordAllHistoryLocked(true);
+        bi.forceRecordAllHistory();
+        bi.mForceOnBattery = true;
+
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
+
+        WorkSource ws = new WorkSource();
+        ws.createWorkChain().addNode(UID, "com.foo");
+
+        bi.noteGpsChangedLocked(new WorkSource(), ws);
+        DualTimer t = bi.getUidStatsLocked(UID).getSensorTimerLocked(Sensor.GPS, false);
+        assertNotNull(t);
+        assertTrue(t.isRunningLocked());
+
+        bi.noteGpsChangedLocked(ws, new WorkSource());
+        t = bi.getUidStatsLocked(UID).getSensorTimerLocked(Sensor.GPS, false);
+        assertFalse(t.isRunningLocked());
+    }
 }
diff --git a/services/core/java/com/android/server/LocationManagerService.java b/services/core/java/com/android/server/LocationManagerService.java
index 9aa588f..bd93b179 100644
--- a/services/core/java/com/android/server/LocationManagerService.java
+++ b/services/core/java/com/android/server/LocationManagerService.java
@@ -67,6 +67,7 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.os.WorkSource;
+import android.os.WorkSource.WorkChain;
 import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.ArrayMap;
@@ -829,7 +830,7 @@
             }
             mAllowedResolutionLevel = getAllowedResolutionLevel(pid, uid);
             mIdentity = new Identity(uid, pid, packageName);
-            if (workSource != null && workSource.size() <= 0) {
+            if (workSource != null && workSource.isEmpty()) {
                 workSource = null;
             }
             mWorkSource = workSource;
@@ -1814,13 +1815,11 @@
 
                         if (locationRequest.getInterval() <= thresholdInterval) {
                             if (record.mReceiver.mWorkSource != null
-                                    && record.mReceiver.mWorkSource.size() > 0
-                                    && record.mReceiver.mWorkSource.getName(0) != null) {
-                                // Assign blame to another work source.
-                                // Can only assign blame if the WorkSource contains names.
+                                    && isValidWorkSource(record.mReceiver.mWorkSource)) {
                                 worksource.add(record.mReceiver.mWorkSource);
                             } else {
-                                // Assign blame to caller.
+                                // Assign blame to caller if there's no WorkSource associated with
+                                // the request or if it's invalid.
                                 worksource.add(
                                         record.mReceiver.mIdentity.mUid,
                                         record.mReceiver.mIdentity.mPackageName);
@@ -1835,6 +1834,23 @@
         p.setRequest(providerRequest, worksource);
     }
 
+    /**
+     * Whether a given {@code WorkSource} associated with a Location request is valid.
+     */
+    private static boolean isValidWorkSource(WorkSource workSource) {
+        if (workSource.size() > 0) {
+            // If the WorkSource has one or more non-chained UIDs, make sure they're accompanied
+            // by tags.
+            return workSource.getName(0) != null;
+        } else {
+            // For now, make sure callers have supplied an attribution tag for use with
+            // AppOpsManager. This might be relaxed in the future.
+            final ArrayList<WorkChain> workChains = workSource.getWorkChains();
+            return workChains != null && !workChains.isEmpty() &&
+                    workChains.get(0).getAttributionTag() != null;
+        }
+    }
+
     @Override
     public String[] getBackgroundThrottlingWhitelist() {
         synchronized (mLock) {
@@ -2057,7 +2073,7 @@
         checkResolutionLevelIsSufficientForProviderUse(allowedResolutionLevel,
                 request.getProvider());
         WorkSource workSource = request.getWorkSource();
-        if (workSource != null && workSource.size() > 0) {
+        if (workSource != null && !workSource.isEmpty()) {
             checkDeviceStatsAllowed();
         }
         boolean hideFromAppOps = request.getHideFromAppOps();
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 04b49ba..6b380f1 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -582,17 +582,11 @@
         }
     }
 
-    public void noteStartGps(int uid) {
+    @Override
+    public void noteGpsChanged(WorkSource oldWs, WorkSource newWs) {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteStartGpsLocked(uid);
-        }
-    }
-
-    public void noteStopGps(int uid) {
-        enforceCallingPermission();
-        synchronized (mStats) {
-            mStats.noteStopGpsLocked(uid);
+            mStats.noteGpsChangedLocked(oldWs, newWs);
         }
     }
 
diff --git a/services/core/java/com/android/server/location/GnssLocationProvider.java b/services/core/java/com/android/server/location/GnssLocationProvider.java
index 48d275c..55c0f5a 100644
--- a/services/core/java/com/android/server/location/GnssLocationProvider.java
+++ b/services/core/java/com/android/server/location/GnssLocationProvider.java
@@ -64,6 +64,7 @@
 import android.os.SystemProperties;
 import android.os.UserHandle;
 import android.os.WorkSource;
+import android.os.WorkSource.WorkChain;
 import android.provider.Settings;
 import android.provider.Telephony.Carriers;
 import android.provider.Telephony.Sms.Intents;
@@ -346,7 +347,8 @@
 
     // Current request from underlying location clients.
     private ProviderRequest mProviderRequest = null;
-    // Current list of underlying location clients.
+    // The WorkSource associated with the most recent client request (i.e, most recent call to
+    // setRequest).
     private WorkSource mWorkSource = null;
     // True if gps should be disabled (used to support battery saver mode in settings).
     private boolean mDisableGps = false;
@@ -408,6 +410,7 @@
     private final IAppOpsService mAppOpsService;
     private final IBatteryStats mBatteryStats;
 
+    // Current list of underlying location clients.
     // only modified on handler thread
     private WorkSource mClientSource = new WorkSource();
 
@@ -1345,46 +1348,80 @@
     }
 
     private void updateClientUids(WorkSource source) {
-        // Update work source.
-        WorkSource[] changes = mClientSource.setReturningDiffs(source);
-        if (changes == null) {
+        if (source.equals(mClientSource)) {
             return;
         }
-        WorkSource newWork = changes[0];
-        WorkSource goneWork = changes[1];
 
-        // Update sources that were not previously tracked.
-        if (newWork != null) {
-            int lastuid = -1;
-            for (int i = 0; i < newWork.size(); i++) {
-                try {
-                    int uid = newWork.get(i);
-                    mAppOpsService.startOperation(AppOpsManager.getToken(mAppOpsService),
-                            AppOpsManager.OP_GPS, uid, newWork.getName(i));
-                    if (uid != lastuid) {
-                        lastuid = uid;
-                        mBatteryStats.noteStartGps(uid);
-                    }
-                } catch (RemoteException e) {
-                    Log.w(TAG, "RemoteException", e);
-                }
-            }
+        // (1) Inform BatteryStats that the list of IDs we're tracking changed.
+        try {
+            mBatteryStats.noteGpsChanged(mClientSource, source);
+        } catch (RemoteException e) {
+            Log.w(TAG, "RemoteException", e);
         }
 
-        // Update sources that are no longer tracked.
-        if (goneWork != null) {
-            int lastuid = -1;
-            for (int i = 0; i < goneWork.size(); i++) {
-                try {
-                    int uid = goneWork.get(i);
-                    mAppOpsService.finishOperation(AppOpsManager.getToken(mAppOpsService),
-                            AppOpsManager.OP_GPS, uid, goneWork.getName(i));
-                    if (uid != lastuid) {
-                        lastuid = uid;
-                        mBatteryStats.noteStopGps(uid);
+        // (2) Inform AppOps service about the list of changes to UIDs.
+
+        List<WorkChain>[] diffs = WorkSource.diffChains(mClientSource, source);
+        if (diffs != null) {
+            List<WorkChain> newChains = diffs[0];
+            List<WorkChain> goneChains = diffs[1];
+
+            if (newChains != null) {
+                for (int i = 0; i < newChains.size(); ++i) {
+                    final WorkChain newChain = newChains.get(i);
+                    try {
+                        mAppOpsService.startOperation(AppOpsManager.getToken(mAppOpsService),
+                                AppOpsManager.OP_GPS, newChain.getAttributionUid(),
+                                newChain.getAttributionTag());
+                    } catch (RemoteException e) {
+                        Log.w(TAG, "RemoteException", e);
                     }
-                } catch (RemoteException e) {
-                    Log.w(TAG, "RemoteException", e);
+                }
+            }
+
+            if (goneChains != null) {
+                for (int i = 0; i < goneChains.size(); i++) {
+                    final WorkChain goneChain = goneChains.get(i);
+                    try {
+                        mAppOpsService.finishOperation(AppOpsManager.getToken(mAppOpsService),
+                                AppOpsManager.OP_GPS, goneChain.getAttributionUid(),
+                                goneChain.getAttributionTag());
+                    } catch (RemoteException e) {
+                        Log.w(TAG, "RemoteException", e);
+                    }
+                }
+            }
+
+            mClientSource.transferWorkChains(source);
+        }
+
+        // Update the flat UIDs and names list and inform app-ops of all changes.
+        WorkSource[] changes = mClientSource.setReturningDiffs(source);
+        if (changes != null) {
+            WorkSource newWork = changes[0];
+            WorkSource goneWork = changes[1];
+
+            // Update sources that were not previously tracked.
+            if (newWork != null) {
+                for (int i = 0; i < newWork.size(); i++) {
+                    try {
+                        mAppOpsService.startOperation(AppOpsManager.getToken(mAppOpsService),
+                                AppOpsManager.OP_GPS, newWork.get(i), newWork.getName(i));
+                    } catch (RemoteException e) {
+                        Log.w(TAG, "RemoteException", e);
+                    }
+                }
+            }
+
+            // Update sources that are no longer tracked.
+            if (goneWork != null) {
+                for (int i = 0; i < goneWork.size(); i++) {
+                    try {
+                        mAppOpsService.finishOperation(AppOpsManager.getToken(mAppOpsService),
+                                AppOpsManager.OP_GPS, goneWork.get(i), goneWork.getName(i));
+                    } catch (RemoteException e) {
+                        Log.w(TAG, "RemoteException", e);
+                    }
                 }
             }
         }