improve performance of auth invalidation

Parkreiner · Parkreiner · commit 74f6cd906371 · 2024-06-03T14:59:14.000Z
diff --git a/plugins/backstage-plugin-coder/src/components/CoderProvider/CoderAuthProvider.tsx b/plugins/backstage-plugin-coder/src/components/CoderProvider/CoderAuthProvider.tsx
@@ -27,6 +27,7 @@ import {
 import { coderClientApiRef } from '../../api/CoderClient';
 import { CoderLogo } from '../CoderLogo';
 import { CoderAuthFormDialog } from '../CoderAuthFormDialog';
+import { AxiosError } from 'axios';
 
 const BACKSTAGE_APP_ROOT_ID = '#root';
 const FALLBACK_UI_OVERRIDE_CLASS_NAME = 'backstage-root-override';
@@ -143,30 +144,30 @@ function useAuthState(): CoderAuth {
   // outside React because we let the user connect their own queryClient
   const queryClient = useQueryClient();
   useEffect(() => {
-    // Pseudo-mutex; makes sure that if we get a bunch of errors, only one
-    // revalidation will be processed at a time
-    let isRevalidating = false;
-
-    const revalidateTokenOnError = async (event: QueryCacheNotifyEvent) => {
-      const queryError = event.query.state.error;
+    if (!isAuthenticated) {
+      return undefined;
+    }
 
-      const shouldRevalidate =
-        isAuthenticated &&
-        !isRevalidating &&
-        BackstageHttpError.isInstance(queryError) &&
-        queryError.status === 401;
+    const onCoderQueryError = (event: QueryCacheNotifyEvent) => {
+      const queryKey = event.query.queryKey;
+      const isCoderQuery =
+        Array.isArray(queryKey) && queryKey[0] === CODER_QUERY_KEY_PREFIX;
 
-      if (!shouldRevalidate) {
+      if (!isCoderQuery) {
         return;
       }
 
-      isRevalidating = true;
-      await queryClient.refetchQueries({ queryKey: sharedAuthQueryKey });
-      isRevalidating = false;
+      const queryError = event.query.state.error;
+      const failedBecauseOfInvalidAuth =
+        queryError instanceof AxiosError && queryError.status === 401;
+
+      if (failedBecauseOfInvalidAuth) {
+        queryClient.setQueryData(sharedAuthQueryKey, false);
+      }
     };
 
     const queryCache = queryClient.getQueryCache();
-    const unsubscribe = queryCache.subscribe(revalidateTokenOnError);
+    const unsubscribe = queryCache.subscribe(onCoderQueryError);
     return unsubscribe;
   }, [queryClient, isAuthenticated]);
 
