Extract watching workspace out of xstate

BrunoQuaresma · BrunoQuaresma · commit 3d96b94496a4 · 2023-11-13T16:36:32.000Z
diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts
@@ -18,10 +18,7 @@ export const workspaceByOwnerAndNameKey = (owner: string, name: string) => [
   "settings",
 ];
 
-export const workspaceByOwnerAndName = (
-  owner: string,
-  name: string,
-): QueryOptions<Workspace> => {
+export const workspaceByOwnerAndName = (owner: string, name: string) => {
   return {
     queryKey: workspaceByOwnerAndNameKey(owner, name),
     queryFn: () =>
diff --git a/site/src/pages/WorkspacePage/Workspace.stories.tsx b/site/src/pages/WorkspacePage/Workspace.stories.tsx
@@ -747,10 +747,3 @@ function makeFailedBuildLogs(): ProvisionerJobLog[] {
     },
   ];
 }
-
-export const UnsupportedWorkspace: Story = {
-  args: {
-    ...Running.args,
-    templateWarnings: ["UNSUPPORTED_WORKSPACES"],
-  },
-};
diff --git a/site/src/pages/WorkspacePage/WorkspacePage.tsx b/site/src/pages/WorkspacePage/WorkspacePage.tsx
@@ -1,6 +1,6 @@
 import { useMachine } from "@xstate/react";
 import { Loader } from "components/Loader/Loader";
-import { FC } from "react";
+import { FC, useEffect, useRef } from "react";
 import { useParams } from "react-router-dom";
 import { workspaceMachine } from "xServices/workspace/workspaceXService";
 import { WorkspaceReadyPage } from "./WorkspaceReadyPage";
@@ -9,14 +9,17 @@ import { ErrorAlert } from "components/Alert/ErrorAlert";
 import { useOrganizationId } from "hooks";
 import { isAxiosError } from "axios";
 import { Margins } from "components/Margins/Margins";
-import { useInfiniteQuery, useQuery } from "react-query";
+import { useInfiniteQuery, useQuery, useQueryClient } from "react-query";
 import { infiniteWorkspaceBuilds } from "api/queries/workspaceBuilds";
 import { templateByName } from "api/queries/templates";
 import { workspaceByOwnerAndName } from "api/queries/workspaces";
 import { checkAuthorization } from "api/queries/authCheck";
 import { WorkspacePermissions, workspaceChecks } from "./permissions";
+import { watchWorkspace } from "api/api";
+import { Workspace } from "api/typesGenerated";
 
 export const WorkspacePage: FC = () => {
+  const queryClient = useQueryClient();
   const params = useParams() as {
     username: string;
     workspace: string;
@@ -37,9 +40,11 @@ export const WorkspacePage: FC = () => {
     },
   });
 
-  const workspaceQuery = useQuery(
-    workspaceByOwnerAndName(username, workspaceName),
+  const workspaceQueryOptions = workspaceByOwnerAndName(
+    username,
+    workspaceName,
   );
+  const workspaceQuery = useQuery(workspaceQueryOptions);
   const workspace = workspaceQuery.data;
 
   const templateQuery = useQuery({
@@ -65,6 +70,43 @@ export const WorkspacePage: FC = () => {
     workspaceQuery.error ?? templateQuery.error ?? permissionsQuery.error;
   const isLoading = !workspace || !template || !permissions;
 
+  // Watch workspace changes
+  const workspaceEventSource = useRef<EventSource | null>(null);
+  useEffect(() => {
+    // If there is an event source, we are already watching the workspace
+    if (!workspace || workspaceEventSource.current) {
+      return;
+    }
+
+    const eventSource = watchWorkspace(workspace.id);
+    workspaceEventSource.current = eventSource;
+
+    eventSource.addEventListener("data", async (event) => {
+      const newWorkspaceData = JSON.parse(event.data) as Workspace;
+      queryClient.setQueryData(
+        workspaceQueryOptions.queryKey,
+        newWorkspaceData,
+      );
+
+      const hasNewBuild =
+        newWorkspaceData.latest_build.id !== workspace.latest_build.id;
+      const lastBuildHasChanged =
+        newWorkspaceData.latest_build.status !== workspace.latest_build.status;
+
+      if (hasNewBuild || lastBuildHasChanged) {
+        await buildsQuery.refetch();
+      }
+    });
+
+    eventSource.addEventListener("error", (event) => {
+      console.error("Error on getting workspace changes.", event);
+    });
+
+    return () => {
+      eventSource.close();
+    };
+  }, [buildsQuery, queryClient, workspace, workspaceQueryOptions.queryKey]);
+
   if (pageError) {
     return (
       <Margins>
diff --git a/site/src/xServices/workspace/workspaceXService.ts b/site/src/xServices/workspace/workspaceXService.ts
@@ -104,35 +104,6 @@ export const workspaceMachine = createMachine(
           },
         },
         states: {
-          listening: {
-            initial: "gettingEvents",
-            states: {
-              gettingEvents: {
-                entry: ["initializeEventSource"],
-                exit: "closeEventSource",
-                invoke: {
-                  src: "listening",
-                  id: "listening",
-                },
-                on: {
-                  REFRESH_WORKSPACE: {
-                    actions: ["refreshWorkspace"],
-                  },
-                  EVENT_SOURCE_ERROR: {
-                    target: "error",
-                  },
-                },
-              },
-              error: {
-                entry: "logWatchWorkspaceWarning",
-                after: {
-                  "2000": {
-                    target: "gettingEvents",
-                  },
-                },
-              },
-            },
-          },
           build: {
             initial: "idle",
             states: {
@@ -359,20 +330,6 @@ export const workspaceMachine = createMachine(
       clearCancellationError: assign({
         cancellationError: (_) => undefined,
       }),
-      // SSE related actions
-      // open a new EventSource so we can stream SSE
-      initializeEventSource: assign({
-        eventSource: (context) =>
-          context.workspace && API.watchWorkspace(context.workspace.id),
-      }),
-      closeEventSource: (context) =>
-        context.eventSource && context.eventSource.close(),
-      refreshWorkspace: assign({
-        workspace: (_, event) => event.data,
-      }),
-      logWatchWorkspaceWarning: (_, event) => {
-        console.error("Watch workspace error:", event);
-      },
       displayActivateError: (_, { data }) => {
         const message = getErrorMessage(data, "Error activate workspace.");
         displayError(message);
@@ -502,44 +459,7 @@ export const workspaceMachine = createMachine(
           throw Error("Cannot activate workspace without workspace id");
         }
       },
-      listening: (context) => (send) => {
-        if (!context.eventSource) {
-          send({ type: "EVENT_SOURCE_ERROR", error: "error initializing sse" });
-          return;
-        }
 
-        context.eventSource.addEventListener("data", (event) => {
-          const newWorkspaceData = JSON.parse(event.data) as TypesGen.Workspace;
-          // refresh our workspace with each SSE
-          send({ type: "REFRESH_WORKSPACE", data: newWorkspaceData });
-
-          const currentWorkspace = context.workspace!;
-          const hasNewBuild =
-            newWorkspaceData.latest_build.id !==
-            currentWorkspace.latest_build.id;
-          const lastBuildHasChanged =
-            newWorkspaceData.latest_build.status !==
-            currentWorkspace.latest_build.status;
-
-          if (hasNewBuild || lastBuildHasChanged) {
-            send({ type: "REFRESH_TIMELINE" });
-          }
-        });
-
-        // handle any error events returned by our sse
-        context.eventSource.addEventListener("error", (event) => {
-          send({ type: "EVENT_SOURCE_ERROR", error: event });
-        });
-
-        // handle any sse implementation exceptions
-        context.eventSource.onerror = () => {
-          send({ type: "EVENT_SOURCE_ERROR", error: "sse error" });
-        };
-
-        return () => {
-          context.eventSource?.close();
-        };
-      },
     },
   },
 );
