diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index f182a7ad6d714..1e88c5647666f 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -56,9 +56,9 @@ public abstract class BatteryStats implements Parcelable {
     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 class BatteryStats implements Parcelable {
          */
         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 class BatteryStats implements Parcelable {
         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 @@ public abstract class BatteryStats implements Parcelable {
      *
      * {@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 @@ public abstract class BatteryStats implements Parcelable {
         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 @@ public abstract class BatteryStats implements Parcelable {
             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 @@ public abstract class BatteryStats implements Parcelable {
         
         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 @@ public abstract class BatteryStats implements Parcelable {
             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 @@ public abstract class BatteryStats implements Parcelable {
             }
             
             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 bd87a0d05261c..351714ec85225 100644
--- a/core/java/com/android/internal/app/IBatteryStats.aidl
+++ b/core/java/com/android/internal/app/IBatteryStats.aidl
@@ -49,10 +49,11 @@ interface IBatteryStats {
     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 753dbf0e41416..c2d003e96293d 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -66,7 +66,7 @@ public final class BatteryStatsImpl extends BatteryStats {
     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 @@ public final class BatteryStatsImpl extends BatteryStats {
     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 @@ public final class BatteryStatsImpl extends BatteryStats {
     StopwatchTimer mWifiOnTimer;
     int mWifiOnUid = -1;
 
-    boolean mWifiRunning;
-    StopwatchTimer mWifiRunningTimer;
+    boolean mGlobalWifiRunning;
+    StopwatchTimer mGlobalWifiRunningTimer;
     
     boolean mBluetoothOn;
     StopwatchTimer mBluetoothOnTimer;
@@ -1769,7 +1773,7 @@ public final class BatteryStatsImpl extends BatteryStats {
         }
     }
     
-    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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
             mWifiOnTimer.stopRunningLocked(this);
         }
         if (mWifiOnUid >= 0) {
-            getUidStatsLocked(mWifiOnUid).noteWifiTurnedOffLocked();
+            getUidStatsLocked(mWifiOnUid).noteWifiStoppedLocked();
             mWifiOnUid = -1;
         }
     }
@@ -1850,25 +1847,52 @@ public final class BatteryStatsImpl extends BatteryStats {
         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 @@ public final class BatteryStatsImpl extends BatteryStats {
         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 @@ public final class BatteryStatsImpl extends BatteryStats {
         long mStartedTcpBytesReceived = -1;
         long mStartedTcpBytesSent = -1;
         
-        boolean mWifiTurnedOn;
-        StopwatchTimer mWifiTurnedOnTimer;
+        boolean mWifiRunning;
+        StopwatchTimer mWifiRunningTimer;
         
         boolean mFullWifiLockOut;
         StopwatchTimer mFullWifiLockTimer;
@@ -2137,14 +2161,14 @@ public final class BatteryStatsImpl extends BatteryStats {
 
         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 @@ public final class BatteryStatsImpl extends BatteryStats {
         }
         
         @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 @@ public final class BatteryStatsImpl extends BatteryStats {
                 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 @@ public final class BatteryStatsImpl extends BatteryStats {
                 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 @@ public final class BatteryStatsImpl extends BatteryStats {
                 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 @@ public final class BatteryStatsImpl extends BatteryStats {
         }
 
         @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 @@ public final class BatteryStatsImpl extends BatteryStats {
         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 @@ public final class BatteryStatsImpl extends BatteryStats {
             mPids.clear();
 
             if (!active) {
-                if (mWifiTurnedOnTimer != null) {
-                    mWifiTurnedOnTimer.detach();
+                if (mWifiRunningTimer != null) {
+                    mWifiRunningTimer.detach();
                 }
                 if (mFullWifiLockTimer != null) {
                     mFullWifiLockTimer.detach();
@@ -2580,9 +2604,9 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
         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 @@ public final class BatteryStatsImpl extends BatteryStats {
         }
         mWifiOn = false;
         mWifiOnTimer.readSummaryFromParcelLocked(in);
-        mWifiRunning = false;
-        mWifiRunningTimer.readSummaryFromParcelLocked(in);
+        mGlobalWifiRunning = false;
+        mGlobalWifiRunningTimer.readSummaryFromParcelLocked(in);
         mBluetoothOn = false;
         mBluetoothOnTimer.readSummaryFromParcelLocked(in);
 
@@ -4471,9 +4546,9 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
                 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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
                     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 @@ public final class BatteryStatsImpl extends BatteryStats {
         }
         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 @@ public final class BatteryStatsImpl extends BatteryStats {
         mPartialTimers.clear();
         mFullTimers.clear();
         mWindowTimers.clear();
+        mWifiRunningTimers.clear();
+        mFullWifiLockTimers.clear();
+        mScanWifiLockTimers.clear();
+        mWifiMulticastTimers.clear();
 
         sNumSpeedSteps = in.readInt();
 
@@ -4908,7 +5005,7 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 @@ public final class BatteryStatsImpl extends BatteryStats {
             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 f11c0f74c0f2f..371f22ee085ec 100644
--- a/services/java/com/android/server/WifiService.java
+++ b/services/java/com/android/server/WifiService.java
@@ -189,6 +189,12 @@ public class WifiService extends IWifiManager.Stub {
     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 @@ public class WifiService extends IWifiManager.Stub {
         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 @@ public class WifiService extends IWifiManager.Stub {
         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 @@ public class WifiService extends IWifiManager.Stub {
                 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 @@ public class WifiService extends IWifiManager.Stub {
             } 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 @@ public class WifiService extends IWifiManager.Stub {
         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 @@ public class WifiService extends IWifiManager.Stub {
     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 @@ public class WifiService extends IWifiManager.Stub {
                     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 @@ public class WifiService extends IWifiManager.Stub {
             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 bb40967005b6a..73a5435f1254f 100644
--- a/services/java/com/android/server/am/BatteryStatsService.java
+++ b/services/java/com/android/server/am/BatteryStatsService.java
@@ -218,17 +218,17 @@ public final class BatteryStatsService extends IBatteryStats.Stub {
         }
     }
 
-    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 final class BatteryStatsService extends IBatteryStats.Stub {
         }
     }
 
-    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 22dbda3bc7eaa..9d27bde124c72 100644
--- a/wifi/java/android/net/wifi/WifiStateTracker.java
+++ b/wifi/java/android/net/wifi/WifiStateTracker.java
@@ -38,6 +38,7 @@ import android.os.SystemProperties;
 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;
@@ -322,6 +323,21 @@ public class WifiStateTracker extends NetworkStateTracker {
 
     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
@@ -632,12 +648,35 @@ public class WifiStateTracker extends NetworkStateTracker {
         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 @@ public class WifiStateTracker extends NetworkStateTracker {
         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 @@ public class WifiStateTracker extends NetworkStateTracker {
 
             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 @@ public class WifiStateTracker extends NetworkStateTracker {
                     mWM.setWifiEnabled(true);
                     break;
                 }
-                noteRunState();
+                synchronized (this) {
+                    updateBatteryWorkSourceLocked(null);
+                }
                 break;
 
             case EVENT_PASSWORD_KEY_MAY_BE_INCORRECT: