Improve power tracking of WIFI use.

We now distribute "wifi started" time across all apps that are
holding WIFI locks that cause it to be started.  But only when
WIFI would not normally be running.  Also have a mechanism to
distribute other WIFI work that has happened across those processes
based on their use.

Also fixed a bug where we were not retaining the CPU speed step
stats across boots...!

Change-Id: I00e3153b98429166273750512cc37e7975211ab9
diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index f182a7a..1e88c56 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -56,9 +56,9 @@
     public static final int SENSOR = 3;
     
     /**
-     * A constant indicating a a wifi turn on timer
+     * A constant indicating a a wifi running timer
      */
-    public static final int WIFI_TURNED_ON = 4;
+    public static final int WIFI_RUNNING = 4;
     
     /**
      * A constant indicating a full wifi lock timer
@@ -249,8 +249,8 @@
          */
         public abstract long getTcpBytesSent(int which);
         
-        public abstract void noteWifiTurnedOnLocked();
-        public abstract void noteWifiTurnedOffLocked();
+        public abstract void noteWifiRunningLocked();
+        public abstract void noteWifiStoppedLocked();
         public abstract void noteFullWifiLockAcquiredLocked();
         public abstract void noteFullWifiLockReleasedLocked();
         public abstract void noteScanWifiLockAcquiredLocked();
@@ -261,7 +261,7 @@
         public abstract void noteAudioTurnedOffLocked();
         public abstract void noteVideoTurnedOnLocked();
         public abstract void noteVideoTurnedOffLocked();
