Chained attribution for GnssLocationProvider.

The change is slightly more involved than similar changes in this area
because GnssLocationProvider doesn't treat WorkSources as opaque, it
tries to compute diffs between them (using the only call to
WorkSource.setReturningDiffs in the platform). It also uses WorkSources
tags to make calls to AppOps.

Some of this can be improved generally, but this change is narrower in
scope and only tries to improve battery related attribution.

Test: BatteryStatsTests
Bug: 62390666

Change-Id: Ie28bc91d74d3a0c3b6aaef6ba7c59fc58633efb1
diff --git a/core/java/android/os/WorkSource.java b/core/java/android/os/WorkSource.java
index d0c2870..ddcb2d5 100644
--- a/core/java/android/os/WorkSource.java
+++ b/core/java/android/os/WorkSource.java
@@ -479,6 +479,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;
@@ -866,6 +889,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 039b66d..dffe8ac 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);
     }
 
@@ -9830,9 +9873,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);
             }
@@ -9843,14 +9884,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 5a4f7ca..57cee08 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);
+                    }
                 }
             }
         }