fix: improve race condition connectivity handling (#2342)

sanity-io · Dec 21, 2024 · a98d6cf · a98d6cf
1 parent 20a46b5
commit a98d6cf
Show file tree

Hide file tree

Showing 7 changed files with 48 additions and 42 deletions.
diff --git a/packages/comlink/src/node.ts b/packages/comlink/src/node.ts
@@ -75,8 +75,8 @@ export type Node<S extends Message, R extends Message> = {
     handler: (event: U['data']) => U['response'],
   ) => () => void
   onStatus: (
-    handler: (status: Omit<Status, 'disconnected'>) => void,
-    filter?: Omit<Status, 'disconnected'>,
+    handler: (status: Exclude<Status, 'disconnected'>) => void,
+    filter?: Exclude<Status, 'disconnected'>,
   ) => () => void
   post: <T extends S['type'], U extends Extract<S, {type: T}>>(
     ...params: (U['data'] extends undefined ? [T] : never) | [T, U['data']]
@@ -134,7 +134,7 @@ export const createNodeMachine = <
         | HeartbeatEmitEvent
         | MessageEmitEvent<R>
         | ReceivedEmitEvent<R>
-        | (StatusEmitEvent & {status: Omit<Status, 'disconnected'>})
+        | (StatusEmitEvent & {status: Exclude<Status, 'disconnected'>})
       events:
         | {type: 'heartbeat.received'; message: MessageEvent<ProtocolMessage<HeartbeatMessage>>}
         | {type: 'message.received'; message: MessageEvent<ProtocolMessage<R>>}
@@ -249,7 +249,7 @@ export const createNodeMachine = <
         return {
           type: '_status',
           status: params.status,
-        } satisfies StatusEmitEvent & {status: Omit<Status, 'disconnected'>}
+        } satisfies StatusEmitEvent & {status: Exclude<Status, 'disconnected'>}
       }),
       'flush buffer': enqueueActions(({enqueue}) => {
         enqueue.raise(({context}) => ({
@@ -520,20 +520,26 @@ export const createNode = <S extends Message, R extends Message>(
     return unsubscribe
   }
 
+  let cachedStatus: Exclude<Status, 'disconnected'>
   const onStatus = (
-    handler: (status: Omit<Status, 'disconnected'>) => void,
-    filter?: Omit<Status, 'disconnected'>,
+    handler: (status: Exclude<Status, 'disconnected'>) => void,
+    filter?: Exclude<Status, 'disconnected'>,
   ) => {
     const {unsubscribe} = actor.on(
       // @ts-expect-error @todo ReceivedEmitEvent causes this
       '_status',
-      (event: StatusEmitEvent & {status: Omit<Status, 'disconnected'>}) => {
+      (event: StatusEmitEvent & {status: Exclude<Status, 'disconnected'>}) => {
+        cachedStatus = event.status
         if (filter && event.status !== filter) {
           return
         }
         handler(event.status)
       },
     )
+    // Call the handler immediately with the current status, if we have one
+    if (cachedStatus) {
+      handler(cachedStatus)
+    }
     return unsubscribe
   }
 

diff --git a/packages/core-loader/src/live-mode/enableLiveMode.ts b/packages/core-loader/src/live-mode/enableLiveMode.ts
@@ -59,13 +59,9 @@ export function enableLiveMode(options: LazyEnableLiveModeOptions): () => void {
     }),
   )
 
-  comlink.onStatus((status) => {
-    if (status === 'connected') {
-      $connected.set(true)
-    } else if (status === 'disconnected') {
-      $connected.set(false)
-    }
-  })
+  comlink.onStatus(() => {
+    $connected.set(true)
+  }, 'connected')
 
   comlink.on('loader/perspective', (data) => {
     if (data.projectId === projectId && data.dataset === dataset) {
@@ -239,5 +235,6 @@ export function enableLiveMode(options: LazyEnableLiveModeOptions): () => void {
     unsetFetcher?.()
     unlistenConnection()
     stop()
+    $connected.set(false)
   }
 }
diff --git a/packages/presentation/src/PresentationTool.tsx b/packages/presentation/src/PresentationTool.tsx
@@ -517,6 +517,8 @@ export default function PresentationTool(props: {
                     <Flex direction="column" flex={1} height="fill" ref={setBoundaryElement}>
                       <BoundaryElementProvider element={boundaryElement}>
                         <Preview
+                          // Make sure the iframe is unmounted if the targetOrigin has changed
+                          key={targetOrigin}
                           canSharePreviewAccess={canSharePreviewAccess}
                           canToggleSharePreviewAccess={canToggleSharePreviewAccess}
                           canUseSharedPreviewAccess={canUseSharedPreviewAccess}

diff --git a/packages/preview-kit-compat/src/useDocumentsInUse.ts b/packages/preview-kit-compat/src/useDocumentsInUse.ts
@@ -29,18 +29,15 @@ export function useDocumentsInUse(
       }),
     )
 
-    comlink.onStatus((status) => {
-      if (status === 'connected') {
-        setConnected(true)
-      } else if (status === 'disconnected') {
-        setConnected(false)
-      }
-    })
+    comlink.onStatus(() => {
+      setConnected(true)
+    }, 'connected')
 
     const timeout = setTimeout(() => setComlink(comlink), 0)
     const stop = comlink.start()
     return () => {
       stop()
+      setConnected(false)
       setComlink(null)
       clearTimeout(timeout)
     }

diff --git a/packages/visual-editing/src/ui/VisualEditing.tsx b/packages/visual-editing/src/ui/VisualEditing.tsx
@@ -1,4 +1,4 @@
-import {type FunctionComponent} from 'react'
+import {useEffect, useState, type FunctionComponent} from 'react'
 import type {VisualEditingOptions} from '../types'
 import {History} from './History'
 import {Meta} from './Meta'
@@ -12,19 +12,22 @@ import {useDatasetMutator} from './useDatasetMutator'
  */
 export const VisualEditing: FunctionComponent<VisualEditingOptions> = (props) => {
   const {components, history, refresh, zIndex} = props
-  const inFrame = window.self !== window.top || window.opener
+  const [inFrame, setInFrame] = useState<boolean | null>(null)
+  useEffect(() => setInFrame(window.self !== window.top || Boolean(window.opener)), [])
 
-  const comlink = useComlink(inFrame)
+  const comlink = useComlink(inFrame === true)
   useDatasetMutator(comlink)
 
   return (
     <>
-      <Overlays
-        comlink={comlink}
-        componentResolver={components}
-        inFrame={inFrame}
-        zIndex={zIndex}
-      />
+      {inFrame !== null && (
+        <Overlays
+          comlink={comlink}
+          componentResolver={components}
+          inFrame={inFrame}
+          zIndex={zIndex}
+        />
+      )}
       {comlink && (
         <>
           <History comlink={comlink} history={history} />

diff --git a/packages/visual-editing/src/ui/useComlink.tsx b/packages/visual-editing/src/ui/useComlink.tsx
@@ -26,10 +26,19 @@ export function useComlink(active: boolean = true): VisualEditingNode | undefine
       }),
     )
 
-    setNode(instance)
+    let timeout = 0
     const stop = instance.start()
+    // Wait with forwarding the comlink until the connection is established
+    const unsubscribe = instance.onStatus(() => {
+      // Due to race conditions in when Presentation Tool loads up components with handlers for comlink, we need to wait a bit before forwarding the comlink instance
+      timeout = window.setTimeout(() => {
+        setNode(instance)
+      }, 3_000)
+    }, 'connected')
 
     return () => {
+      clearTimeout(timeout)
+      unsubscribe()
       stop()
       setNode(undefined)
     }

diff --git a/packages/visual-editing/src/ui/useDatasetMutator.ts b/packages/visual-editing/src/ui/useDatasetMutator.ts
@@ -1,6 +1,6 @@
-import {useEffect, useState} from 'react'
+import {useEffect} from 'react'
 import {createActor} from 'xstate'
-import {setActor, type MutatorActor} from '../optimistic/context'
+import {setActor} from '../optimistic/context'
 import {createSharedListener} from '../optimistic/state/createSharedListener'
 import {createDatasetMutator} from '../optimistic/state/datasetMutator'
 import type {VisualEditingNode} from '../types'
@@ -9,11 +9,7 @@ import type {VisualEditingNode} from '../types'
  * Hook for maintaining a channel between overlays and the presentation tool
  * @internal
  */
-export function useDatasetMutator(
-  comlink: VisualEditingNode | undefined,
-): MutatorActor | undefined {
-  const [mutator, setMutator] = useState<MutatorActor>()
-
+export function useDatasetMutator(comlink: VisualEditingNode | undefined): void {
   useEffect(() => {
     if (!comlink) return
     const listener = createSharedListener(comlink)
@@ -23,7 +19,6 @@ export function useDatasetMutator(
       input: {client: {withConfig: () => {}}, sharedListener: listener},
     })
 
-    setMutator(mutator)
     mutator.start()
 
     // Fetch features to determine if optimistic updates are supported
@@ -51,9 +46,6 @@ export function useDatasetMutator(
       mutator.stop()
       featuresFetch.abort()
       unsub()
-      setMutator(undefined)
     }
   }, [comlink])
-
-  return mutator
 }