Fix an update issue caused by React being out of scope at a bad time (#1468)

Summary: Moving PixieAPIManager out of React's scope was a risky move.
As it turns out, PixieAPIContext wasn't catching one of the updates
during the embed authentication procedure because of this.
By implementing an unholy hack to tell React when this happens, the
`authorized` network call still fires with a bearer token.

Type of change: /kind bugfix

Test Plan: Try to embed Pixie using `embedPixieToken` to authenticate.
Before, it would give up trying right before the `postMessage` comes
through. After, it should try once more as soon as the auth token is
provided to Pixie.

---------

Signed-off-by: Nick Lanam <nlanam@pixielabs.ai>

Author: NickLanam

src/ui/src/api/api-context.tsx

@@ -21,27 +21,50 @@ import * as React from 'react';
 import { ApolloProvider } from '@apollo/client/react';
 
 import { AuthContext } from 'app/common/auth-context';
+import { SetStateFunc } from 'app/context/common';
 import { WithChildren } from 'app/utils/react-boilerplate';
 
 import { PixieAPIClient, PixieAPIClientAbstract } from './api-client';
 import { PixieAPIManager } from './api-manager';
 import { PixieAPIClientOptions } from './api-options';
 
-
 export const PixieAPIContext = React.createContext<PixieAPIClientAbstract>(null);
 PixieAPIContext.displayName = 'PixieAPIContext';
 
 export type PixieAPIContextProviderProps = WithChildren<PixieAPIClientOptions>;
 
+declare global {
+  interface Window {
+    setApiContextUpdatesFromOutsideReact?: SetStateFunc<number>;
+  }
+}
+
 export const PixieAPIContextProvider: React.FC<PixieAPIContextProviderProps> = React.memo(({ children, ...opts }) => {
   const { authToken } = React.useContext(AuthContext);
 
+  // PixieAPIManager exists outside of React's scope. If it replaces its instance, we'll never know.
+  // By putting a setState function somewhere that PixieAPIManager can reach, we can force the update from there.
+  // This is not a typical or conventional thing to need to do, so please don't do this anywhere else.
+  const [updateFromOutsideReact, setUpdateFromOutsideReact] = React.useState(0);
+  React.useEffect(() => {
+    // Technically this means PixieAPIContextProvider has to be singleton - and it already is, in practice.
+    if (!window.setApiContextUpdatesFromOutsideReact) {
+      window.setApiContextUpdatesFromOutsideReact = setUpdateFromOutsideReact;
+    }
+    return () => {
+      delete window.setApiContextUpdatesFromOutsideReact;
+    };
+  }, []);
+  React.useEffect(() => { /* Just need to update this context, nothing more */ }, [updateFromOutsideReact]);
+
   // PixieAPIManager already reinitializes the API client when options change, making this context just a wrapper.
   React.useEffect(() => { PixieAPIManager.uri = opts.uri; }, [opts.uri]);
   React.useEffect(() => { PixieAPIManager.authToken = authToken; }, [authToken]);
   React.useEffect(() => { PixieAPIManager.onUnauthorized = opts.onUnauthorized; }, [opts.onUnauthorized]);
 
   const instance = PixieAPIManager.instance as PixieAPIClient;
+  const gqlClient = instance.getCloudClient().graphQL;
+  React.useEffect(() => { /* Similar to the above trick, must re-render to update ApolloProvider */ }, [gqlClient]);
 
   return !instance ? null : (
     <PixieAPIContext.Provider value={instance}>

src/ui/src/api/api-manager.ts

@@ -46,6 +46,8 @@ export class PixieAPIManager {
     if (PixieAPIManager._authToken != null) opts.authToken = PixieAPIManager._authToken;
     if (PixieAPIManager._onUnauthorized != null) opts.onUnauthorized = PixieAPIManager._onUnauthorized;
     PixieAPIManager._instance = PixieAPIClient.create(opts);
+    // See api-context.tsx for why this exists
+    if (window.setApiContextUpdatesFromOutsideReact) window.setApiContextUpdatesFromOutsideReact((prev) => prev + 1);
   }
 
   public static get uri() { return PixieAPIManager._uri; }

src/ui/src/containers/App/live.tsx

@@ -283,7 +283,8 @@ export default function PixieWithContext(): React.ReactElement {
   if (errMsg) {
     // This is an error with pixie cloud, it is probably not relevant to the user.
     // Show a generic error message instead.
-    showSnackbar({ message: 'There was a problem connecting to Pixie', autoHideDuration: 5000 });
+    // Wait one update cycle, since this can happen in the middle of updating other components in some cases
+    setTimeout(() => showSnackbar({ message: 'There was a problem connecting to Pixie', autoHideDuration: 5000 }));
     // eslint-disable-next-line no-console
     console.error(errMsg);
   }