Streamline from-browser initialization

src/actions/app.js

@@ -45,6 +45,7 @@ import type {
   UploadedProfileInformation,
 } from 'firefox-profiler/types';
 import type { TabSlug } from 'firefox-profiler/app-logic/tabs-handling';
+import type { BrowserConnection } from 'firefox-profiler/app-logic/browser-connection';
 
 export function changeSelectedTab(selectedTab: TabSlug): ThunkAction<void> {
   return (dispatch, getState) => {
@@ -121,7 +122,8 @@ export function setHasZoomedViaMousewheel() {
  */
 export function setupInitialUrlState(
   location: Location,
-  profile: Profile | null
+  profile: Profile | null,
+  browserConnection: BrowserConnection | null
 ): ThunkAction<void> {
   return (dispatch) => {
     let urlState;
@@ -157,7 +159,7 @@ export function setupInitialUrlState(
     // load process.
     withHistoryReplaceStateSync(() => {
       dispatch(updateUrlState(urlState));
-      dispatch(finalizeProfileView());
+      dispatch(finalizeProfileView(browserConnection));
       dispatch(urlSetupDone());
     });
   };

src/actions/receive-profile.js

@@ -965,7 +965,9 @@ export async function doSymbolicateProfile(
   dispatch(doneSymbolicating());
 }
 
-export function retrieveProfileFromBrowser(): ThunkAction<Promise<void>> {
+export function retrieveProfileFromBrowser(
+  initialLoad: boolean = false
+): ThunkAction<Promise<BrowserConnection | null>> {
   return async (dispatch) => {
     try {
       // Attempt to establish a connection to the browser.
@@ -1017,10 +1019,12 @@ export function retrieveProfileFromBrowser(): ThunkAction<Promise<void>> {
         rawGeckoProfile
       );
       const profile = processGeckoProfile(unpackedProfile);
-      await dispatch(loadProfile(profile, { browserConnection }));
+      await dispatch(loadProfile(profile, { browserConnection }, initialLoad));
+      return browserConnection;
     } catch (error) {
       dispatch(fatalError(error));
       console.error(error);
+      return null;
     }
   };
 }
@@ -1508,9 +1512,12 @@ export function retrieveProfilesToCompare(
 // and loads the profile in that given location, then returns the profile data.
 // This function is being used to get the initial profile data before upgrading
 // the url and processing the UrlState.
-export function retrieveProfileForRawUrl(
-  location: Location
-): ThunkAction<Promise<Profile | null>> {
+export function retrieveProfileForRawUrl(location: Location): ThunkAction<
+  Promise<{|
+    profile: Profile | null,
+    browserConnection: BrowserConnection | null,
+  |}>
+> {
   return async (dispatch, getState) => {
     const pathParts = location.pathname.split('/').filter((d) => d);
     let possibleDataSource = pathParts[0];
@@ -1528,14 +1535,11 @@ export function retrieveProfileForRawUrl(
     }
     dispatch(setDataSource(dataSource));
 
+    let browserConnection = null;
+
     switch (dataSource) {
       case 'from-browser':
-        // We don't need to `await` the result because there's no url upgrading
-        // when retrieving the profile from the browser and we don't need to wait
-        // for the process. Moreover we don't want to wait for the end of
-        // symbolication and rather want to show the UI as soon as we get
-        // the profile data.
-        dispatch(retrieveProfileFromBrowser());
+        browserConnection = await dispatch(retrieveProfileFromBrowser(true));
         break;
       case 'public':
         await dispatch(retrieveProfileFromStore(pathParts[1], true));
@@ -1568,8 +1572,10 @@ export function retrieveProfileForRawUrl(
         );
     }
 
-    // Profile may be null only for the `from-browser` dataSource since we do
-    // not `await` for retrieveProfileFromBrowser function.
-    return getProfileOrNull(getState());
+    // Profile may be null if the response was a zip file.
+    return {
+      profile: getProfileOrNull(getState()),
+      browserConnection,
+    };
   };
 }

src/components/app/UrlManager.js

@@ -95,17 +95,21 @@ class UrlManagerImpl extends React.PureComponent<Props> {
 
     try {
       // Process the raw url and fetch the profile.
-      // We try to fetch the profile before setting the url state, because
-      // while processing and especially upgrading the url information we may
-      // need the profile data.
+      // We try to fetch the profile before setting the url state, because we
+      // may need the profile data during URL upgrading.
       //
-      // Also note the profile may be null for the `from-browser` dataSource since
-      // we do not `await` for retrieveProfileFromBrowser function, but also in
-      // case of fatal errors in the process of retrieving and processing a
-      // profile. To handle the latter case properly, we won't `pushState` if
-      // we're in a FATAL_ERROR state.
-      const profile = await retrieveProfileForRawUrl(window.location);
-      setupInitialUrlState(window.location, profile);
+      // In some cases, the returned profile will be null:
+      //  - If the response is a zip file (even if the URL tells us which file
+      //    to pick from the zip file).
+      //  - If a fatal error was encountered in the process of retrieving and
+      //    processing the profile.
+      //
+      // To handle the latter case properly, we won't `pushState` if we're in
+      // a FATAL_ERROR state.
+      const { profile, browserConnection } = await retrieveProfileForRawUrl(
+        window.location
+      );
+      setupInitialUrlState(window.location, profile, browserConnection);
     } catch (error) {
       // Complete the URL setup, as values can come from the user, so we should
       // still proceed with loading the app.

src/test/components/UrlManager.test.js

@@ -31,6 +31,7 @@ import {
 jest.mock('../../profile-logic/symbol-store');
 
 import { TextEncoder, TextDecoder } from 'util';
+import { simulateOldWebChannelAndFrameScript } from '../fixtures/mocks/web-channel';
 
 describe('UrlManager', function () {
   autoMockFullNavigation();
@@ -87,8 +88,8 @@ describe('UrlManager', function () {
     window.fetch = jest
       .fn()
       .mockRejectedValue(new Error('Simulated network error'));
-    window.geckoProfilerPromise = Promise.resolve(geckoProfiler);
     window.TextDecoder = TextDecoder;
+    simulateOldWebChannelAndFrameScript(geckoProfiler);
   });
 
   afterEach(function () {

src/test/store/receive-profile.test.js

@@ -2035,7 +2035,9 @@ describe('actions/receive-profile', function () {
       jest.spyOn(console, 'warn').mockImplementation(() => {});
 
       const store = blankStore();
-      await store.dispatch(retrieveProfileForRawUrl(location));
+      const { browserConnection } = await store.dispatch(
+        retrieveProfileForRawUrl(location)
+      );
 
       // To find stupid mistakes more easily, check that we didn't get a fatal
       // error here. If we got one, let's rethrow the error.
@@ -2059,6 +2061,7 @@ describe('actions/receive-profile', function () {
         geckoProfile,
         waitUntilPhase,
         waitUntilSymbolication,
+        browserConnection,
         ...store,
       };
     }
@@ -2147,19 +2150,25 @@ describe('actions/receive-profile', function () {
     });
 
     it('retrieves profile from a `from-browser` data source and loads it', async function () {
-      const { geckoProfile, getState, waitUntilPhase } = await setup(
-        {
-          pathname: '/from-browser/',
-          search: '',
-          hash: '',
-        },
-        0
-      );
+      const { geckoProfile, getState, dispatch, browserConnection } =
+        await setup(
+          {
+            pathname: '/from-browser/',
+            search: '',
+            hash: '',
+          },
+          0
+        );
+
+      // Check if we loaded the profile data successfully.
+      expect(getView(getState()).phase).toBe('PROFILE_LOADED');
+
+      // Check if we can successfully finalize the profile view.
+      await dispatch(finalizeProfileView(browserConnection));
+      expect(getView(getState()).phase).toBe('DATA_LOADED');
 
-      // Differently, `from-browser` calls the finalizeProfileView internally,
-      // we don't need to call it again.
-      await waitUntilPhase('DATA_LOADED');
       const processedProfile = processGeckoProfile(geckoProfile);
+      processedProfile.meta.symbolicated = true;
       expect(ProfileViewSelectors.getProfile(getState())).toEqual(
         processedProfile
       );