Bug 1124492 - Allow for distribution intent processing to occur after first use. r=margaret

mobile/android/base/GeckoProfile.java

@@ -19,6 +19,7 @@ import java.util.regex.Pattern;
 
 import org.mozilla.gecko.GeckoProfileDirectories.NoMozillaDirectoryException;
 import org.mozilla.gecko.GeckoProfileDirectories.NoSuchProfileException;
+import org.mozilla.gecko.db.BrowserContract;
 import org.mozilla.gecko.db.BrowserDB;
 import org.mozilla.gecko.db.LocalBrowserDB;
 import org.mozilla.gecko.db.StubBrowserDB;
@@ -831,9 +832,14 @@ public final class GeckoProfile {
 
         // Add everything when we're done loading the distribution.
         final Distribution distribution = Distribution.getInstance(context);
-        distribution.addOnDistributionReadyCallback(new Runnable() {
+        distribution.addOnDistributionReadyCallback(new Distribution.ReadyCallback() {
             @Override
-            public void run() {
+            public void distributionNotFound() {
+                this.distributionFound(null);
+            }
+
+            @Override
+            public void distributionFound(Distribution distribution) {
                 Log.d(LOGTAG, "Running post-distribution task: bookmarks.");
 
                 final ContentResolver cr = context.getContentResolver();
@@ -853,10 +859,29 @@ public final class GeckoProfile {
                     // bookmarks as there are favicons, we can also guarantee that
                     // the favicon IDs won't overlap.
                     final LocalBrowserDB db = new LocalBrowserDB(getName());
-                    final int offset = db.addDistributionBookmarks(cr, distribution, 0);
+                    final int offset = distribution == null ? 0 : db.addDistributionBookmarks(cr, distribution, 0);
                     db.addDefaultBookmarks(context, cr, offset);
                 }
             }
+
+            @Override
+            public void distributionArrivedLate(Distribution distribution) {
+                Log.d(LOGTAG, "Running late distribution task: bookmarks.");
+                // Recover as best we can.
+                synchronized (GeckoProfile.this) {
+                    // Skip initialization if the profile directory has been removed.
+                    if (!profileDir.exists()) {
+                        return;
+                    }
+
+                    final LocalBrowserDB db = new LocalBrowserDB(getName());
+                    // We assume we've been called very soon after startup, and so our offset
+                    // into "Mobile Bookmarks" is the number of bookmarks in the DB.
+                    final ContentResolver cr = context.getContentResolver();
+                    final int offset = db.getCount(cr, "bookmarks");
+                    db.addDistributionBookmarks(cr, distribution, offset);
+                }
+            }
         });
     }
 }

mobile/android/base/db/SuggestedSites.java

@@ -288,17 +288,16 @@ public class SuggestedSites {
             return;
         }
 
-        distribution.addOnDistributionReadyCallback(new Runnable() {
+        distribution.addOnDistributionReadyCallback(new Distribution.ReadyCallback() {
             @Override
-            public void run() {
-                Log.d(LOGTAG, "Running post-distribution task: suggested sites.");
-
+            public void distributionNotFound() {
                 // If distribution doesn't exist, simply continue to load
                 // suggested sites directly from resources. See refresh().
-                if (!distribution.exists()) {
-                    return;
-                }
+            }
 
+            @Override
+            public void distributionFound(Distribution distribution) {
+                Log.d(LOGTAG, "Running post-distribution task: suggested sites.");
                 // Merge suggested sites from distribution with the
                 // default ones. Distribution takes precedence.
                 Map<String, Site> sites = loadFromDistribution(distribution);
@@ -320,6 +319,11 @@ public class SuggestedSites {
                 final ContentResolver cr = context.getContentResolver();
                 cr.notifyChange(BrowserContract.SuggestedSites.CONTENT_URI, null);
             }
+
+            @Override
+            public void distributionArrivedLate(Distribution distribution) {
+                distributionFound(distribution);
+            }
         });
     }
 

mobile/android/base/distribution/Distribution.java

@@ -36,12 +36,14 @@ import org.apache.http.protocol.HTTP;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
+import org.mozilla.gecko.AppConstants;
 import org.mozilla.gecko.GeckoAppShell;
 import org.mozilla.gecko.GeckoEvent;
 import org.mozilla.gecko.GeckoSharedPrefs;
 import org.mozilla.gecko.Telemetry;
 import org.mozilla.gecko.mozglue.RobocopTarget;
 import org.mozilla.gecko.util.FileUtils;
+import org.mozilla.gecko.util.HardwareUtils;
 import org.mozilla.gecko.util.ThreadUtils;
 
 import android.app.Activity;
@@ -104,9 +106,29 @@ public class Distribution {
     // Corresponds to the high value in Histograms.json.
     private static final long MAX_DOWNLOAD_TIME_MSEC = 40000;    // 40 seconds.
 
-    // Wait just a little while for the system to send a referrer intent after install.
-    private static final long DELAY_WAIT_FOR_REFERRER_MSEC = 400;
+    // If this is true, ready callbacks that arrive after our state is initially determined
+    // will be queued for delayed running.
+    // This should only be the case on first run, when we're in STATE_NONE.
+    // Implicitly accessed from any non-UI threads via Distribution.doInit, but in practice only one
+    // will actually perform initialization, and "non-UI thread" really means "background thread".
+    private volatile boolean shouldDelayLateCallbacks = false;
 
+    /**
+     * These tasks can be queued to run when a distribution is available.
+     *
+     * If <code>distributionFound</code> is called, it will be the only call.
+     * If <code>distributionNotFound</code> is called, it might be followed by
+     * a call to <code>distributionArrivedLate</code>.
+     *
+     * When <code>distributionNotFound</code> is called,
+     * {@link org.mozilla.gecko.distribution.Distribution#exists()} will return
+     * false. In the other two callbacks, it will return true.
+     */
+    public interface ReadyCallback {
+        void distributionNotFound();
+        void distributionFound(Distribution distribution);
+        void distributionArrivedLate(Distribution distribution);
+    }
 
     /**
      * Used as a drop-off point for ReferrerReceiver. Checked when we process
@@ -123,10 +145,14 @@ public class Distribution {
     private final String packagePath;
     private final String prefsBranch;
 
-    private volatile int state = STATE_UNKNOWN;
+    volatile int state = STATE_UNKNOWN;
     private File distributionDir;
 
-    private final Queue<Runnable> onDistributionReady = new ConcurrentLinkedQueue<Runnable>();
+    private final Queue<ReadyCallback> onDistributionReady = new ConcurrentLinkedQueue<>();
+
+    // Callbacks in this queue have been invoked once as distributionNotFound.
+    // If they're invoked again, it'll be with distributionArrivedLate.
+    private final Queue<ReadyCallback> onLateReady = new ConcurrentLinkedQueue<>();
 
     /**
      * This is a little bit of a bad singleton, because in principle a Distribution
@@ -246,9 +272,51 @@ public class Distribution {
      *
      * @param ref a parsed referrer value from the store-supplied intent.
      */
-    public static void onReceivedReferrer(ReferrerDescriptor ref) {
+    public static void onReceivedReferrer(final Context context, final ReferrerDescriptor ref) {
         // Track the referrer object for distribution handling.
         referrer = ref;
+
+        ThreadUtils.postToBackgroundThread(new Runnable() {
+            @Override
+            public void run() {
+                final Distribution distribution = Distribution.getInstance(context);
+
+                // This will bail if we aren't delayed, or we already have a distribution.
+                distribution.processDelayedReferrer(ref);
+            }
+        });
+    }
+
+    /**
+     * Handle a referrer intent that arrives after first use of the distribution.
+     */
+    private void processDelayedReferrer(final ReferrerDescriptor ref) {
+        ThreadUtils.assertOnBackgroundThread();
+        if (state != STATE_NONE) {
+            return;
+        }
+
+        Log.i(LOGTAG, "Processing delayed referrer.");
+
+        if (!checkIntentDistribution(ref)) {
+            // Oh well. No sense keeping these tasks around.
+            this.onLateReady.clear();
+            return;
+        }
+
+        // Persist our new state.
+        this.state = STATE_SET;
+        getSharedPreferences().edit().putInt(getKeyName(), this.state).apply();
+
+        // Just in case this isn't empty but doInit has finished.
+        runReadyQueue();
+
+        // Now process any tasks that already ran while we were in STATE_NONE
+        // to tell them of our good news.
+        runLateReadyQueue();
+
+        // Make sure that changes to search defaults are applied immediately.
+        GeckoAppShell.sendEventToGecko(GeckoEvent.createBroadcastEvent("Distribution:Changed", ""));
     }
 
     /**
@@ -344,15 +412,11 @@ public class Distribution {
 
         // Bail if we've already tried to initialize the distribution, and
         // there wasn't one.
-        final SharedPreferences settings;
-        if (prefsBranch == null) {
-            settings = GeckoSharedPrefs.forApp(context);
-        } else {
-            settings = context.getSharedPreferences(prefsBranch, Activity.MODE_PRIVATE);
-        }
+        final SharedPreferences settings = getSharedPreferences();
 
-        String keyName = context.getPackageName() + ".distribution_state";
+        final String keyName = getKeyName();
         this.state = settings.getInt(keyName, STATE_UNKNOWN);
+
         if (this.state == STATE_NONE) {
             runReadyQueue();
             return false;
@@ -368,10 +432,14 @@ public class Distribution {
 
         // We try the install intent, then the APK, then the system directory.
         final boolean distributionSet =
-                checkIntentDistribution() ||
+                checkIntentDistribution(referrer) ||
                 checkAPKDistribution() ||
                 checkSystemDistribution();
 
+        // If this is our first run -- and thus we weren't already in STATE_NONE or STATE_SET above --
+        // and we didn't find a distribution already, then we should hold on to callbacks in case we
+        // get a late distribution.
+        this.shouldDelayLateCallbacks = !distributionSet;
         this.state = distributionSet ? STATE_SET : STATE_NONE;
         settings.edit().putInt(keyName, this.state).apply();
 
@@ -385,18 +453,9 @@ public class Distribution {
      *
      * @return true if a referrer-supplied distribution was selected.
      */
-    private boolean checkIntentDistribution() {
+    private boolean checkIntentDistribution(final ReferrerDescriptor referrer) {
         if (referrer == null) {
-            // Wait a predetermined time and try again.
-            // Just block the thread, because it's the simplest solution.
-            try {
-                Thread.sleep(DELAY_WAIT_FOR_REFERRER_MSEC);
-            } catch (InterruptedException e) {
-                // Good enough.
-            }
-            if (referrer == null) {
-                return false;
-            }
+            return false;
         }
 
         URI uri = getReferredDistribution(referrer);
@@ -408,9 +467,21 @@ public class Distribution {
         Log.v(LOGTAG, "Downloading referred distribution: " + uri);
 
         try {
-            HttpURLConnection connection = (HttpURLConnection) uri.toURL().openConnection();
+            final HttpURLConnection connection = (HttpURLConnection) uri.toURL().openConnection();
 
-            connection.setRequestProperty(HTTP.USER_AGENT, GeckoAppShell.getGeckoInterface().getDefaultUAString());
+            // If the Search Activity starts, and we handle the referrer intent, this'll return
+            // null. Recover gracefully in this case.
+            final GeckoAppShell.GeckoInterface geckoInterface = GeckoAppShell.getGeckoInterface();
+            final String ua;
+            if (geckoInterface == null) {
+                // Fall back to GeckoApp's default implementation.
+                ua = HardwareUtils.isTablet() ? AppConstants.USER_AGENT_FENNEC_TABLET :
+                                                AppConstants.USER_AGENT_FENNEC_MOBILE;
+            } else {
+                ua = geckoInterface.getDefaultUAString();
+            }
+
+            connection.setRequestProperty(HTTP.USER_AGENT, ua);
             connection.setRequestProperty("Accept", EXPECTED_CONTENT_TYPE);
 
             try {
@@ -535,18 +606,6 @@ public class Distribution {
         Telemetry.addToHistogram(HISTOGRAM_CODE_CATEGORY, CODE_CATEGORY_FETCH_EXCEPTION);
     }
 
-    /**
-     * Execute tasks that wanted to run when we were done loading
-     * the distribution. These tasks are expected to call {@link #exists()}
-     * to find out whether there's a distribution or not.
-     */
-    private void runReadyQueue() {
-        Runnable task;
-        while ((task = onDistributionReady.poll()) != null) {
-            ThreadUtils.postToBackgroundThread(task);
-        }
-    }
-
     /**
      * @return true if we copied files out of the APK. Sets distributionDir in that case.
      */
@@ -759,21 +818,77 @@ public class Distribution {
     }
 
     /**
-     * The provided <code>Runnable</code> will be queued for execution after
+     * The provided <code>ReadyCallback</code> will be queued for execution after
      * the distribution is ready, or queued for immediate execution if the
      * distribution has already been processed.
      *
-     * Each <code>Runnable</code> will be executed on the background thread.
+     * Each <code>ReadyCallback</code> will be executed on the background thread.
      */
-    public void addOnDistributionReadyCallback(Runnable runnable) {
+    public void addOnDistributionReadyCallback(final ReadyCallback callback) {
         if (state == STATE_UNKNOWN) {
-            this.onDistributionReady.add(runnable);
+            // Queue for later.
+            onDistributionReady.add(callback);
         } else {
-            // If we're already initialized, just queue up the runnable.
-            ThreadUtils.postToBackgroundThread(runnable);
+            invokeCallbackDelayed(callback);
         }
     }
 
+    /**
+     * Run our delayed queue, after a delayed distribution arrives.
+     */
+    private void runLateReadyQueue() {
+        ReadyCallback task;
+        while ((task = onLateReady.poll()) != null) {
+            invokeLateCallbackDelayed(task);
+        }
+    }
+
+    /**
+     * Execute tasks that wanted to run when we were done loading
+     * the distribution.
+     */
+    private void runReadyQueue() {
+        ReadyCallback task;
+        while ((task = onDistributionReady.poll()) != null) {
+            invokeCallbackDelayed(task);
+        }
+    }
+
+    private void invokeLateCallbackDelayed(final ReadyCallback callback) {
+        ThreadUtils.postToBackgroundThread(new Runnable() {
+            @Override
+            public void run() {
+                // Sanity.
+                if (state != STATE_SET) {
+                    Log.w(LOGTAG, "Refusing to invoke late distro callback in state " + state);
+                    return;
+                }
+                callback.distributionArrivedLate(Distribution.this);
+            }
+        });
+    }
+
+    private void invokeCallbackDelayed(final ReadyCallback callback) {
+        ThreadUtils.postToBackgroundThread(new Runnable() {
+            @Override
+            public void run() {
+                switch (state) {
+                    case STATE_SET:
+                        callback.distributionFound(Distribution.this);
+                        break;
+                    case STATE_NONE:
+                        callback.distributionNotFound();
+                        if (shouldDelayLateCallbacks) {
+                            onLateReady.add(callback);
+                        }
+                        break;
+                    default:
+                        throw new IllegalStateException("Expected STATE_NONE or STATE_SET, got " + state);
+                }
+            }
+        });
+    }
+
     /**
      * A safe way for callers to determine if this Distribution instance
      * represents a real live distribution.
@@ -781,4 +896,18 @@ public class Distribution {
     public boolean exists() {
         return state == STATE_SET;
     }
+
+    private String getKeyName() {
+        return context.getPackageName() + ".distribution_state";
+    }
+
+    private SharedPreferences getSharedPreferences() {
+        final SharedPreferences settings;
+        if (prefsBranch == null) {
+            settings = GeckoSharedPrefs.forApp(context);
+        } else {
+            settings = context.getSharedPreferences(prefsBranch, Activity.MODE_PRIVATE);
+        }
+        return settings;
+    }
 }

mobile/android/base/distribution/ReferrerReceiver.java

@@ -49,7 +49,7 @@ public class ReferrerReceiver extends BroadcastReceiver {
 
         // Track the referrer object for distribution handling.
         if (TextUtils.equals(referrer.campaign, DISTRIBUTION_UTM_CAMPAIGN)) {
-            Distribution.onReceivedReferrer(referrer);
+            Distribution.onReceivedReferrer(context, referrer);
         } else {
             Log.d(LOGTAG, "Not downloading distribution: non-matching campaign.");
         }

mobile/android/base/health/BrowserHealthRecorder.java

@@ -531,17 +531,50 @@ public class BrowserHealthRecorder implements HealthRecorder, GeckoEventListener
         // Because the distribution lookup can take some time, do it at the end of
         // our background startup work, along with the Gecko snapshot fetch.
         final Distribution distribution = Distribution.getInstance(context);
-        distribution.addOnDistributionReadyCallback(new Runnable() {
+        distribution.addOnDistributionReadyCallback(new Distribution.ReadyCallback() {
+            private void requestGeckoFields() {
+                Log.d(LOG_TAG, "Requesting all add-ons and FHR prefs from Gecko.");
+                dispatcher.registerGeckoThreadListener(BrowserHealthRecorder.this, EVENT_SNAPSHOT);
+                GeckoAppShell.sendEventToGecko(GeckoEvent.createBroadcastEvent("HealthReport:RequestSnapshot", null));
+            }
+
             @Override
-            public void run() {
+            public void distributionNotFound() {
+                requestGeckoFields();
+            }
+
+            @Override
+            public void distributionFound(Distribution distribution) {
                 Log.d(LOG_TAG, "Running post-distribution task: health recorder.");
                 final DistributionDescriptor desc = distribution.getDescriptor();
                 if (desc != null && desc.valid) {
                     profileCache.setDistributionString(desc.id, desc.version);
                 }
-                Log.d(LOG_TAG, "Requesting all add-ons and FHR prefs from Gecko.");
-                dispatcher.registerGeckoThreadListener(BrowserHealthRecorder.this, EVENT_SNAPSHOT);
-                GeckoAppShell.sendEventToGecko(GeckoEvent.createBroadcastEvent("HealthReport:RequestSnapshot", null));
+                requestGeckoFields();
+            }
+
+            @Override
+            public void distributionArrivedLate(Distribution distribution) {
+                profileCache.beginInitialization();
+
+                final DistributionDescriptor desc = distribution.getDescriptor();
+                if (desc != null && desc.valid) {
+                    profileCache.setDistributionString(desc.id, desc.version);
+                }
+
+                // Now rebuild.
+                try {
+                    profileCache.completeInitialization();
+
+                    if (state == State.INITIALIZING) {
+                        initializeStorage();
+                    } else {
+                        onEnvironmentChanged();
+                    }
+                } catch (Exception e) {
+                    // Well, we tried.
+                    Log.e(LOG_TAG, "Couldn't complete profile cache init.", e);
+                }
             }
         });
     }

mobile/android/base/tests/testDistribution.java

@@ -173,7 +173,22 @@ public class testDistribution extends ContentProviderTest {
         BrowserLocaleManager.storeAndNotifyOSLocale(GeckoSharedPrefs.forProfile(mActivity), locale);
     }
 
-    private void doReferrerTest(String ref, final TestableDistribution distribution, final Runnable distributionReady) throws InterruptedException {
+    private abstract class ExpectNoDistributionCallback implements Distribution.ReadyCallback {
+            @Override
+            public void distributionFound(final Distribution distribution) {
+                mAsserter.ok(false, "No distributionFound.", "Wasn't expecting a distribution!");
+                synchronized (distribution) {
+                    distribution.notifyAll();
+                }
+            }
+
+            @Override
+            public void distributionArrivedLate(final Distribution distribution) {
+                mAsserter.ok(false, "No distributionArrivedLate.", "Wasn't expecting a late distribution!");
+            }
+    }
+
+    private void doReferrerTest(String ref, final TestableDistribution distribution, final Distribution.ReadyCallback distributionReady) throws InterruptedException {
         final Intent intent = new Intent(ACTION_INSTALL_REFERRER);
         intent.setClassName(AppConstants.ANDROID_PACKAGE_NAME, CLASS_REFERRER_RECEIVER);
         intent.putExtra("referrer", ref);
@@ -215,10 +230,10 @@ public class testDistribution extends ContentProviderTest {
         //              --es "referrer" "utm_source=mozilla&utm_medium=testmedium&utm_term=testterm&utm_content=testcontent&utm_campaign=distribution"
         final String ref = "utm_source=mozilla&utm_medium=testmedium&utm_term=testterm&utm_content=testcontent&utm_campaign=distribution";
         final TestableDistribution distribution = new TestableDistribution(mActivity);
-        final Runnable distributionReady = new Runnable() {
+        final Distribution.ReadyCallback distributionReady = new ExpectNoDistributionCallback() {
             @Override
-            public void run() {
-                Log.i(LOGTAG, "Test told distribution is ready.");
+            public void distributionNotFound() {
+                Log.i(LOGTAG, "Test told distribution processing is done.");
                 mAsserter.ok(!distribution.exists(), "Not processed.", "No download because we're offline.");
                 ReferrerDescriptor referrerValue = TestableDistribution.getReferrerDescriptorForTesting();
                 mAsserter.dumpLog("Referrer was " + referrerValue);
@@ -246,9 +261,9 @@ public class testDistribution extends ContentProviderTest {
         //              --es "referrer" "utm_source=mozilla&utm_medium=testmedium&utm_term=testterm&utm_content=testcontent&utm_campaign=testname"
         final String ref = "utm_source=mozilla&utm_medium=testmedium&utm_term=testterm&utm_content=testcontent&utm_campaign=testname";
         final TestableDistribution distribution = new TestableDistribution(mActivity);
-        final Runnable distributionReady = new Runnable() {
+        final Distribution.ReadyCallback distributionReady = new ExpectNoDistributionCallback() {
             @Override
-            public void run() {
+            public void distributionNotFound() {
                 mAsserter.ok(!distribution.exists(), "Not processed.", "No download because campaign was wrong.");
                 ReferrerDescriptor referrerValue = TestableDistribution.getReferrerDescriptorForTesting();
                 mAsserter.is(referrerValue, null, "No referrer.");

mobile/android/base/widget/ActivityChooserModel.java

@@ -1049,15 +1049,13 @@ public class ActivityChooserModel extends DataSetObservable {
             readHistoricalDataFromStream(new FileInputStream(f));
         } catch (FileNotFoundException fnfe) {
             final Distribution dist = Distribution.getInstance(mContext);
-            dist.addOnDistributionReadyCallback(new Runnable() {
+            dist.addOnDistributionReadyCallback(new Distribution.ReadyCallback() {
                 @Override
-                public void run() {
-                    Log.d(LOGTAG, "Running post-distribution task: quickshare.");
-
-                    if (!dist.exists()) {
-                        return;
-                    }
+                public void distributionNotFound() {
+                }
 
+                @Override
+                public void distributionFound(Distribution distribution) {
                     try {
                         File distFile = dist.getDistributionFile("quickshare/" + mHistoryFileName);
                         if (distFile == null) {
@@ -1074,6 +1072,11 @@ public class ActivityChooserModel extends DataSetObservable {
                         return;
                     }
                 }
+
+                @Override
+                public void distributionArrivedLate(Distribution distribution) {
+                    distributionFound(distribution);
+                }
             });
         }
     }

mobile/android/chrome/content/browser.js

@@ -7531,6 +7531,7 @@ var Distribution = {
   _file: null,
 
   init: function dc_init() {
+    Services.obs.addObserver(this, "Distribution:Changed", false);
     Services.obs.addObserver(this, "Distribution:Set", false);
     Services.obs.addObserver(this, "prefservice:after-app-defaults", false);
     Services.obs.addObserver(this, "Campaign:Set", false);
@@ -7544,6 +7545,15 @@ var Distribution = {
 
   observe: function dc_observe(aSubject, aTopic, aData) {
     switch (aTopic) {
+      case "Distribution:Changed":
+        // Re-init the search service.
+        try {
+          Services.search._asyncReInit();
+        } catch (e) {
+          console.log("Unable to reinit search service.");
+        }
+        // Fall through.
+
       case "Distribution:Set":
         // Reload the default prefs so we can observe "prefservice:after-app-defaults"
         Services.prefs.QueryInterface(Ci.nsIObserver).observe(null, "reload-default-prefs", null);

mobile/android/search/java/org/mozilla/search/providers/SearchEngineManager.java

@@ -67,9 +67,9 @@ public class SearchEngineManager implements SharedPreferences.OnSharedPreference
     /**
      * Sets a callback to be called when the default engine changes.
      *
-     * @param callback SearchEngineCallback to be called after the search engine
-     *                 changed. This will run on the UI thread.
-     *                 Note: callback may be called with null engine.
+     * @param changeCallback SearchEngineCallback to be called after the search engine
+     *                       changed. This will run on the UI thread.
+     *                       Note: callback may be called with null engine.
      */
     public void setChangeCallback(SearchEngineCallback changeCallback) {
         this.changeCallback = changeCallback;
@@ -141,9 +141,41 @@ public class SearchEngineManager implements SharedPreferences.OnSharedPreference
      */
     private void getDefaultEngine(final SearchEngineCallback callback) {
         // This runnable is posted to the background thread.
-        distribution.addOnDistributionReadyCallback(new Runnable() {
+        distribution.addOnDistributionReadyCallback(new Distribution.ReadyCallback() {
             @Override
-            public void run() {
+            public void distributionNotFound() {
+                defaultBehavior();
+            }
+
+            @Override
+            public void distributionFound(Distribution distribution) {
+                defaultBehavior();
+            }
+
+            @Override
+            public void distributionArrivedLate(Distribution distribution) {
+                // Let's see if there's a name in the distro.
+                // If so, just this once we'll override the saved value.
+                final String name = getDefaultEngineNameFromDistribution();
+
+                if (name == null) {
+                    return;
+                }
+
+                // Store the default engine name for the future.
+                // Increment an 'ignore' counter so that this preference change
+                // won't cause getDefaultEngine to be called again.
+                ignorePreferenceChange++;
+                GeckoSharedPrefs.forApp(context)
+                        .edit()
+                        .putString(PREF_DEFAULT_ENGINE_KEY, name)
+                        .apply();
+
+                final SearchEngine engine = createEngineFromName(name);
+                runCallback(engine, callback);
+            }
+
+            private void defaultBehavior() {
                 // First look for a default name stored in shared preferences.
                 String name = GeckoSharedPrefs.forApp(context).getString(PREF_DEFAULT_ENGINE_KEY, null);
 
@@ -159,7 +191,7 @@ public class SearchEngineManager implements SharedPreferences.OnSharedPreference
 
                     // Store the default engine name for the future.
                     // Increment an 'ignore' counter so that this preference change
-                    // won'tcause getDefaultEngine to be called again.
+                    // won't cause getDefaultEngine to be called again.
                     ignorePreferenceChange++;
                     GeckoSharedPrefs.forApp(context)
                                     .edit()