-        public abstract long getWifiTurnedOnTime(long batteryRealtime, int which);
+        public abstract long getWifiRunningTime(long batteryRealtime, int which);
         public abstract long getFullWifiLockTime(long batteryRealtime, int which);
         public abstract long getScanWifiLockTime(long batteryRealtime, int which);
         public abstract long getWifiMulticastTime(long batteryRealtime,
@@ -701,7 +701,7 @@
      *
      * {@hide}
      */
-    public abstract long getWifiRunningTime(long batteryRealtime, int which);
+    public abstract long getGlobalWifiRunningTime(long batteryRealtime, int which);
 
     /**
      * Returns the time in microseconds that bluetooth has been on while the device was
@@ -977,7 +977,7 @@
         final long screenOnTime = getScreenOnTime(batteryRealtime, which);
         final long phoneOnTime = getPhoneOnTime(batteryRealtime, which);
         final long wifiOnTime = getWifiOnTime(batteryRealtime, which);
-        final long wifiRunningTime = getWifiRunningTime(batteryRealtime, which);
+        final long wifiRunningTime = getGlobalWifiRunningTime(batteryRealtime, which);
         final long bluetoothOnTime = getBluetoothOnTime(batteryRealtime, which);
        
         StringBuilder sb = new StringBuilder(128);
@@ -1091,14 +1091,14 @@
             long tx = u.getTcpBytesSent(which);
             long fullWifiLockOnTime = u.getFullWifiLockTime(batteryRealtime, which);
             long scanWifiLockOnTime = u.getScanWifiLockTime(batteryRealtime, which);
-            long wifiTurnedOnTime = u.getWifiTurnedOnTime(batteryRealtime, which);
+            long uidWifiRunningTime = u.getWifiRunningTime(batteryRealtime, which);
             
             if (rx > 0 || tx > 0) dumpLine(pw, uid, category, NETWORK_DATA, rx, tx);
             
             if (fullWifiLockOnTime != 0 || scanWifiLockOnTime != 0
-                    || wifiTurnedOnTime != 0) {
+                    || uidWifiRunningTime != 0) {
                 dumpLine(pw, uid, category, WIFI_LOCK_DATA, 
-                        fullWifiLockOnTime, scanWifiLockOnTime, wifiTurnedOnTime);
+                        fullWifiLockOnTime, scanWifiLockOnTime, uidWifiRunningTime);
             }
 
             if (u.hasUserActivity()) {
@@ -1240,7 +1240,7 @@
         
         final long screenOnTime = getScreenOnTime(batteryRealtime, which);
         final long phoneOnTime = getPhoneOnTime(batteryRealtime, which);
-        final long wifiRunningTime = getWifiRunningTime(batteryRealtime, which);
+        final long wifiRunningTime = getGlobalWifiRunningTime(batteryRealtime, which);
         final long wifiOnTime = getWifiOnTime(batteryRealtime, which);
         final long bluetoothOnTime = getBluetoothOnTime(batteryRealtime, which);
         sb.setLength(0);
@@ -1449,7 +1449,7 @@
             long tcpSent = u.getTcpBytesSent(which);
             long fullWifiLockOnTime = u.getFullWifiLockTime(batteryRealtime, which);
             long scanWifiLockOnTime = u.getScanWifiLockTime(batteryRealtime, which);
-            long wifiTurnedOnTime = u.getWifiTurnedOnTime(batteryRealtime, which);
+            long uidWifiRunningTime = u.getWifiRunningTime(batteryRealtime, which);
             
             if (tcpReceived != 0 || tcpSent != 0) {
                 pw.print(prefix); pw.print("    Network: ");
@@ -1480,11 +1480,11 @@
             }
             
             if (fullWifiLockOnTime != 0 || scanWifiLockOnTime != 0
-                    || wifiTurnedOnTime != 0) {
+                    || uidWifiRunningTime != 0) {
                 sb.setLength(0);
-                sb.append(prefix); sb.append("    Turned Wifi On: "); 
-                        formatTimeMs(sb, wifiTurnedOnTime / 1000); 
-                        sb.append("("); sb.append(formatRatioLocked(wifiTurnedOnTime, 
+                sb.append(prefix); sb.append("    Wifi Running: ");
+                        formatTimeMs(sb, uidWifiRunningTime / 1000);
+                        sb.append("("); sb.append(formatRatioLocked(uidWifiRunningTime,
                                 whichBatteryRealtime)); sb.append(")\n");
                 sb.append(prefix); sb.append("    Full Wifi Lock: "); 
                         formatTimeMs(sb, fullWifiLockOnTime / 1000); 
diff --git a/core/java/com/android/internal/app/IBatteryStats.aidl b/core/java/com/android/internal/app/IBatteryStats.aidl
index bd87a0d..351714e 100644
--- a/core/java/com/android/internal/app/IBatteryStats.aidl
+++ b/core/java/com/android/internal/app/IBatteryStats.aidl
@@ -49,10 +49,11 @@
     void notePhoneSignalStrength(in SignalStrength signalStrength);
     void notePhoneDataConnectionState(int dataType, boolean hasData);
     void notePhoneState(int phoneState);
-    void noteWifiOn(int uid);
-    void noteWifiOff(int uid);
-    void noteWifiRunning();
-    void noteWifiStopped();
+    void noteWifiOn();
+    void noteWifiOff();
+    void noteWifiRunning(in WorkSource ws);
+    void noteWifiRunningChanged(in WorkSource oldWs, in WorkSource newWs);
+    void noteWifiStopped(in WorkSource ws);
     void noteBluetoothOn();
     void noteBluetoothOff();
     void noteFullWifiLockAcquired(int uid);
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 753dbf0..c2d003e 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -66,7 +66,7 @@
     private static final int MAGIC = 0xBA757475; // 'BATSTATS' 
 
     // Current on-disk Parcel version
-    private static final int VERSION = 50;
+    private static final int VERSION = 51;
 
     // Maximum number of items we will record in the history.
     private static final int MAX_HISTORY_ITEMS = 2000;
@@ -129,6 +129,10 @@
     final ArrayList<StopwatchTimer> mWindowTimers = new ArrayList<StopwatchTimer>();
     final SparseArray<ArrayList<StopwatchTimer>> mSensorTimers
             = new SparseArray<ArrayList<StopwatchTimer>>();
+    final ArrayList<StopwatchTimer> mWifiRunningTimers = new ArrayList<StopwatchTimer>();
+    final ArrayList<StopwatchTimer> mFullWifiLockTimers = new ArrayList<StopwatchTimer>();
+    final ArrayList<StopwatchTimer> mScanWifiLockTimers = new ArrayList<StopwatchTimer>();
+    final ArrayList<StopwatchTimer> mWifiMulticastTimers = new ArrayList<StopwatchTimer>();
 
     // Last partial timers we use for distributing CPU usage.
     final ArrayList<StopwatchTimer> mLastPartialTimers = new ArrayList<StopwatchTimer>();
@@ -194,8 +198,8 @@
     StopwatchTimer mWifiOnTimer;
     int mWifiOnUid = -1;
 
-    boolean mWifiRunning;
-    StopwatchTimer mWifiRunningTimer;
+    boolean mGlobalWifiRunning;
+    StopwatchTimer mGlobalWifiRunningTimer;
     
     boolean mBluetoothOn;
     StopwatchTimer mBluetoothOnTimer;
@@ -1769,7 +1773,7 @@
         }
     }
     
-    public void noteWifiOnLocked(int uid) {
+    public void noteWifiOnLocked() {
         if (!mWifiOn) {
             mHistoryCur.states |= HistoryItem.STATE_WIFI_ON_FLAG;
             if (DEBUG_HISTORY) Slog.v(TAG, "WIFI on to: "
@@ -1778,16 +1782,9 @@
             mWifiOn = true;
             mWifiOnTimer.startRunningLocked(this);
         }
-        if (mWifiOnUid != uid) {
-            if (mWifiOnUid >= 0) {
-                getUidStatsLocked(mWifiOnUid).noteWifiTurnedOffLocked();
-            }
-            mWifiOnUid = uid;
-            getUidStatsLocked(uid).noteWifiTurnedOnLocked();
-        }
     }
     
-    public void noteWifiOffLocked(int uid) {
+    public void noteWifiOffLocked() {
         if (mWifiOn) {
             mHistoryCur.states &= ~HistoryItem.STATE_WIFI_ON_FLAG;
             if (DEBUG_HISTORY) Slog.v(TAG, "WIFI off to: "
@@ -1797,7 +1794,7 @@
             mWifiOnTimer.stopRunningLocked(this);
         }
         if (mWifiOnUid >= 0) {
-            getUidStatsLocked(mWifiOnUid).noteWifiTurnedOffLocked();
+            getUidStatsLocked(mWifiOnUid).noteWifiStoppedLocked();
             mWifiOnUid = -1;
         }
     }
@@ -1850,25 +1847,52 @@
         getUidStatsLocked(uid).noteVideoTurnedOffLocked();
     }
 
-    public void noteWifiRunningLocked() {
-        if (!mWifiRunning) {
+    public void noteWifiRunningLocked(WorkSource ws) {
+        if (!mGlobalWifiRunning) {
             mHistoryCur.states |= HistoryItem.STATE_WIFI_RUNNING_FLAG;
             if (DEBUG_HISTORY) Slog.v(TAG, "WIFI running to: "
                     + Integer.toHexString(mHistoryCur.states));
             addHistoryRecordLocked(SystemClock.elapsedRealtime());
-            mWifiRunning = true;
-            mWifiRunningTimer.startRunningLocked(this);
+            mGlobalWifiRunning = true;
+            mGlobalWifiRunningTimer.startRunningLocked(this);
+            int N = ws.size();
+            for (int i=0; i<N; i++) {
+                getUidStatsLocked(ws.get(i)).noteWifiRunningLocked();
+            }
+        } else {
+            Log.w(TAG, "noteWifiRunningLocked -- called while WIFI running");
         }
     }
 
-    public void noteWifiStoppedLocked() {
-        if (mWifiRunning) {
+    public void noteWifiRunningChangedLocked(WorkSource oldWs, WorkSource newWs) {
+        if (mGlobalWifiRunning) {
+            int N = oldWs.size();
+            for (int i=0; i<N; i++) {
+                getUidStatsLocked(oldWs.get(i)).noteWifiStoppedLocked();
+            }
+            N = newWs.size();
+            for (int i=0; i<N; i++) {
+                getUidStatsLocked(newWs.get(i)).noteWifiRunningLocked();
+            }
+        } else {
+            Log.w(TAG, "noteWifiRunningChangedLocked -- called while WIFI not running");
+        }
+    }
+
+    public void noteWifiStoppedLocked(WorkSource ws) {
+        if (mGlobalWifiRunning) {
             mHistoryCur.states &= ~HistoryItem.STATE_WIFI_RUNNING_FLAG;
             if (DEBUG_HISTORY) Slog.v(TAG, "WIFI stopped to: "
                     + Integer.toHexString(mHistoryCur.states));
             addHistoryRecordLocked(SystemClock.elapsedRealtime());
-            mWifiRunning = false;
-            mWifiRunningTimer.stopRunningLocked(this);
+            mGlobalWifiRunning = false;
+            mGlobalWifiRunningTimer.stopRunningLocked(this);
+            int N = ws.size();
+            for (int i=0; i<N; i++) {
+                getUidStatsLocked(ws.get(i)).noteWifiStoppedLocked();
+            }
+        } else {
+            Log.w(TAG, "noteWifiStoppedLocked -- called while WIFI not running");
         }
     }
 
@@ -2056,8 +2080,8 @@
         return mWifiOnTimer.getTotalTimeLocked(batteryRealtime, which);
     }
     
-    @Override public long getWifiRunningTime(long batteryRealtime, int which) {
-        return mWifiRunningTimer.getTotalTimeLocked(batteryRealtime, which);
+    @Override public long getGlobalWifiRunningTime(long batteryRealtime, int which) {
+        return mGlobalWifiRunningTimer.getTotalTimeLocked(batteryRealtime, which);
     }
 
     @Override public long getBluetoothOnTime(long batteryRealtime, int which) {
@@ -2090,8 +2114,8 @@
         long mStartedTcpBytesReceived = -1;
         long mStartedTcpBytesSent = -1;
         
-        boolean mWifiTurnedOn;
-        StopwatchTimer mWifiTurnedOnTimer;
+        boolean mWifiRunning;
+        StopwatchTimer mWifiRunningTimer;
         
         boolean mFullWifiLockOut;
         StopwatchTimer mFullWifiLockTimer;
@@ -2137,14 +2161,14 @@
 
         public Uid(int uid) {
             mUid = uid;
-            mWifiTurnedOnTimer = new StopwatchTimer(Uid.this, WIFI_TURNED_ON,
-                    null, mUnpluggables);
+            mWifiRunningTimer = new StopwatchTimer(Uid.this, WIFI_RUNNING,
+                    mWifiRunningTimers, mUnpluggables);
             mFullWifiLockTimer = new StopwatchTimer(Uid.this, FULL_WIFI_LOCK,
-                    null, mUnpluggables);
+                    mFullWifiLockTimers, mUnpluggables);
             mScanWifiLockTimer = new StopwatchTimer(Uid.this, SCAN_WIFI_LOCK,
-                    null, mUnpluggables);
+                    mScanWifiLockTimers, mUnpluggables);
             mWifiMulticastTimer = new StopwatchTimer(Uid.this, WIFI_MULTICAST_ENABLED,
-                    null, mUnpluggables);
+                    mWifiMulticastTimers, mUnpluggables);
             mAudioTurnedOnTimer = new StopwatchTimer(Uid.this, AUDIO_TURNED_ON,
                     null, mUnpluggables);
             mVideoTurnedOnTimer = new StopwatchTimer(Uid.this, VIDEO_TURNED_ON,
@@ -2212,22 +2236,22 @@
         }
         
         @Override
-        public void noteWifiTurnedOnLocked() {
-            if (!mWifiTurnedOn) {
-                mWifiTurnedOn = true;
-                if (mWifiTurnedOnTimer == null) {
-                    mWifiTurnedOnTimer = new StopwatchTimer(Uid.this, WIFI_TURNED_ON,
-                            null, mUnpluggables);
+        public void noteWifiRunningLocked() {
+            if (!mWifiRunning) {
+                mWifiRunning = true;
+                if (mWifiRunningTimer == null) {
+                    mWifiRunningTimer = new StopwatchTimer(Uid.this, WIFI_RUNNING,
+                            mWifiRunningTimers, mUnpluggables);
                 }
-                mWifiTurnedOnTimer.startRunningLocked(BatteryStatsImpl.this);
+                mWifiRunningTimer.startRunningLocked(BatteryStatsImpl.this);
             }
         }
         
         @Override
-        public void noteWifiTurnedOffLocked() {
-            if (mWifiTurnedOn) {
-                mWifiTurnedOn = false;
-                mWifiTurnedOnTimer.stopRunningLocked(BatteryStatsImpl.this);
+        public void noteWifiStoppedLocked() {
+            if (mWifiRunning) {
+                mWifiRunning = false;
+                mWifiRunningTimer.stopRunningLocked(BatteryStatsImpl.this);
             }
         }
         
@@ -2237,7 +2261,7 @@
                 mFullWifiLockOut = true;
                 if (mFullWifiLockTimer == null) {
                     mFullWifiLockTimer = new StopwatchTimer(Uid.this, FULL_WIFI_LOCK,
-                            null, mUnpluggables);
+                            mFullWifiLockTimers, mUnpluggables);
                 }
                 mFullWifiLockTimer.startRunningLocked(BatteryStatsImpl.this);
             }
@@ -2257,7 +2281,7 @@
                 mScanWifiLockOut = true;
                 if (mScanWifiLockTimer == null) {
                     mScanWifiLockTimer = new StopwatchTimer(Uid.this, SCAN_WIFI_LOCK,
-                            null, mUnpluggables);
+                            mScanWifiLockTimers, mUnpluggables);
                 }
                 mScanWifiLockTimer.startRunningLocked(BatteryStatsImpl.this);
             }
@@ -2277,7 +2301,7 @@
                 mWifiMulticastEnabled = true;
                 if (mWifiMulticastTimer == null) {
                     mWifiMulticastTimer = new StopwatchTimer(Uid.this, WIFI_MULTICAST_ENABLED,
-                            null, mUnpluggables);
+                            mWifiMulticastTimers, mUnpluggables);
                 }
                 mWifiMulticastTimer.startRunningLocked(BatteryStatsImpl.this);
             }
@@ -2332,11 +2356,11 @@
         }
 
         @Override 
-        public long getWifiTurnedOnTime(long batteryRealtime, int which) {
-            if (mWifiTurnedOnTimer == null) {
+        public long getWifiRunningTime(long batteryRealtime, int which) {
+            if (mWifiRunningTimer == null) {
                 return 0;
             }
-            return mWifiTurnedOnTimer.getTotalTimeLocked(batteryRealtime, which);
+            return mWifiRunningTimer.getTotalTimeLocked(batteryRealtime, which);
         }
 
         @Override 
@@ -2422,9 +2446,9 @@
         boolean reset() {
             boolean active = false;
             
-            if (mWifiTurnedOnTimer != null) {
-                active |= !mWifiTurnedOnTimer.reset(BatteryStatsImpl.this, false);
-                active |= mWifiTurnedOn;
+            if (mWifiRunningTimer != null) {
+                active |= !mWifiRunningTimer.reset(BatteryStatsImpl.this, false);
+                active |= mWifiRunning;
             }
             if (mFullWifiLockTimer != null) {
                 active |= !mFullWifiLockTimer.reset(BatteryStatsImpl.this, false);
@@ -2517,8 +2541,8 @@
             mPids.clear();
 
             if (!active) {
-                if (mWifiTurnedOnTimer != null) {
-                    mWifiTurnedOnTimer.detach();
+                if (mWifiRunningTimer != null) {
+                    mWifiRunningTimer.detach();
                 }
                 if (mFullWifiLockTimer != null) {
                     mFullWifiLockTimer.detach();
@@ -2580,9 +2604,9 @@
             out.writeLong(computeCurrentTcpBytesSent());
             out.writeLong(mTcpBytesReceivedAtLastUnplug);
             out.writeLong(mTcpBytesSentAtLastUnplug);
-            if (mWifiTurnedOnTimer != null) {
+            if (mWifiRunningTimer != null) {
                 out.writeInt(1);
-                mWifiTurnedOnTimer.writeToParcel(out, batteryRealtime);
+                mWifiRunningTimer.writeToParcel(out, batteryRealtime);
             } else {
                 out.writeInt(0);
             }
@@ -2674,31 +2698,31 @@
             mCurrentTcpBytesSent = in.readLong();
             mTcpBytesReceivedAtLastUnplug = in.readLong();
             mTcpBytesSentAtLastUnplug = in.readLong();
-            mWifiTurnedOn = false;
+            mWifiRunning = false;
             if (in.readInt() != 0) {
-                mWifiTurnedOnTimer = new StopwatchTimer(Uid.this, WIFI_TURNED_ON,
-                        null, mUnpluggables, in);
+                mWifiRunningTimer = new StopwatchTimer(Uid.this, WIFI_RUNNING,
+                        mWifiRunningTimers, mUnpluggables, in);
             } else {
-                mWifiTurnedOnTimer = null;
+                mWifiRunningTimer = null;
             }
             mFullWifiLockOut = false;
             if (in.readInt() != 0) {
                 mFullWifiLockTimer = new StopwatchTimer(Uid.this, FULL_WIFI_LOCK,
-                        null, mUnpluggables, in);
+                        mFullWifiLockTimers, mUnpluggables, in);
             } else {
                 mFullWifiLockTimer = null;
             }
             mScanWifiLockOut = false;
             if (in.readInt() != 0) {
                 mScanWifiLockTimer = new StopwatchTimer(Uid.this, SCAN_WIFI_LOCK,
-                        null, mUnpluggables, in);
+                        mScanWifiLockTimers, mUnpluggables, in);
             } else {
                 mScanWifiLockTimer = null;
             }
             mWifiMulticastEnabled = false;
             if (in.readInt() != 0) {
                 mWifiMulticastTimer = new StopwatchTimer(Uid.this, WIFI_MULTICAST_ENABLED,
-                        null, mUnpluggables, in);
+                        mWifiMulticastTimers, mUnpluggables, in);
             } else {
                 mWifiMulticastTimer = null;
             }
@@ -3771,7 +3795,7 @@
             mPhoneDataConnectionsTimer[i] = new StopwatchTimer(null, -300-i, null, mUnpluggables);
         }
         mWifiOnTimer = new StopwatchTimer(null, -3, null, mUnpluggables);
-        mWifiRunningTimer = new StopwatchTimer(null, -4, null, mUnpluggables);
+        mGlobalWifiRunningTimer = new StopwatchTimer(null, -4, null, mUnpluggables);
         mBluetoothOnTimer = new StopwatchTimer(null, -5, null, mUnpluggables);
         mAudioOnTimer = new StopwatchTimer(null, -6, null, mUnpluggables);
         mVideoOnTimer = new StopwatchTimer(null, -7, null, mUnpluggables);
@@ -3861,7 +3885,7 @@
             mPhoneDataConnectionsTimer[i].reset(this, false);
         }
         mWifiOnTimer.reset(this, false);
-        mWifiRunningTimer.reset(this, false);
+        mGlobalWifiRunningTimer.reset(this, false);
         mBluetoothOnTimer.reset(this, false);
         
         for (int i=0; i<mUidStats.size(); i++) {
@@ -4280,6 +4304,57 @@
         return u.getServiceStatsLocked(pkg, name);
     }
 
+    /**
+     * Massage data to distribute any reasonable work down to more specific
+     * owners.  Must only be called on a dead BatteryStats object!
+     */
+    public void distributeWorkLocked(int which) {
+        // Aggregate all CPU time associated with WIFI.
+        Uid wifiUid = mUidStats.get(Process.WIFI_UID);
+        if (wifiUid != null) {
+            long uSecTime = computeBatteryRealtime(SystemClock.elapsedRealtime() * 1000, which);
+            for (Uid.Proc proc : wifiUid.mProcessStats.values()) {
+                long totalRunningTime = getGlobalWifiRunningTime(uSecTime, which);
+                for (int i=0; i<mUidStats.size(); i++) {
+                    Uid uid = mUidStats.valueAt(i);
+                    if (uid.mUid != Process.WIFI_UID) {
+                        long uidRunningTime = uid.getWifiRunningTime(uSecTime, which);
+                        if (uidRunningTime > 0) {
+                            Uid.Proc uidProc = uid.getProcessStatsLocked("*wifi*");
+                            long time = proc.getUserTime(which);
+                            time = (time*uidRunningTime)/totalRunningTime;
+                            uidProc.mUserTime += time;
+                            proc.mUserTime -= time;
+                            time = proc.getSystemTime(which);
+                            time = (time*uidRunningTime)/totalRunningTime;
+                            uidProc.mSystemTime += time;
+                            proc.mSystemTime -= time;
+                            time = proc.getForegroundTime(which);
+                            time = (time*uidRunningTime)/totalRunningTime;
+                            uidProc.mForegroundTime += time;
+                            proc.mForegroundTime -= time;
+                            for (int sb=0; sb<proc.mSpeedBins.length; sb++) {
+                                SamplingCounter sc = proc.mSpeedBins[sb];
+                                if (sc != null) {
+                                    time = sc.getCountLocked(which);
+                                    time = (time*uidRunningTime)/totalRunningTime;
+                                    SamplingCounter uidSc = uidProc.mSpeedBins[sb];
+                                    if (uidSc == null) {
+                                        uidSc = new SamplingCounter(mUnpluggables);
+                                        uidProc.mSpeedBins[sb] = uidSc;
+                                    }
+                                    uidSc.mCount.addAndGet((int)time);
+                                    sc.mCount.addAndGet((int)-time);
+                                }
+                            }
+                            totalRunningTime -= uidRunningTime;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     public void shutdownLocked() {
         writeLocked();
         mShuttingDown = true;
@@ -4442,8 +4517,8 @@
         }
         mWifiOn = false;
         mWifiOnTimer.readSummaryFromParcelLocked(in);
-        mWifiRunning = false;
-        mWifiRunningTimer.readSummaryFromParcelLocked(in);
+        mGlobalWifiRunning = false;
+        mGlobalWifiRunningTimer.readSummaryFromParcelLocked(in);
         mBluetoothOn = false;
         mBluetoothOnTimer.readSummaryFromParcelLocked(in);
 
@@ -4471,9 +4546,9 @@
             Uid u = new Uid(uid);
             mUidStats.put(uid, u);
 
-            u.mWifiTurnedOn = false;
+            u.mWifiRunning = false;
             if (in.readInt() != 0) {
-                u.mWifiTurnedOnTimer.readSummaryFromParcelLocked(in);
+                u.mWifiRunningTimer.readSummaryFromParcelLocked(in);
             }
             u.mFullWifiLockOut = false;
             if (in.readInt() != 0) {
@@ -4547,6 +4622,14 @@
                 p.mUserTime = p.mLoadedUserTime = in.readLong();
                 p.mSystemTime = p.mLoadedSystemTime = in.readLong();
                 p.mStarts = p.mLoadedStarts = in.readInt();
+                int NSB = in.readInt();
+                p.mSpeedBins = new SamplingCounter[NSB];
+                for (int i=0; i<NSB; i++) {
+                    if (in.readInt() != 0) {
+                        p.mSpeedBins[i] = new SamplingCounter(mUnpluggables);
+                        p.mSpeedBins[i].readSummaryFromParcelLocked(in);
+                    }
+                }
                 p.readExcessiveWakeFromParcelLocked(in);
             }
 
@@ -4614,7 +4697,7 @@
             mPhoneDataConnectionsTimer[i].writeSummaryFromParcelLocked(out, NOWREAL);
         }
         mWifiOnTimer.writeSummaryFromParcelLocked(out, NOWREAL);
-        mWifiRunningTimer.writeSummaryFromParcelLocked(out, NOWREAL);
+        mGlobalWifiRunningTimer.writeSummaryFromParcelLocked(out, NOWREAL);
         mBluetoothOnTimer.writeSummaryFromParcelLocked(out, NOWREAL);
 
         out.writeInt(mKernelWakelockStats.size());
@@ -4636,9 +4719,9 @@
             out.writeInt(mUidStats.keyAt(iu));
             Uid u = mUidStats.valueAt(iu);
             
-            if (u.mWifiTurnedOnTimer != null) {
+            if (u.mWifiRunningTimer != null) {
                 out.writeInt(1);
-                u.mWifiTurnedOnTimer.writeSummaryFromParcelLocked(out, NOWREAL);
+                u.mWifiRunningTimer.writeSummaryFromParcelLocked(out, NOWREAL);
             } else {
                 out.writeInt(0);
             }
@@ -4736,6 +4819,16 @@
                     out.writeLong(ps.mUserTime);
                     out.writeLong(ps.mSystemTime);
                     out.writeInt(ps.mStarts);
+                    final int N = ps.mSpeedBins.length;
+                    out.writeInt(N);
+                    for (int i=0; i<N; i++) {
+                        if (ps.mSpeedBins[i] != null) {
+                            out.writeInt(1);
+                            ps.mSpeedBins[i].writeSummaryFromParcelLocked(out);
+                        } else {
+                            out.writeInt(0);
+                        }
+                    }
                     ps.writeExcessiveWakeToParcelLocked(out);
                 }
             }
@@ -4806,8 +4899,8 @@
         }
         mWifiOn = false;
         mWifiOnTimer = new StopwatchTimer(null, -2, null, mUnpluggables, in);
-        mWifiRunning = false;
-        mWifiRunningTimer = new StopwatchTimer(null, -2, null, mUnpluggables, in);
+        mGlobalWifiRunning = false;
+        mGlobalWifiRunningTimer = new StopwatchTimer(null, -2, null, mUnpluggables, in);
         mBluetoothOn = false;
         mBluetoothOnTimer = new StopwatchTimer(null, -2, null, mUnpluggables, in);
         mUptime = in.readLong();
@@ -4859,6 +4952,10 @@
         mPartialTimers.clear();
         mFullTimers.clear();
         mWindowTimers.clear();
+        mWifiRunningTimers.clear();
+        mFullWifiLockTimers.clear();
+        mScanWifiLockTimers.clear();
+        mWifiMulticastTimers.clear();
 
         sNumSpeedSteps = in.readInt();
 
@@ -4908,7 +5005,7 @@
             mPhoneDataConnectionsTimer[i].writeToParcel(out, batteryRealtime);
         }
         mWifiOnTimer.writeToParcel(out, batteryRealtime);
-        mWifiRunningTimer.writeToParcel(out, batteryRealtime);
+        mGlobalWifiRunningTimer.writeToParcel(out, batteryRealtime);
         mBluetoothOnTimer.writeToParcel(out, batteryRealtime);
         out.writeLong(mUptime);
         out.writeLong(mUptimeStart);
@@ -5006,7 +5103,7 @@
             pr.println("*** Wifi timer:");
             mWifiOnTimer.logState(pr, "  ");
             pr.println("*** WifiRunning timer:");
-            mWifiRunningTimer.logState(pr, "  ");
+            mGlobalWifiRunningTimer.logState(pr, "  ");
             pr.println("*** Bluetooth timer:");
             mBluetoothOnTimer.logState(pr, "  ");
         }
diff --git a/services/java/com/android/server/WifiService.java b/services/java/com/android/server/WifiService.java
index f11c0f7..371f22e 100644
--- a/services/java/com/android/server/WifiService.java
+++ b/services/java/com/android/server/WifiService.java
@@ -189,6 +189,12 @@
     private static final int SCAN_RESULT_BUFFER_SIZE = 512;
     private boolean mNeedReconfig;
 
+    /**
+     * Temporary for computing UIDS that are responsible for starting WIFI.
+     * Protected by mWifiStateTracker lock.
+     */
+    private final WorkSource mTmpWorkSource = new WorkSource();
+
     /*
      * Last UID that asked to enable WIFI.
      */
@@ -529,9 +535,9 @@
         long ident = Binder.clearCallingIdentity();
         try {
             if (wifiState == WIFI_STATE_ENABLED) {
-                mBatteryStats.noteWifiOn(uid);
+                mBatteryStats.noteWifiOn();
             } else if (wifiState == WIFI_STATE_DISABLED) {
-                mBatteryStats.noteWifiOff(uid);
+                mBatteryStats.noteWifiOff();
             }
         } catch (RemoteException e) {
         } finally {
@@ -788,9 +794,9 @@
         long ident = Binder.clearCallingIdentity();
         try {
             if (wifiAPState == WIFI_AP_STATE_ENABLED) {
-                mBatteryStats.noteWifiOn(uid);
+                mBatteryStats.noteWifiOn();
             } else if (wifiAPState == WIFI_AP_STATE_DISABLED) {
-                mBatteryStats.noteWifiOff(uid);
+                mBatteryStats.noteWifiOff();
             }
         } catch (RemoteException e) {
         } finally {
@@ -1664,6 +1670,9 @@
                 mAlarmManager.cancel(mIdleIntent);
                 mDeviceIdle = false;
                 mScreenOff = false;
+                // Once the screen is on, we are not keeping WIFI running
+                // because of any locks so clear that tracking immediately.
+                reportStartWorkSource();
                 mWifiStateTracker.enableRssiPolling(true);
                 /* DHCP or other temporary failures in the past can prevent
                  * a disabled network from being connected to, enable on screen on
@@ -1705,6 +1714,7 @@
             } else if (action.equals(ACTION_DEVICE_IDLE)) {
                 Slog.d(TAG, "got ACTION_DEVICE_IDLE");
                 mDeviceIdle = true;
+                reportStartWorkSource();
             } else if (action.equals(Intent.ACTION_BATTERY_CHANGED)) {
                 /*
                  * Set a timer to put Wi-Fi to sleep, but only if the screen is off
@@ -1806,6 +1816,18 @@
         Message.obtain(mWifiHandler, MESSAGE_ENABLE_NETWORKS).sendToTarget();
     }
 
+    private void reportStartWorkSource() {
+        synchronized (mWifiStateTracker) {
+            mTmpWorkSource.clear();
+            if (mDeviceIdle) {
+                for (int i=0; i<mLocks.mList.size(); i++) {
+                    mTmpWorkSource.add(mLocks.mList.get(i).mWorkSource);
+                }
+            }
+            mWifiStateTracker.updateBatteryWorkSourceLocked(mTmpWorkSource);
+        }
+    }
+
     private void updateWifiState() {
         // send a message so it's all serialized
         Message.obtain(mWifiHandler, MESSAGE_UPDATE_STATE, 0, 0).sendToTarget();
@@ -1814,7 +1836,12 @@
     private void doUpdateWifiState() {
         boolean wifiEnabled = getPersistedWifiEnabled();
         boolean airplaneMode = isAirplaneModeOn() && !mAirplaneModeOverwridden;
-        boolean lockHeld = mLocks.hasLocks();
+
+        boolean lockHeld;
+        synchronized (mLocks) {
+            lockHeld = mLocks.hasLocks();
+        }
+
         int strongestLockMode = WifiManager.WIFI_MODE_FULL;
         boolean wifiShouldBeEnabled = wifiEnabled && !airplaneMode;
         boolean wifiShouldBeStarted = !mDeviceIdle || lockHeld;
@@ -1922,6 +1949,7 @@
                     break;
 
                 case MESSAGE_START_WIFI:
+                    reportStartWorkSource();
                     mWifiStateTracker.setScanOnlyMode(msg.arg1 == WifiManager.WIFI_MODE_SCAN_ONLY);
                     mWifiStateTracker.restart();
                     mWifiStateTracker.setHighPerfMode(msg.arg1 ==
@@ -2198,6 +2226,10 @@
             Binder.restoreCallingIdentity(ident);
         }
 
+        // Be aggressive about adding new locks into the accounted state...
+        // we want to over-report rather than under-report.
+        reportStartWorkSource();
+
         updateWifiState();
         return true;
     }
diff --git a/services/java/com/android/server/am/BatteryStatsService.java b/services/java/com/android/server/am/BatteryStatsService.java
index bb40967..73a5435 100644
--- a/services/java/com/android/server/am/BatteryStatsService.java
+++ b/services/java/com/android/server/am/BatteryStatsService.java
@@ -218,17 +218,17 @@
         }
     }
 
-    public void noteWifiOn(int uid) {
+    public void noteWifiOn() {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteWifiOnLocked(uid);
+            mStats.noteWifiOnLocked();
         }
     }
     
-    public void noteWifiOff(int uid) {
+    public void noteWifiOff() {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteWifiOffLocked(uid);
+            mStats.noteWifiOffLocked();
         }
     }
 
@@ -260,17 +260,24 @@
         }
     }
 
-    public void noteWifiRunning() {
+    public void noteWifiRunning(WorkSource ws) {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteWifiRunningLocked();
+            mStats.noteWifiRunningLocked(ws);
         }
     }
 
-    public void noteWifiStopped() {
+    public void noteWifiRunningChanged(WorkSource oldWs, WorkSource newWs) {
         enforceCallingPermission();
         synchronized (mStats) {
-            mStats.noteWifiStoppedLocked();
+            mStats.noteWifiRunningChangedLocked(oldWs, newWs);
+        }
+    }
+
+    public void noteWifiStopped(WorkSource ws) {
+        enforceCallingPermission();
+        synchronized (mStats) {
+            mStats.noteWifiStoppedLocked(ws);
         }
     }
 
diff --git a/wifi/java/android/net/wifi/WifiStateTracker.java b/wifi/java/android/net/wifi/WifiStateTracker.java
index 22dbda3..9d27bde 100644
--- a/wifi/java/android/net/wifi/WifiStateTracker.java
+++ b/wifi/java/android/net/wifi/WifiStateTracker.java
@@ -38,6 +38,7 @@
 import android.os.Looper;
 import android.os.RemoteException;
 import android.os.ServiceManager;
+import android.os.WorkSource;
 import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.EventLog;
@@ -323,6 +324,21 @@
     private static String[] sDnsPropNames;
 
     /**
+     * Keep track of whether we last told the battery stats we had started.
+     */
+    private boolean mReportedRunning = false;
+
+    /**
+     * Most recently set source of starting WIFI.
+     */
+    private final WorkSource mRunningWifiUids = new WorkSource();
+
+    /**
+     * The last reported UIDs that were responsible for starting WIFI.
+     */
+    private final WorkSource mLastRunningWifiUids = new WorkSource();
+
+    /**
      * A structure for supplying information about a supplicant state
      * change in the STATE_CHANGE event message that comes from the
      * WifiMonitor
@@ -632,12 +648,35 @@
         return mRunState == RUN_STATE_STOPPED || mRunState == RUN_STATE_STOPPING;
     }
 
-    private void noteRunState() {
+    public void updateBatteryWorkSourceLocked(WorkSource newSource) {
         try {
+            if (newSource != null) {
+                mRunningWifiUids.set(newSource);
+            }
             if (mRunState == RUN_STATE_RUNNING) {
-                mBatteryStats.noteWifiRunning();
+                if (mReportedRunning) {
+                    // If the work source has changed since last time, need
+                    // to remove old work from battery stats.
+                    if (mLastRunningWifiUids.diff(mRunningWifiUids)) {
+                        mBatteryStats.noteWifiRunningChanged(mLastRunningWifiUids,
+                                mRunningWifiUids);
+                        mLastRunningWifiUids.set(mRunningWifiUids);
+                    }
+                } else {
+                    // Now being started, report it.
+                    mBatteryStats.noteWifiRunning(mRunningWifiUids);
+                    mLastRunningWifiUids.set(mRunningWifiUids);
+                    mReportedRunning = true;
+                }
             } else if (mRunState == RUN_STATE_STOPPED) {
-                mBatteryStats.noteWifiStopped();
+                if (mReportedRunning) {
+                    // Last reported we were running, time to stop.
+                    mBatteryStats.noteWifiStopped(mLastRunningWifiUids);
+                    mLastRunningWifiUids.clear();
+                    mReportedRunning = false;
+                }
+            } else {
+                // State in transition -- nothing to update yet.
             }
         } catch (RemoteException ignore) {
         }
@@ -801,7 +840,9 @@
         switch (msg.what) {
             case EVENT_SUPPLICANT_CONNECTION:
                 mRunState = RUN_STATE_RUNNING;
-                noteRunState();
+                synchronized (this) {
+                    updateBatteryWorkSourceLocked(null);
+                }
                 checkUseStaticIp();
                 /* Reset notification state on new connection */
                 resetNotificationTimer();
@@ -875,7 +916,9 @@
 
             case EVENT_SUPPLICANT_DISCONNECT:
                 mRunState = RUN_STATE_STOPPED;
-                noteRunState();
+                synchronized (this) {
+                    updateBatteryWorkSourceLocked(null);
+                }
                 boolean died = mWifiState.get() != WIFI_STATE_DISABLED &&
                                mWifiState.get() != WIFI_STATE_DISABLING;
                 if (died) {
@@ -1267,7 +1310,9 @@
                     mWM.setWifiEnabled(true);
                     break;
                 }
-                noteRunState();
+                synchronized (this) {
+                    updateBatteryWorkSourceLocked(null);
+                }
                 break;
 
             case EVENT_PASSWORD_KEY_MAY_BE_INCORRECT: