From 5b2f23fd6dda0e57ce3d82e3a073329c0a43f9b9 Mon Sep 17 00:00:00 2001
From: Nick Lanam <314133+NickLanam@users.noreply.github.com>
Date: Mon, 12 Jun 2023 17:14:13 -0700
Subject: [PATCH] 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>
---
 src/ui/src/api/api-context.tsx     | 25 ++++++++++++++++++++++++-
 src/ui/src/api/api-manager.ts      |  2 ++
 src/ui/src/containers/App/live.tsx |  3 ++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/ui/src/api/api-context.tsx b/src/ui/src/api/api-context.tsx
index e3b3a22c28f..508b42825b6 100644
--- a/src/ui/src/api/api-context.tsx
+++ b/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}>
diff --git a/src/ui/src/api/api-manager.ts b/src/ui/src/api/api-manager.ts
index 70a4af9a372..afae5e62759 100644
--- a/src/ui/src/api/api-manager.ts
+++ b/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; }
diff --git a/src/ui/src/containers/App/live.tsx b/src/ui/src/containers/App/live.tsx
index 66ba531db3b..c0280272e9f 100644
--- a/src/ui/src/containers/App/live.tsx
+++ b/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);
   }