Acquire wakelock during network stats collection.

Bug: 4517275
Change-Id: I01dfe3b13d7aca84605446e9e9927d0648623552
diff --git a/services/java/com/android/server/net/NetworkStatsService.java b/services/java/com/android/server/net/NetworkStatsService.java
index 7ec6b81..6cc01f4 100644
--- a/services/java/com/android/server/net/NetworkStatsService.java
+++ b/services/java/com/android/server/net/NetworkStatsService.java
@@ -60,6 +60,7 @@
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.INetworkManagementService;
+import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.SystemClock;
 import android.provider.Settings;
@@ -112,6 +113,8 @@
     private final TrustedTime mTime;
     private final NetworkStatsSettings mSettings;
 
+    private final PowerManager.WakeLock mWakeLock;
+
     private IConnectivityManager mConnManager;
 
     // @VisibleForTesting
@@ -191,6 +194,10 @@
         mTime = checkNotNull(time, "missing TrustedTime");
         mSettings = checkNotNull(settings, "missing NetworkStatsSettings");
 
+        final PowerManager powerManager = (PowerManager) context.getSystemService(
+                Context.POWER_SERVICE);
+        mWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
+
         mHandlerThread = new HandlerThread(TAG);
         mHandlerThread.start();
         mHandler = new Handler(mHandlerThread.getLooper());
@@ -408,7 +415,12 @@
             // on background handler thread, and verified CONNECTIVITY_INTERNAL
             // permission above.
             synchronized (mStatsLock) {
-                updateIfacesLocked();
+                mWakeLock.acquire();
+                try {
+                    updateIfacesLocked();
+                } finally {
+                    mWakeLock.release();
+                }
             }
         }
     };
@@ -419,8 +431,12 @@
             // on background handler thread, and verified UPDATE_DEVICE_STATS
             // permission above.
             synchronized (mStatsLock) {
-                // TODO: acquire wakelock while performing poll
-                performPollLocked(true, false);
+                mWakeLock.acquire();
+                try {
+                    performPollLocked(true, false);
+                } finally {
+                    mWakeLock.release();
+                }
             }
         }
     };
@@ -433,7 +449,12 @@
             final int uid = intent.getIntExtra(EXTRA_UID, 0);
             synchronized (mStatsLock) {
                 // TODO: perform one last stats poll for UID
-                removeUidLocked(uid);
+                mWakeLock.acquire();
+                try {
+                    removeUidLocked(uid);
+                } finally {
+                    mWakeLock.release();
+                }
             }
         }
     };