diff --git a/services/core/java/com/android/server/DeviceIdleController.java b/services/core/java/com/android/server/DeviceIdleController.java
index 8be859f..0c2cf60 100644
--- a/services/core/java/com/android/server/DeviceIdleController.java
+++ b/services/core/java/com/android/server/DeviceIdleController.java
@@ -85,6 +85,7 @@
 import com.android.internal.util.FastXmlSerializer;
 import com.android.internal.util.XmlUtils;
 import com.android.server.am.BatteryStatsService;
+import com.android.server.net.NetworkPolicyManagerInternal;
 
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
@@ -303,10 +304,7 @@
     private final SparseArray<Pair<MutableLong, String>> mTempWhitelistAppIdEndTimes
             = new SparseArray<>();
 
-    /**
-     * Callback to the NetworkPolicyManagerService to tell it that the temp whitelist has changed.
-     */
-    Runnable mNetworkPolicyTempWhitelistCallback;
+    private NetworkPolicyManagerInternal mNetworkPolicyManagerInternal;
 
     /**
      * Current app IDs of temporarily whitelist apps for high-priority messages.
@@ -1018,6 +1016,7 @@
     private static final int MSG_TEMP_APP_WHITELIST_TIMEOUT = 6;
     private static final int MSG_REPORT_MAINTENANCE_ACTIVITY = 7;
     private static final int MSG_FINISH_IDLE_OP = 8;
+    private static final int MSG_REPORT_TEMP_APP_WHITELIST_CHANGED = 9;
 
     final class MyHandler extends Handler {
         MyHandler(Looper looper) {
@@ -1133,6 +1132,11 @@
                     // mActiveIdleWakeLock is held at this point
                     decActiveIdleOps();
                 } break;
+                case MSG_REPORT_TEMP_APP_WHITELIST_CHANGED: {
+                    final int appId = msg.arg1;
+                    final boolean added = (msg.arg2 == 1);
+                    mNetworkPolicyManagerInternal.onTempPowerSaveWhitelistChange(appId, added);
+                } break;
             }
         }
     }
@@ -1283,10 +1287,6 @@
             return mConstants.NOTIFICATION_WHITELIST_DURATION;
         }
 
-        public void setNetworkPolicyTempWhitelistCallback(Runnable callback) {
-            setNetworkPolicyTempWhitelistCallbackInternal(callback);
-        }
-
         public void setJobsActive(boolean active) {
             DeviceIdleController.this.setJobsActive(active);
         }
@@ -1411,6 +1411,7 @@
                 mLocalAlarmManager = getLocalService(AlarmManagerService.LocalService.class);
                 mNetworkPolicyManager = INetworkPolicyManager.Stub.asInterface(
                         ServiceManager.getService(Context.NETWORK_POLICY_SERVICE));
+                mNetworkPolicyManagerInternal = getLocalService(NetworkPolicyManagerInternal.class);
                 mDisplayManager = (DisplayManager) getContext().getSystemService(
                         Context.DISPLAY_SERVICE);
                 mSensorManager = (SensorManager) getContext().getSystemService(Context.SENSOR_SERVICE);
@@ -1699,7 +1700,7 @@
     void addPowerSaveTempWhitelistAppDirectInternal(int callingUid, int appId,
             long duration, boolean sync, String reason) {
         final long timeNow = SystemClock.elapsedRealtime();
-        Runnable networkPolicyTempWhitelistCallback = null;
+        boolean informWhitelistChanged = false;
         synchronized (this) {
             int callingAppId = UserHandle.getAppId(callingUid);
             if (callingAppId >= Process.FIRST_APPLICATION_UID) {
@@ -1729,24 +1730,17 @@
                 }
                 postTempActiveTimeoutMessage(appId, duration);
                 updateTempWhitelistAppIdsLocked(appId, true);
-                if (mNetworkPolicyTempWhitelistCallback != null) {
-                    if (!sync) {
-                        mHandler.post(mNetworkPolicyTempWhitelistCallback);
-                    } else {
-                        networkPolicyTempWhitelistCallback = mNetworkPolicyTempWhitelistCallback;
-                    }
+                if (sync) {
+                    informWhitelistChanged = true;
+                } else {
+                    mHandler.obtainMessage(MSG_REPORT_TEMP_APP_WHITELIST_CHANGED, appId, 1)
+                            .sendToTarget();
                 }
                 reportTempWhitelistChangedLocked();
             }
         }
-        if (networkPolicyTempWhitelistCallback != null) {
-            networkPolicyTempWhitelistCallback.run();
-        }
-    }
-
-    public void setNetworkPolicyTempWhitelistCallbackInternal(Runnable callback) {
-        synchronized (this) {
-            mNetworkPolicyTempWhitelistCallback = callback;
+        if (informWhitelistChanged) {
+            mNetworkPolicyManagerInternal.onTempPowerSaveWhitelistChange(appId, true);
         }
     }
 
@@ -1775,9 +1769,8 @@
                     Slog.d(TAG, "Removing UID " + uid + " from temp whitelist");
                 }
                 updateTempWhitelistAppIdsLocked(uid, false);
-                if (mNetworkPolicyTempWhitelistCallback != null) {
-                    mHandler.post(mNetworkPolicyTempWhitelistCallback);
-                }
+                mHandler.obtainMessage(MSG_REPORT_TEMP_APP_WHITELIST_CHANGED, uid, 0)
+                        .sendToTarget();
                 reportTempWhitelistChangedLocked();
                 try {
                     mBatteryStats.noteEvent(BatteryStats.HistoryItem.EVENT_TEMP_WHITELIST_FINISH,
diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerInternal.java b/services/core/java/com/android/server/net/NetworkPolicyManagerInternal.java
index dc2ebb4..7934a96 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerInternal.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerInternal.java
@@ -38,4 +38,13 @@
      * to current networking policies.
      */
     public abstract boolean isUidNetworkingBlocked(int uid, String ifname);
