Merge changes I5fd02510,I062f0005 into pi-dev

* changes:
  Change onboarding flow.
  Fix WifiWake locking behavior.

6 files changed, 309 insertions, 79 deletions

diff --git a/service/java/com/android/server/wifi/WakeupConfigStoreData.java b/service/java/com/android/server/wifi/WakeupConfigStoreData.java
index b936a4c6c..d98567766 100644
--- a/service/java/com/android/server/wifi/WakeupConfigStoreData.java
+++ b/service/java/com/android/server/wifi/WakeupConfigStoreData.java
@@ -39,12 +39,14 @@ public class WakeupConfigStoreData implements StoreData {
     private static final String XML_TAG_FEATURE_STATE_SECTION = "FeatureState";
     private static final String XML_TAG_IS_ACTIVE = "IsActive";
     private static final String XML_TAG_IS_ONBOARDED = "IsOnboarded";
+    private static final String XML_TAG_NOTIFICATIONS_SHOWN = "NotificationsShown";
     private static final String XML_TAG_NETWORK_SECTION = "Network";
     private static final String XML_TAG_SSID = "SSID";
     private static final String XML_TAG_SECURITY = "Security";
 
     private final DataSource<Boolean> mIsActiveDataSource;
     private final DataSource<Boolean> mIsOnboardedDataSource;
+    private final DataSource<Integer> mNotificationsDataSource;
     private final DataSource<Set<ScanResultMatchInfo>> mNetworkDataSource;
     private boolean mHasBeenRead = false;
 
@@ -76,9 +78,11 @@ public class WakeupConfigStoreData implements StoreData {
     public WakeupConfigStoreData(
             DataSource<Boolean> isActiveDataSource,
             DataSource<Boolean> isOnboardedDataSource,
+            DataSource<Integer> notificationsDataSource,
             DataSource<Set<ScanResultMatchInfo>> networkDataSource) {
         mIsActiveDataSource = isActiveDataSource;
         mIsOnboardedDataSource = isOnboardedDataSource;
+        mNotificationsDataSource = notificationsDataSource;
         mNetworkDataSource = networkDataSource;
     }
 
@@ -116,6 +120,8 @@ public class WakeupConfigStoreData implements StoreData {
 
         XmlUtil.writeNextValue(out, XML_TAG_IS_ACTIVE, mIsActiveDataSource.getData());
         XmlUtil.writeNextValue(out, XML_TAG_IS_ONBOARDED, mIsOnboardedDataSource.getData());
+        XmlUtil.writeNextValue(out, XML_TAG_NOTIFICATIONS_SHOWN,
+                mNotificationsDataSource.getData());
 
         XmlUtil.writeNextSectionEnd(out, XML_TAG_FEATURE_STATE_SECTION);
     }
@@ -185,6 +191,7 @@ public class WakeupConfigStoreData implements StoreData {
             throws IOException, XmlPullParserException {
         boolean isActive = false;
         boolean isOnboarded = false;
+        int notificationsShown = 0;
 
         while (!XmlUtil.isNextSectionEnd(in, outerTagDepth)) {
             String[] valueName = new String[1];
@@ -199,6 +206,9 @@ public class WakeupConfigStoreData implements StoreData {
                 case XML_TAG_IS_ONBOARDED:
                     isOnboarded = (Boolean) value;
                     break;
+                case XML_TAG_NOTIFICATIONS_SHOWN:
+                    notificationsShown = (Integer) value;
+                    break;
                 default:
                     throw new XmlPullParserException("Unknown value found: " + valueName[0]);
             }
@@ -206,6 +216,7 @@ public class WakeupConfigStoreData implements StoreData {
 
         mIsActiveDataSource.setData(isActive);
         mIsOnboardedDataSource.setData(isOnboarded);
+        mNotificationsDataSource.setData(notificationsShown);
     }
 
     /**
@@ -248,6 +259,7 @@ public class WakeupConfigStoreData implements StoreData {
             mNetworkDataSource.setData(Collections.emptySet());
             mIsActiveDataSource.setData(false);
             mIsOnboardedDataSource.setData(false);
+            mNotificationsDataSource.setData(0);
         }
     }
 
diff --git a/service/java/com/android/server/wifi/WakeupController.java b/service/java/com/android/server/wifi/WakeupController.java
index 9743390ed..e8adbb27c 100644
--- a/service/java/com/android/server/wifi/WakeupController.java
+++ b/service/java/com/android/server/wifi/WakeupController.java
@@ -26,7 +26,6 @@ import android.net.wifi.WifiScanner;
 import android.os.Handler;
 import android.os.Looper;
 import android.provider.Settings;
-import android.util.ArraySet;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -38,6 +37,7 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 
 /**
@@ -127,24 +127,30 @@ public class WakeupController {
         mContentObserver = new ContentObserver(mHandler) {
             @Override
             public void onChange(boolean selfChange) {
-                mWifiWakeupEnabled = mFrameworkFacade.getIntegerSetting(
-                                mContext, Settings.Global.WIFI_WAKEUP_ENABLED, 0) == 1;
-                Log.d(TAG, "WifiWake " + (mWifiWakeupEnabled ? "enabled" : "disabled"));
+                readWifiWakeupEnabledFromSettings();
+                mWakeupOnboarding.setOnboarded();
             }
         };
         mFrameworkFacade.registerContentObserver(mContext, Settings.Global.getUriFor(
                 Settings.Global.WIFI_WAKEUP_ENABLED), true, mContentObserver);
-        mContentObserver.onChange(false /* selfChange */);
+        readWifiWakeupEnabledFromSettings();
 
         // registering the store data here has the effect of reading the persisted value of the
         // data sources after system boot finishes
         mWakeupConfigStoreData = new WakeupConfigStoreData(
                 new IsActiveDataSource(),
-                mWakeupOnboarding.getDataSource(),
+                mWakeupOnboarding.getIsOnboadedDataSource(),
+                mWakeupOnboarding.getNotificationsDataSource(),
                 mWakeupLock.getDataSource());
         wifiConfigStore.registerStoreData(mWakeupConfigStoreData);
     }
 
+    private void readWifiWakeupEnabledFromSettings() {
+        mWifiWakeupEnabled = mFrameworkFacade.getIntegerSetting(
+                mContext, Settings.Global.WIFI_WAKEUP_ENABLED, 0) == 1;
+        Log.d(TAG, "WifiWake " + (mWifiWakeupEnabled ? "enabled" : "disabled"));
+    }
+
     private void setActive(boolean isActive) {
         if (mIsActive != isActive) {
             Log.d(TAG, "Setting active to " + isActive);
@@ -175,14 +181,14 @@ public class WakeupController {
         if (isEnabled()) {
             mWakeupOnboarding.maybeShowNotification();
 
-            Set<ScanResultMatchInfo> mostRecentSavedScanResults = getMostRecentSavedScanResults();
-
+            Set<ScanResultMatchInfo> savedNetworksFromLatestScan = getSavedNetworksFromLatestScan();
             if (mVerboseLoggingEnabled) {
-                Log.d(TAG, "Saved networks in most recent scan:" + mostRecentSavedScanResults);
+                Log.d(TAG, "Saved networks in most recent scan:" + savedNetworksFromLatestScan);
             }
 
-            mWifiWakeMetrics.recordStartEvent(mostRecentSavedScanResults.size());
-            mWakeupLock.initialize(mostRecentSavedScanResults);
+            mWifiWakeMetrics.recordStartEvent(savedNetworksFromLatestScan.size());
+            mWakeupLock.setLock(savedNetworksFromLatestScan);
+            // TODO(b/77291248): request low latency scan here
         }
     }
 
@@ -212,18 +218,30 @@ public class WakeupController {
         mWakeupLock.enableVerboseLogging(mVerboseLoggingEnabled);
     }
 
-    /** Returns a list of saved networks from the last full scan. */
-    private Set<ScanResultMatchInfo> getMostRecentSavedScanResults() {
-        Set<ScanResultMatchInfo> goodSavedNetworks = getGoodSavedNetworks();
+    /** Returns a filtered list of saved networks from the last full scan. */
+    private Set<ScanResultMatchInfo> getSavedNetworksFromLatestScan() {
+        Set<ScanResult> filteredScanResults =
+                filterScanResults(mWifiInjector.getWifiScanner().getSingleScanResults());
+        Set<ScanResultMatchInfo> goodMatchInfos  = toMatchInfos(filteredScanResults);
+        goodMatchInfos.retainAll(getGoodSavedNetworks());
+
+        return goodMatchInfos;
+    }
 
-        List<ScanResult> scanResults = mWifiInjector.getWifiScanner().getSingleScanResults();
-        Set<ScanResultMatchInfo> lastSeenNetworks = new HashSet<>(scanResults.size());
-        for (ScanResult scanResult : scanResults) {
-            lastSeenNetworks.add(ScanResultMatchInfo.fromScanResult(scanResult));
+    /** Returns a set of ScanResults with all DFS channels removed. */
+    private Set<ScanResult> filterScanResults(Collection<ScanResult> scanResults) {
+        int[] dfsChannels = mWifiInjector.getWifiNative()
+                .getChannelsForBand(WifiScanner.WIFI_BAND_5_GHZ_DFS_ONLY);
+        if (dfsChannels == null) {
+            dfsChannels = new int[0];
         }
 
-        lastSeenNetworks.retainAll(goodSavedNetworks);
-        return lastSeenNetworks;
+        final Set<Integer> dfsChannelSet = Arrays.stream(dfsChannels).boxed()
+                .collect(Collectors.toSet());
+
+        return scanResults.stream()
+                .filter(scanResult -> !dfsChannelSet.contains(scanResult.frequency))
+                .collect(Collectors.toSet());
     }
 
     /** Returns a filtered list of saved networks from WifiConfigManager. */
@@ -245,16 +263,18 @@ public class WakeupController {
     }
 
     //TODO(b/69271702) implement WAN filtering
-    private boolean isWideAreaNetwork(WifiConfiguration wifiConfiguration) {
+    private static boolean isWideAreaNetwork(WifiConfiguration config) {
         return false;
     }
 
     /**
      * Handles incoming scan results.
      *
-     * <p>The controller updates the WakeupLock with the incoming scan results. If WakeupLock is
-     * empty, it evaluates scan results for a match with saved networks. If a match exists, it
-     * enables wifi.
+     * <p>The controller updates the WakeupLock with the incoming scan results. If WakeupLock is not
+     * yet fully initialized, it adds the current scanResults to the lock and returns. If WakeupLock
+     * is initialized but not empty, the controller updates the lock with the current scan. If it is
+     * both initialized and empty, it evaluates scan results for a match with saved networks. If a
+     * match exists, it enables wifi.
      *
      * <p>The feature must be enabled and the store data must be loaded in order for the controller
      * to handle scan results.
@@ -269,41 +289,22 @@ public class WakeupController {
 
         // only count scan as handled if isEnabled
         mNumScansHandled++;
-
         if (mVerboseLoggingEnabled) {
-            Log.d(TAG, "Incoming scan. Total scans handled: " + mNumScansHandled);
-            Log.d(TAG, "ScanResults: " + scanResults);
+            Log.d(TAG, "Incoming scan #" + mNumScansHandled);
         }
 
-        // need to show notification here in case user enables Wifi Wake when Wifi is off
+        // need to show notification here in case user turns phone on while wifi is off
         mWakeupOnboarding.maybeShowNotification();
-        if (!mWakeupOnboarding.isOnboarded()) {
-            return;
-        }
-
-        // only update the wakeup lock if it's not already empty
-        if (!mWakeupLock.isEmpty()) {
-            if (mVerboseLoggingEnabled) {
-                Log.d(TAG, "WakeupLock not empty. Updating.");
-            }
 
-            Set<ScanResultMatchInfo> networks = new ArraySet<>();
-            for (ScanResult scanResult : scanResults) {
-                networks.add(ScanResultMatchInfo.fromScanResult(scanResult));
-            }
-            mWakeupLock.update(networks);
-
-            // if wakeup lock is still not empty, return
-            if (!mWakeupLock.isEmpty()) {
-                return;
-            }
+        Set<ScanResult> filteredScanResults = filterScanResults(scanResults);
 
-            Log.d(TAG, "WakeupLock emptied");
-            mWifiWakeMetrics.recordUnlockEvent(mNumScansHandled);
+        mWakeupLock.update(toMatchInfos(filteredScanResults));
+        if (!mWakeupLock.isUnlocked()) {
+            return;
         }
 
         ScanResult network =
-                mWakeupEvaluator.findViableNetwork(scanResults, getGoodSavedNetworks());
+                mWakeupEvaluator.findViableNetwork(filteredScanResults, getGoodSavedNetworks());
 
         if (network != null) {
             Log.d(TAG, "Enabling wifi for network: " + network.SSID);
@@ -312,6 +313,15 @@ public class WakeupController {
     }
 
     /**
+     * Converts ScanResults to ScanResultMatchInfos.
+     */
+    private static Set<ScanResultMatchInfo> toMatchInfos(Collection<ScanResult> scanResults) {
+        return scanResults.stream()
+                .map(ScanResultMatchInfo::fromScanResult)
+                .collect(Collectors.toSet());
+    }
+
+    /**
      * Enables wifi.
      *
      * <p>This method ignores all checks and assumes that {@link WifiStateMachine} is currently
diff --git a/service/java/com/android/server/wifi/WakeupLock.java b/service/java/com/android/server/wifi/WakeupLock.java
index 9e617a498..8a7422f46 100644
--- a/service/java/com/android/server/wifi/WakeupLock.java
+++ b/service/java/com/android/server/wifi/WakeupLock.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wifi;
 
+import android.text.format.DateUtils;
 import android.util.ArrayMap;
 import android.util.Log;
 
@@ -29,7 +30,7 @@ import java.util.Map;
 import java.util.Set;
 
 /**
- * A lock to determine whether Auto Wifi can re-enable Wifi.
+ * A lock to determine whether Wifi Wake can re-enable Wifi.
  *
  * <p>Wakeuplock manages a list of networks to determine whether the device's location has changed.
  */
@@ -39,44 +40,144 @@ public class WakeupLock {
 
     @VisibleForTesting
     static final int CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT = 3;
-
+    @VisibleForTesting
+    static final long MAX_LOCK_TIME_MILLIS = 10 * DateUtils.MINUTE_IN_MILLIS;
 
     private final WifiConfigManager mWifiConfigManager;
     private final Map<ScanResultMatchInfo, Integer> mLockedNetworks = new ArrayMap<>();
+    private final WifiWakeMetrics mWifiWakeMetrics;
+    private final Clock mClock;
+
     private boolean mVerboseLoggingEnabled;
+    private long mLockTimestamp;
+    private boolean mIsInitialized;
+    private int mNumScans;
 
-    public WakeupLock(WifiConfigManager wifiConfigManager) {
+    public WakeupLock(WifiConfigManager wifiConfigManager, WifiWakeMetrics wifiWakeMetrics,
+                      Clock clock) {
         mWifiConfigManager = wifiConfigManager;
+        mWifiWakeMetrics = wifiWakeMetrics;
+        mClock = clock;
     }
 
     /**
-     * Initializes the WakeupLock with the given {@link ScanResultMatchInfo} list.
+     * Sets the WakeupLock with the given {@link ScanResultMatchInfo} list.
      *
-     * <p>This saves the wakeup lock to the store.
+     * <p>This saves the wakeup lock to the store and begins the initialization process.
      *
      * @param scanResultList list of ScanResultMatchInfos to start the lock with
      */
-    public void initialize(Collection<ScanResultMatchInfo> scanResultList) {
+    public void setLock(Collection<ScanResultMatchInfo> scanResultList) {
+        mLockTimestamp = mClock.getElapsedSinceBootMillis();
+        mIsInitialized = false;
+        mNumScans = 0;
+
         mLockedNetworks.clear();
         for (ScanResultMatchInfo scanResultMatchInfo : scanResultList) {
             mLockedNetworks.put(scanResultMatchInfo, CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT);
         }
 
-        Log.d(TAG, "Lock initialized. Number of networks: " + mLockedNetworks.size());
+        Log.d(TAG, "Lock set. Number of networks: " + mLockedNetworks.size());
 
         mWifiConfigManager.saveToStore(false /* forceWrite */);
     }
 
     /**
-     * Updates the lock with the given {@link ScanResultMatchInfo} list.
+     * Maybe sets the WakeupLock as initialized based on total scans handled.
+     *
+     * @param numScans total number of elapsed scans in the current WifiWake session
+     * @return Whether the lock transitioned into its initialized state.
+     */
+    private boolean maybeSetInitializedByScans(int numScans) {
+        if (mIsInitialized) {
+            return false;
+        }
+        boolean shouldBeInitialized = numScans >= CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT;
+        if (shouldBeInitialized) {
+            mIsInitialized = true;
+
+            Log.d(TAG, "Lock initialized by handled scans. Scans: " + numScans);
+            if (mVerboseLoggingEnabled) {
+                Log.d(TAG, "State of lock: " + mLockedNetworks);
+            }
+        }
+        return mIsInitialized;
+    }
+
+    /**
+     * Maybe sets the WakeupLock as initialized based on elapsed time.
+     *
+     * @param timestampMillis current timestamp
+     * @return Whether the lock transitioned into its initialized state.
+     */
+    private boolean maybeSetInitializedByTimeout(long timestampMillis) {
+        if (mIsInitialized) {
+            return false;
+        }
+        long elapsedTime = timestampMillis - mLockTimestamp;
+        boolean shouldBeInitialized = elapsedTime > MAX_LOCK_TIME_MILLIS;
+
+        if (shouldBeInitialized) {
+            mIsInitialized = true;
+
+            Log.d(TAG, "Lock initialized by timeout. Elapsed time: " + elapsedTime);
+            if (mVerboseLoggingEnabled) {
+                Log.d(TAG, "State of lock: " + mLockedNetworks);
+            }
+        }
+        return mIsInitialized;
+    }
+
+    /** Returns whether the lock has been fully initialized. */
+    public boolean isInitialized() {
+        return mIsInitialized;
+    }
+
+    /**
+     * Adds the given networks to the lock.
+     *
+     * <p>This is called during the initialization step.
+     *
+     * @param networkList The list of networks to be added
+     */
+    private void addToLock(Collection<ScanResultMatchInfo> networkList) {
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Initializing lock with networks: " + networkList);
+        }
+
+        boolean hasChanged = false;
+
+        for (ScanResultMatchInfo network : networkList) {
+            if (!mLockedNetworks.containsKey(network)) {
+                mLockedNetworks.put(network, CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT);
+                hasChanged = true;
+            }
+        }
+
+        if (hasChanged) {
+            mWifiConfigManager.saveToStore(false /* forceWrite */);
+        }
+
+        // Set initialized if the lock has handled enough scans, and log the event
+        if (maybeSetInitializedByScans(mNumScans)) {
+            // TODO(easchwar) log metric
+        }
+    }
+
+    /**
+     * Removes networks from the lock if not present in the given {@link ScanResultMatchInfo} list.
      *
      * <p>If a network in the lock is not present in the list, reduce the number of scans
      * required to evict by one. Remove any entries in the list with 0 scans required to evict. If
      * any entries in the lock are removed, the store is updated.
      *
-     * @param scanResultList list of present ScanResultMatchInfos to update the lock with
+     * @param networkList list of present ScanResultMatchInfos to update the lock with
      */
-    public void update(Collection<ScanResultMatchInfo> scanResultList) {
+    private void removeFromLock(Collection<ScanResultMatchInfo> networkList) {
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Filtering lock with networks: " + networkList);
+        }
+
         boolean hasChanged = false;
         Iterator<Map.Entry<ScanResultMatchInfo, Integer>> it =
                 mLockedNetworks.entrySet().iterator();
@@ -84,7 +185,7 @@ public class WakeupLock {
             Map.Entry<ScanResultMatchInfo, Integer> entry = it.next();
 
             // if present in scan list, reset to max
-            if (scanResultList.contains(entry.getKey())) {
+            if (networkList.contains(entry.getKey())) {
                 if (mVerboseLoggingEnabled) {
                     Log.d(TAG, "Found network in lock: " + entry.getKey().networkSsid);
                 }
@@ -104,13 +205,50 @@ public class WakeupLock {
         if (hasChanged) {
             mWifiConfigManager.saveToStore(false /* forceWrite */);
         }
+
+        if (isUnlocked()) {
+            Log.d(TAG, "Lock emptied. Recording unlock event.");
+            mWifiWakeMetrics.recordUnlockEvent(mNumScans);
+        }
     }
 
     /**
-     * Returns whether the internal network set is empty.
+     * Updates the lock with the given {@link ScanResultMatchInfo} list.
+     *
+     * <p>Based on the current initialization state of the lock, either adds or removes networks
+     * from the lock.
+     *
+     * <p>The lock is initialized after {@link #CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT}
+     * scans have been handled, or after {@link #MAX_LOCK_TIME_MILLIS} milliseconds have elapsed
+     * since {@link #setLock(Collection, long)}.
+     *
+     * @param networkList list of present ScanResultMatchInfos to update the lock with
      */
-    public boolean isEmpty() {
-        return mLockedNetworks.isEmpty();
+    public void update(Collection<ScanResultMatchInfo> networkList) {
+        // update is no-op if already unlocked
+        if (isUnlocked()) {
+            return;
+        }
+        mNumScans++;
+
+        // Before checking mIsInitialized, we check to see whether we've exceeded the maximum time
+        // allowed for initialization. If so, we set initialized and treat this scan as a
+        // "removeFromLock()" instead of an "addToLock()".
+        if (maybeSetInitializedByTimeout(mClock.getElapsedSinceBootMillis())) {
+            // TODO(easchwar) log metric
+        }
+
+        // add or remove networks based on initialized status
+        if (mIsInitialized) {
+            removeFromLock(networkList);
+        } else {
+            addToLock(networkList);
+        }
+    }
+
+    /** Returns whether the WakeupLock is unlocked */
+    public boolean isUnlocked() {
+        return mIsInitialized && mLockedNetworks.isEmpty();
     }
 
     /** Returns the data source for the WakeupLock config store data. */
@@ -121,6 +259,8 @@ public class WakeupLock {
     /** Dumps wakeup lock contents. */
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("WakeupLock: ");
+        pw.println("mNumScans: " + mNumScans);
+        pw.println("mIsInitialized: " + mIsInitialized);
         pw.println("Locked networks: " + mLockedNetworks.size());
         for (Map.Entry<ScanResultMatchInfo, Integer> entry : mLockedNetworks.entrySet()) {
             pw.println(entry.getKey() + ", scans to evict: " + entry.getValue());
@@ -146,7 +286,8 @@ public class WakeupLock {
             for (ScanResultMatchInfo network : data) {
                 mLockedNetworks.put(network, CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT);
             }
-
+            // lock is considered initialized if loaded from store
+            mIsInitialized = true;
         }
     }
 }
diff --git a/service/java/com/android/server/wifi/WakeupNotificationFactory.java b/service/java/com/android/server/wifi/WakeupNotificationFactory.java
index 42ae46707..23f31a7db 100644
--- a/service/java/com/android/server/wifi/WakeupNotificationFactory.java
+++ b/service/java/com/android/server/wifi/WakeupNotificationFactory.java
@@ -22,6 +22,7 @@ import android.content.Context;
 import android.content.Intent;
 
 import com.android.internal.R;
+import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
 import com.android.internal.notification.SystemNotificationChannels;
 
 
@@ -37,6 +38,9 @@ public class WakeupNotificationFactory {
     public static final String ACTION_TURN_OFF_WIFI_WAKE =
             "com.android.server.wifi.wakeup.TURN_OFF_WIFI_WAKE";
 
+    /** Notification channel ID for onboarding messages. */
+    public static final int ONBOARD_ID = SystemMessage.NOTE_WIFI_WAKE_ONBOARD;
+
     private final Context mContext;
     private final FrameworkFacade mFrameworkFacade;
 
diff --git a/service/java/com/android/server/wifi/WakeupOnboarding.java b/service/java/com/android/server/wifi/WakeupOnboarding.java
index d4caa0fdb..b6bcbc3c0 100644
--- a/service/java/com/android/server/wifi/WakeupOnboarding.java
+++ b/service/java/com/android/server/wifi/WakeupOnboarding.java
@@ -27,22 +27,31 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.SystemClock;
 import android.provider.Settings;
+import android.text.format.DateUtils;
 import android.util.Log;
 
-import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
+import com.android.internal.annotations.VisibleForTesting;
 
 /**
  * Manages the WiFi Wake onboarding notification.
  *
  * <p>If a user disables wifi with Wifi Wake enabled, this notification is shown to explain that
- * wifi may turn back on automatically. Wifi will not automatically turn back on until after the
- * user interacts with the onboarding notification in some way (e.g. dismiss, tap).
+ * wifi may turn back on automatically. It will be displayed up to 3 times, or until the
+ * user either interacts with the onboarding notification in some way (e.g. dismiss, tap) or
+ * manually enables/disables the feature in WifiSettings.
  */
 public class WakeupOnboarding {
 
     private static final String TAG = "WakeupOnboarding";
 
+    @VisibleForTesting
+    static final int NOTIFICATIONS_UNTIL_ONBOARDED = 3;
+    @VisibleForTesting
+    static final long REQUIRED_NOTIFICATION_DELAY = DateUtils.DAY_IN_MILLIS;
+    private static final long NOT_SHOWN_TIMESTAMP = -1;
+
     private final Context mContext;
     private final WakeupNotificationFactory mWakeupNotificationFactory;
     private NotificationManager mNotificationManager;
@@ -52,6 +61,8 @@ public class WakeupOnboarding {
     private final FrameworkFacade mFrameworkFacade;
 
     private boolean mIsOnboarded;
+    private int mTotalNotificationsShown;
+    private long mLastShownTimestamp = NOT_SHOWN_TIMESTAMP;
     private boolean mIsNotificationShowing;
 
     private final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
@@ -104,17 +115,46 @@ public class WakeupOnboarding {
 
     /** Shows the onboarding notification if applicable. */
     public void maybeShowNotification() {
-        if (isOnboarded() || mIsNotificationShowing) {
+        maybeShowNotification(SystemClock.elapsedRealtime());
+    }
+
+    @VisibleForTesting
+    void maybeShowNotification(long timestamp) {
+        if (!shouldShowNotification(timestamp)) {
             return;
         }
-
         Log.d(TAG, "Showing onboarding notification.");
 
+        incrementTotalNotificationsShown();
+        mIsNotificationShowing = true;
+        mLastShownTimestamp = timestamp;
+
         mContext.registerReceiver(mBroadcastReceiver, mIntentFilter,
                 null /* broadcastPermission */, mHandler);
-        getNotificationManager().notify(SystemMessage.NOTE_WIFI_WAKE_ONBOARD,
+        getNotificationManager().notify(WakeupNotificationFactory.ONBOARD_ID,
                 mWakeupNotificationFactory.createOnboardingNotification());
-        mIsNotificationShowing = true;
+    }
+
+    /**
+     * Increment the total number of shown notifications and onboard the user if reached the
+     * required amount.
+     */
+    private void incrementTotalNotificationsShown() {
+        mTotalNotificationsShown++;
+        if (mTotalNotificationsShown >= NOTIFICATIONS_UNTIL_ONBOARDED) {
+            setOnboarded();
+        } else {
+            mWifiConfigManager.saveToStore(false /* forceWrite */);
+        }
+    }
+
+    private boolean shouldShowNotification(long timestamp) {
+        if (isOnboarded() || mIsNotificationShowing) {
+            return false;
+        }
+
+        return mLastShownTimestamp == NOT_SHOWN_TIMESTAMP
+                || (timestamp - mLastShownTimestamp) > REQUIRED_NOTIFICATION_DELAY;
     }
 
     /** Handles onboarding cleanup on stop. */
@@ -132,11 +172,15 @@ public class WakeupOnboarding {
         }
 
         mContext.unregisterReceiver(mBroadcastReceiver);
-        getNotificationManager().cancel(SystemMessage.NOTE_WIFI_WAKE_ONBOARD);
+        getNotificationManager().cancel(WakeupNotificationFactory.ONBOARD_ID);
         mIsNotificationShowing = false;
     }
 
-    private void setOnboarded() {
+    /** Sets the user as onboarded and persists to store. */
+    public void setOnboarded() {
+        if (mIsOnboarded) {
+            return;
+        }
         Log.d(TAG, "Setting user as onboarded.");
         mIsOnboarded = true;
         mWifiConfigManager.saveToStore(false /* forceWrite */);
@@ -150,12 +194,17 @@ public class WakeupOnboarding {
         return mNotificationManager;
     }
 
-    /** Returns the {@link WakeupConfigStoreData.DataSource} for the {@link WifiConfigStore}. */
-    public WakeupConfigStoreData.DataSource<Boolean> getDataSource() {
-        return new OnboardingDataSource();
+    /** Returns the {@link WakeupConfigStoreData.DataSource} for the onboarded status. */
+    public WakeupConfigStoreData.DataSource<Boolean> getIsOnboadedDataSource() {
+        return new IsOnboardedDataSource();
     }
 
-    private class OnboardingDataSource implements WakeupConfigStoreData.DataSource<Boolean> {
+    /** Returns the {@link WakeupConfigStoreData.DataSource} for the notification status. */
+    public WakeupConfigStoreData.DataSource<Integer> getNotificationsDataSource() {
+        return new NotificationsDataSource();
+    }
+
+    private class IsOnboardedDataSource implements WakeupConfigStoreData.DataSource<Boolean> {
 
         @Override
         public Boolean getData() {
@@ -167,4 +216,17 @@ public class WakeupOnboarding {
             mIsOnboarded = data;
         }
     }
+
+    private class NotificationsDataSource implements WakeupConfigStoreData.DataSource<Integer> {
+
+        @Override
+        public Integer getData() {
+            return mTotalNotificationsShown;
+        }
+
+        @Override
+        public void setData(Integer data) {
+            mTotalNotificationsShown = data;
+        }
+    }
 }
diff --git a/service/java/com/android/server/wifi/WifiInjector.java b/service/java/com/android/server/wifi/WifiInjector.java
index 04bb3f48c..9b2df953c 100644
--- a/service/java/com/android/server/wifi/WifiInjector.java
+++ b/service/java/com/android/server/wifi/WifiInjector.java
@@ -273,7 +273,8 @@ public class WifiInjector {
                 mWifiStateMachineHandlerThread.getLooper(), mFrameworkFacade,
                 wakeupNotificationFactory);
         mWakeupController = new WakeupController(mContext,
-                mWifiStateMachineHandlerThread.getLooper(), new WakeupLock(mWifiConfigManager),
+                mWifiStateMachineHandlerThread.getLooper(),
+                new WakeupLock(mWifiConfigManager, mWifiMetrics.getWakeupMetrics(), mClock),
                 WakeupEvaluator.fromContext(mContext), wakeupOnboarding, mWifiConfigManager,
                 mWifiConfigStore, mWifiMetrics.getWakeupMetrics(), this, mFrameworkFacade);
         mLockManager = new WifiLockManager(mContext, BatteryStatsService.getService());