Merge pull request #509 from NERSC/fix-intermittent-draining-issue

fix bad nats connection on page refresh and add nats status indicator

Author: swelborn

frontend/interactEM/src/components/natsstatus.tsx

@@ -0,0 +1,67 @@
+import { Chip, Tooltip } from "@mui/material"
+import type { NatsConnection } from "@nats-io/nats-core"
+import { useMemo } from "react"
+import { useNats } from "../contexts/nats"
+
+type NatsStatus = {
+  color: "success" | "warning" | "error"
+  label: string
+  tooltip: string
+}
+
+const getNatsStatus = (
+  connection: NatsConnection | null,
+  isConnected: boolean,
+): NatsStatus => {
+  if (!connection) {
+    return {
+      color: "error",
+      label: "Disconnected",
+      tooltip: "NATS is not connected",
+    }
+  }
+
+  if (connection.isDraining()) {
+    return {
+      color: "warning",
+      label: "Draining",
+      tooltip: "NATS connection is draining",
+    }
+  }
+
+  if (connection.isClosed()) {
+    return {
+      color: "error",
+      label: "Closed",
+      tooltip: "NATS connection is closed",
+    }
+  }
+
+  if (isConnected) {
+    return {
+      color: "success",
+      label: "Connected",
+      tooltip: "NATS connection is active",
+    }
+  }
+
+  return {
+    color: "warning",
+    label: "Connecting",
+    tooltip: "NATS connection is establishing",
+  }
+}
+
+export const NatsStatusIndicator = () => {
+  const { natsConnection, isConnected } = useNats()
+  const status = useMemo(
+    () => getNatsStatus(natsConnection, isConnected),
+    [natsConnection, isConnected],
+  )
+
+  return (
+    <Tooltip title={status.tooltip} arrow placement="top">
+      <Chip label="NATS" color={status.color} size="small" variant="outlined" />
+    </Tooltip>
+  )
+}

frontend/interactEM/src/components/pipelines/hudcomposer.tsx