+
+    /**
+     * Informs that an appId has been added or removed from the temp-powersave-whitelist so that
+     * that network rules for that appId can be updated.
+     *
+     * @param appId The appId which has been updated in the whitelist.
+     * @param added Denotes whether the {@param appId} has been added or removed from the whitelist.
+     */
+    public abstract void onTempPowerSaveWhitelistChange(int appId, boolean added);
 }
diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
index f70486a..32513db 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -182,7 +182,6 @@
 import com.android.internal.util.FastXmlSerializer;
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.Preconditions;
-import com.android.server.DeviceIdleController;
 import com.android.server.EventLogTags;
 import com.android.server.LocalServices;
 import com.android.server.ServiceThread;
@@ -603,36 +602,6 @@
         return changed;
     }
 
-    void updatePowerSaveTempWhitelistUL() {
-        try {
-            // Clear the states of the current whitelist
-            final int N = mPowerSaveTempWhitelistAppIds.size();
-            for (int i = 0; i < N; i++) {
-                mPowerSaveTempWhitelistAppIds.setValueAt(i, false);
-            }
-            // Update the states with the new whitelist
-            final int[] whitelist = mDeviceIdleController.getAppIdTempWhitelist();
-            if (whitelist != null) {
-                for (int uid : whitelist) {
-                    mPowerSaveTempWhitelistAppIds.put(uid, true);
-                }
-            }
-        } catch (RemoteException e) {
-        }
-    }
-
-    /**
-     * Remove unnecessary entries in the temp whitelist
-     */
-    void purgePowerSaveTempWhitelistUL() {
-        final int N = mPowerSaveTempWhitelistAppIds.size();
-        for (int i = N - 1; i >= 0; i--) {
-            if (mPowerSaveTempWhitelistAppIds.valueAt(i) == false) {
-                mPowerSaveTempWhitelistAppIds.removeAt(i);
-            }
-        }
-    }
-
     private void initService(CountDownLatch initCompleteSignal) {
         Trace.traceBegin(Trace.TRACE_TAG_NETWORK, "systemReady");
         final int oldPriority = Process.getThreadPriority(Process.myTid());
@@ -730,10 +699,6 @@
                     PowerManager.ACTION_POWER_SAVE_WHITELIST_CHANGED);
             mContext.registerReceiver(mPowerSaveWhitelistReceiver, whitelistFilter, null, mHandler);
 
-            DeviceIdleController.LocalService deviceIdleService
-                    = LocalServices.getService(DeviceIdleController.LocalService.class);
-            deviceIdleService.setNetworkPolicyTempWhitelistCallback(mTempPowerSaveChangedCallback);
-
             // watch for network interfaces to be claimed
             final IntentFilter connFilter = new IntentFilter(CONNECTIVITY_ACTION);
             mContext.registerReceiver(mConnReceiver, connFilter, CONNECTIVITY_INTERNAL, mHandler);
@@ -838,17 +803,6 @@
         }
     };
 
-    final private Runnable mTempPowerSaveChangedCallback = new Runnable() {
-        @Override
-        public void run() {
-            synchronized (mUidRulesFirstLock) {
-                updatePowerSaveTempWhitelistUL();
-                updateRulesForTempWhitelistChangeUL();
-                purgePowerSaveTempWhitelistUL();
-            }
-        }
-    };
-
     final private BroadcastReceiver mPackageReceiver = new BroadcastReceiver() {
         @Override
         public void onReceive(Context context, Intent intent) {
@@ -3292,20 +3246,18 @@
         }
     }
 
-    private void updateRulesForTempWhitelistChangeUL() {
+    private void updateRulesForTempWhitelistChangeUL(int appId) {
         final List<UserInfo> users = mUserManager.getUsers();
-        for (int i = 0; i < users.size(); i++) {
+        final int numUsers = users.size();
+        for (int i = 0; i < numUsers; i++) {
             final UserInfo user = users.get(i);
-            for (int j = mPowerSaveTempWhitelistAppIds.size() - 1; j >= 0; j--) {
-                int appId = mPowerSaveTempWhitelistAppIds.keyAt(j);
-                int uid = UserHandle.getUid(user.id, appId);
-                // Update external firewall rules.
-                updateRuleForAppIdleUL(uid);
-                updateRuleForDeviceIdleUL(uid);
-                updateRuleForRestrictPowerUL(uid);
-                // Update internal rules.
-                updateRulesForPowerRestrictionsUL(uid);
-            }
+            int uid = UserHandle.getUid(user.id, appId);
+            // Update external firewall rules.
+            updateRuleForAppIdleUL(uid);
+            updateRuleForDeviceIdleUL(uid);
+            updateRuleForRestrictPowerUL(uid);
+            // Update internal rules.
+            updateRulesForPowerRestrictionsUL(uid);
         }
     }
 
@@ -4286,6 +4238,18 @@
             if (LOGV) logUidStatus(uid, "allowed by default");
             return false;
         }
+
+        @Override
+        public void onTempPowerSaveWhitelistChange(int appId, boolean added) {
+            synchronized (mUidRulesFirstLock) {
+                if (added) {
+                    mPowerSaveTempWhitelistAppIds.put(appId, true);
+                } else {
+                    mPowerSaveTempWhitelistAppIds.delete(appId);
+                }
+                updateRulesForTempWhitelistChangeUL(appId);
+            }
+        }
     }
 
     private static boolean hasRule(int uidRules, int rule) {