@@ -9,6 +9,7 @@ import {
 import { useCallback, useRef, useState } from "react"
 import { useActivePipeline } from "../../hooks/api/useActivePipeline"
 import { usePipelineStore } from "../../stores"
+import { NatsStatusIndicator } from "../natsstatus"
 import { DeletePipelineButton } from "./deletebutton"
 import { HudListButton } from "./hudlistbutton"
 import { LaunchPipelineButton } from "./launchbutton"
@@ -126,8 +127,12 @@ export const HudComposer: React.FC = () => {
           onClick={handleTogglePipelineDrawer}
         />
 
-        {/* Pipeline content */}
-        {pipelineDisplayContent()}
+        <Box sx={{ display: "flex", alignItems: "center", flex: 1 }}>
+          {/* Pipeline content */}
+          {pipelineDisplayContent()}
+        </Box>
+
+        <NatsStatusIndicator />
       </Box>
 
       {/* Pipeline List Drawer */}

frontend/interactEM/src/components/pipelines/hudrunning.tsx

@@ -5,6 +5,7 @@ import type { PipelineDeploymentPublic } from "../../client"
 import { useActivePipeline } from "../../hooks/api/useActivePipeline"
 import { useDeployment } from "../../hooks/api/useDeploymentsQuery"
 import { usePipelineStore } from "../../stores"
+import { NatsStatusIndicator } from "../natsstatus"
 import { DeploymentManagementPanel } from "./deploymentmanagementpanel"
 import { HudListButton } from "./hudlistbutton"
 import { HudRuntimeInfo } from "./hudruntimeinfo"
@@ -63,6 +64,8 @@ export const HudRunning: React.FC = () => {
           isLoading={loading}
           error={error}
         />
+
+        <NatsStatusIndicator />
       </Box>
 
       {/* Deployment Management Panel */}

frontend/interactEM/src/contexts/nats.tsx

@@ -65,22 +65,26 @@ export const NatsProvider: React.FC<NatsProviderProps> = ({
     isConnected: false,
   })
 
-  const [natsConnection, setNatsConnection] = useState<NatsConnection | null>(
-    null,
-  )
   const { token, natsJwt, isAuthenticated } = useAuth()
   const tokenRef = useRef(token)
   const natsJwtRef = useRef(natsJwt)
+  const hasConnectedRef = useRef(false)
+  const connectionRef = useRef<NatsConnection | null>(null)
 
   useEffect(() => {
     tokenRef.current = token
   }, [token])
 
+  useEffect(() => {
+    natsJwtRef.current = natsJwt
+  }, [natsJwt])
+
   useEffect(() => {
     if (!isAuthenticated) {
       return
     }
-    async function setupNatsServices(nc: NatsConnection) {
+
+    async function setupNatsServices(nc: NatsConnection): Promise<boolean> {
       try {
         const js = jetstream(nc)
         const jsm = await jetstreamManager(nc)
@@ -93,11 +97,14 @@ export const NatsProvider: React.FC<NatsProviderProps> = ({
           keyValueManager: kvm,
           isConnected: true,
         })
+        return true
       } catch (error) {
         console.error("Failed to setup NATS services:", error)
         setState((prev) => ({ ...prev, isConnected: false }))
+        return false
       }
     }
+
     async function connect() {
       try {
         const servers = Array.isArray(natsServers) ? natsServers : [natsServers]
@@ -125,11 +132,20 @@ export const NatsProvider: React.FC<NatsProviderProps> = ({
           maxReconnectAttempts: 30,
         })
 
+        connectionRef.current = nc
         console.log("NATS connection successful")
 
-        setNatsConnection(nc)
-        await setupNatsServices(nc)
-
+        const setupOk = await setupNatsServices(nc)
+        if (!setupOk) {
+          try {
+            await nc.drain()
+          } catch (err) {
+            console.error("Error draining NATS connection:", err)
+          }
+          connectionRef.current = null
+          hasConnectedRef.current = false
+          return
+        }
         // natsConnection will cycle through the following status sequence when
         // it is disconnected:
         // 1. Error
@@ -145,30 +161,44 @@ export const NatsProvider: React.FC<NatsProviderProps> = ({
                 setState((prev) => ({ ...prev, isConnected: true }))
                 break
               case "error":
-                // TODO: handle error better, maybe with UI update
-                setState((prev) => ({
-                  ...prev,
-                  isConnected: false,
-                }))
+              case "disconnect":
+              case "staleConnection":
+              case "close":
+                setState((prev) => ({ ...prev, isConnected: false }))
                 break
             }
           }
         })().catch(console.error)
       } catch (error) {
         console.error("Failed to connect to NATS:", error)
         setState((prev) => ({ ...prev, isConnected: false }))
+        hasConnectedRef.current = false
+        const nc = connectionRef.current
+        connectionRef.current = null
+        if (nc) {
+          try {
+            await nc.drain()
+          } catch (err) {
+            console.error("Error draining NATS connection:", err)
+          }
+        }
       }
     }
 
-    if (!natsConnection) {
+    if (!hasConnectedRef.current) {
+      hasConnectedRef.current = true
       connect()
     }
+
     return () => {
-      if (natsConnection) {
+      hasConnectedRef.current = false
+      const nc = connectionRef.current
+      connectionRef.current = null
+      if (nc) {
         console.log("Draining NATS connection")
         ;(async () => {
           try {
-            await natsConnection.drain()
+            await nc.drain()
             console.log("NATS connection drained and closed")
           } catch (err) {
             console.error("Error draining NATS connection:", err)
@@ -183,7 +213,7 @@ export const NatsProvider: React.FC<NatsProviderProps> = ({
         isConnected: false,
       })
     }
-  }, [isAuthenticated, natsConnection, natsServers])
+  }, [isAuthenticated, natsServers])
 
   if (!isAuthenticated) {
     return null

frontend/interactEM/src/hooks/nats/natsErrors.ts

@@ -0,0 +1,11 @@
+import {
+  ClosedConnectionError,
+  DrainingConnectionError,
+} from "@nats-io/nats-core"
+
+export function isConnectionError(err: unknown): boolean {
+  return (
+    err instanceof DrainingConnectionError ||
+    err instanceof ClosedConnectionError
+  )
+}

frontend/interactEM/src/hooks/nats/useBucket.ts

@@ -1,39 +1,57 @@
 import type { KV } from "@nats-io/kv"
-import { DrainingConnectionError } from "@nats-io/nats-core"
 import { useEffect, useRef, useState } from "react"
 import { useNats } from "../../contexts/nats"
+import { isConnectionError } from "./natsErrors"
 
 export const useBucket = (bucketName: string): KV | null => {
-  const { keyValueManager } = useNats()
+  const { keyValueManager, isConnected } = useNats()
   const [bucket, setBucket] = useState<KV | null>(null)
-  const isMounted = useRef(true)
+  const isMountedRef = useRef(true)
 
   useEffect(() => {
-    isMounted.current = true
+    isMountedRef.current = true
+    // Create abort signal for this effect instance to handle race conditions
+    // when connection drops mid-open and reconnects before promise settles
+    const abortController = new AbortController()
 
     const openBucket = async () => {
-      if (keyValueManager && !bucket) {
-        try {
-          const openedBucket = await keyValueManager.open(bucketName)
-          if (isMounted.current) {
-            setBucket(openedBucket)
-          }
-        } catch (error) {
-          if (error instanceof DrainingConnectionError) {
-            // quietly ignore if connection is draining
-            return
-          }
-          console.error(`Failed to open bucket "${bucketName}":`, error)
+      if (!keyValueManager || !isConnected) {
+        return
+      }
+
+      try {
+        const openedBucket = await keyValueManager.open(bucketName)
+
+        // Check if this effect instance was aborted before setting state
+        if (abortController.signal.aborted) {
+          return
+        }
+
+        if (isMountedRef.current) {
+          setBucket(openedBucket)
+        }
+      } catch (error) {
+        // Check if this effect instance was aborted before handling error
+        if (abortController.signal.aborted) {
+          return
+        }
+
+        if (isConnectionError(error)) {
+          // Silently ignore connection lifecycle errors
+          return
         }
+        console.error(`Failed to open bucket "${bucketName}":`, error)
       }
     }
 
     openBucket()
 
     return () => {
-      isMounted.current = false
+      isMountedRef.current = false
+      abortController.abort()
+      setBucket(null)
     }
-  }, [keyValueManager, bucket, bucketName])
+  }, [keyValueManager, bucketName, isConnected])
 
   return bucket
 }