Recovery for locations stuck after failed refresh

python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_watch_grpc_server.py

@@ -2,10 +2,15 @@
 from typing import Any
 
 from dagster._core.remote_representation.grpc_server_state_subscriber import (
-    LocationStateChangeEvent,
     LocationStateChangeEventType,
     LocationStateSubscriber,
 )
+from dagster._core.workspace.workspace import (
+    CodeLocationEntry,
+    CodeLocationLoadStatus,
+    DefinitionsSource,
+)
+from dagster._utils.error import SerializableErrorInfo
 
 from dagster_graphql_tests.graphql.graphql_context_test_suite import (
     GraphQLContextVariant,
@@ -27,15 +32,103 @@ def test_grpc_server_handle_message_subscription(self, graphql_context):
         graphql_context.process_context.add_state_subscriber(test_subscriber)
         location.client.shutdown_server()
 
-        # Wait for event
+        # Wait for LOCATION_ERROR event. LOCATION_DISCONNECTED may arrive first since the
+        # watch thread detects the disconnect before exhausting reconnect attempts.
         start_time = time.time()
         timeout = 60
-        while not len(events) > 0:
+        while not any(e.event_type == LocationStateChangeEventType.LOCATION_ERROR for e in events):
             if time.time() - start_time > timeout:
-                raise Exception("Timed out waiting for LocationStateChangeEvent")
+                raise Exception("Timed out waiting for LOCATION_ERROR event")
             time.sleep(1)
 
-        assert len(events) == 1
-        assert isinstance(events[0], LocationStateChangeEvent)
-        assert events[0].event_type == LocationStateChangeEventType.LOCATION_ERROR
-        assert events[0].location_name == location.name
+        error_events = [
+            e for e in events if e.event_type == LocationStateChangeEventType.LOCATION_ERROR
+        ]
+        assert len(error_events) == 1
+        assert error_events[0].location_name == location.name
+
+        # LOCATION_DISCONNECTED should have arrived before LOCATION_ERROR
+        disconnect_events = [
+            e for e in events if e.event_type == LocationStateChangeEventType.LOCATION_DISCONNECTED
+        ]
+        assert len(disconnect_events) == 1
+        assert disconnect_events[0].location_name == location.name
+
+
+RecoveryTestSuite: Any = make_graphql_context_test_suite(
+    context_variants=[GraphQLContextVariant.non_launchable_sqlite_instance_deployed_grpc_env()]
+)
+
+
+class TestGrpcServerRecovery(RecoveryTestSuite):
+    def test_grpc_server_recovery_after_failed_refresh(self, graphql_context):
+        """Integration test: verify that the watch thread triggers recovery when a code location
+        is stuck in an error state after a failed refresh.
+
+        This simulates the K8s rolling deployment race condition end-to-end:
+        1. Directly set the workspace entry to an errored state (simulating a failed refresh)
+        2. Verify the watch thread detects the error via _should_recover_location
+        3. Verify it fires on_disconnect + on_reconnected events
+        4. Verify the event handler calls refresh_code_location (which now succeeds)
+        5. Verify the location recovers
+        """
+        ctx = graphql_context.process_context
+        location = next(iter(ctx.create_request_context().code_locations))
+        location_name = location.name
+
+        events = []
+        test_subscriber = LocationStateSubscriber(events.append)
+        ctx.add_state_subscriber(test_subscriber)
+
+        # Verify location starts healthy
+        assert not ctx.has_code_location_error(location_name)
+
+        # Simulate a failed refresh by directly setting the workspace entry to an errored state.
+        # This is what happens when _load_location fails during the race condition.
+        current_entry = ctx.get_current_workspace().code_location_entries[location_name]
+        errored_entry = CodeLocationEntry(
+            origin=current_entry.origin,
+            code_location=None,
+            load_error=SerializableErrorInfo("Simulated failure", [], "RuntimeError"),
+            load_status=CodeLocationLoadStatus.LOADED,
+            display_metadata=current_entry.origin.get_display_metadata(),
+            update_timestamp=0.0,
+            version_key="stale-version-key",
+            definitions_source=DefinitionsSource.CODE_SERVER,
+        )
+        with ctx._lock:  # noqa: SLF001
+            ctx._current_workspace = (  # noqa: SLF001
+                ctx._current_workspace.with_code_location(  # noqa: SLF001
+                    location_name, errored_entry
+                )
+            )
+
+        # Verify location is now errored
+        assert ctx.has_code_location_error(location_name)
+
+        # The watch thread should detect the error via _should_recover_location and fire
+        # on_disconnect + on_reconnected, which flows through _location_state_events_handler
+        # and triggers refresh_code_location, recovering the location.
+        start_time = time.time()
+        timeout = 30
+        while ctx.has_code_location_error(location_name):
+            if time.time() - start_time > timeout:
+                raise Exception(
+                    f"Timed out waiting for location {location_name} to recover. "
+                    f"Events received: {[(e.event_type, e.location_name) for e in events]}"
+                )
+            time.sleep(0.5)
+
+        # Location should be recovered — no longer errored, code_location is set
+        assert not ctx.has_code_location_error(location_name)
+        assert ctx.has_code_location(location_name)
+
+        # Verify the recovery events were fired
+        disconnect_events = [
+            e for e in events if e.event_type == LocationStateChangeEventType.LOCATION_DISCONNECTED
+        ]
+        reconnected_events = [
+            e for e in events if e.event_type == LocationStateChangeEventType.LOCATION_RECONNECTED
+        ]
+        assert len(disconnect_events) >= 1
+        assert len(reconnected_events) >= 1

python_modules/dagster/dagster/_core/workspace/context.py

@@ -1075,6 +1075,21 @@ def _start_watch_thread(self, origin: GrpcServerCodeLocationOrigin) -> None:
                     ),
                 )
             ),
+            on_disconnect=lambda location_name: self._send_state_event_to_subscribers(
+                LocationStateChangeEvent(
+                    LocationStateChangeEventType.LOCATION_DISCONNECTED,
+                    location_name=location_name,
+                    message="Disconnected from the server.",
+                )
+            ),
+            on_reconnected=lambda location_name: self._send_state_event_to_subscribers(
+                LocationStateChangeEvent(
+                    LocationStateChangeEventType.LOCATION_RECONNECTED,
+                    location_name=location_name,
+                    message="Reconnected to the server.",
+                )
+            ),
+            needs_location_refresh=self._should_recover_location,
         )
         self._watch_thread_shutdown_events[location_name] = shutdown_event
         self._watch_threads[location_name] = watch_thread
@@ -1171,6 +1186,22 @@ def has_code_location_error(self, location_name: str) -> bool:
                 is not None
             )
 
+    def _should_recover_location(self, location_name: str, version_key: str) -> bool:
+        """Check without locking whether a code location needs a recovery refresh.
+
+        Returns True if the location is in an error state or has a stale version key.
+        Called from the watch thread without holding self._lock. This is safe because
+        _current_workspace is replaced atomically (single reference assignment) and we only
+        need a consistent-enough snapshot — correctness doesn't depend on reading the latest
+        value.
+        """
+        check.str_param(location_name, "location_name")
+        check.str_param(version_key, "version_key")
+        entry = self._current_workspace.code_location_entries.get(location_name, None)
+        if entry is None:
+            return False
+        return entry.load_error is not None or entry.version_key != version_key
+
     def reload_code_location(self, name: str) -> None:
         new_entry = self._load_location(
             self._current_workspace.code_location_entries[name].origin, reload=True
@@ -1238,19 +1269,31 @@ def create_request_context(self, source: object | None = None) -> WorkspaceReque
 
     def _location_state_events_handler(self, event: LocationStateChangeEvent) -> None:
         # If the server was updated or we were not able to reconnect, we immediately reload the
-        # location handle
-        if event.event_type in (
+        # location handle. Reconnect/disconnect events are logged but don't trigger a refresh
+        # unless the location is in an error state (in which case reconnect triggers a recovery
+        # refresh).
+        logger = logging.getLogger("dagster-webserver")
+        refresh_event_types = (
             LocationStateChangeEventType.LOCATION_UPDATED,
             LocationStateChangeEventType.LOCATION_ERROR,
-        ):
-            # In case of an updated location, reload the handle to get updated repository data and
-            # re-attach a subscriber
-            # In case of a location error, just reload the handle in order to update the workspace
-            # with the correct error messages
-            logging.getLogger("dagster-webserver").info(
+        )
+
+        location_is_errored = self.has_code_location_error(event.location_name)
+        # Refresh on update/error events, or on any non-disconnect event when the location is
+        # errored (to attempt recovery). Skip disconnect events for errored locations since the
+        # server is unreachable and retries would fail.
+        should_refresh = event.event_type in refresh_event_types or (
+            location_is_errored
+            and event.event_type != LocationStateChangeEventType.LOCATION_DISCONNECTED
+        )
+
+        if should_refresh:
+            logger.info(
                 f"Received {event.event_type} event for location {event.location_name}, refreshing"
             )
             self.refresh_code_location(event.location_name)
+        else:
+            logger.debug(f"Received {event.event_type} event for location {event.location_name}")
 
     def refresh_code_location(self, name: str) -> None:
         # This method reloads the webserver's copy of the code from the remote gRPC server without

python_modules/dagster/dagster/_grpc/server_watcher.py

@@ -17,6 +17,7 @@ def watch_grpc_server_thread(
     on_reconnected: Callable[[str], None],
     on_updated: Callable[[str, str], None],
     on_error: Callable[[str], None],
+    needs_location_refresh: Callable[[str, str], bool],
     shutdown_event: threading.Event,
     watch_interval: float | None = None,
     max_reconnect_attempts: int | None = None,
@@ -28,52 +29,83 @@ def watch_grpc_server_thread(
     1. The server_id has changed
     2. The server is unreachable
 
-    In the case of (1) The server ID has changed, we call `on_updated` and end the thread.
+    In the case of (1) The server ID has changed, we call `on_updated` and continue polling.
+    The thread does not exit — it keeps monitoring for further changes.
 
     In the case of (2) The server is unreachable, we attempt to automatically reconnect. If we
     are able to reconnect, there are two possibilities:
 
     a. The server ID has changed
-        -> In this case, we we call `on_updated` and end the thread.
+        -> In this case, we call `on_updated` and resume polling.
     b. The server ID is the same
-        -> In this case, we we call `on_reconnected`, and we go back to polling the server for
+        -> In this case, we call `on_reconnected`, and we go back to polling the server for
         changes.
 
     If we are unable to reconnect to the server within the specified max_reconnect_attempts, we
-    call on_error.
-
-    Once the on_updated or on_error events are called, this thread shuts down completely. These two
-    events are called at most once, while `on_disconnected` and `on_reconnected` may be called
-    multiple times in order to be properly handle intermittent network failures.
+    call on_error. After on_error, the reconnect loop continues indefinitely — the thread does
+    not shut down, so that if the server eventually comes back, it will be detected via
+    on_updated (not on_reconnected). This is intentional: on_error already notified subscribers
+    of the failure, so recovery must go through on_updated to trigger a refresh that clears the
+    error state. The stored server ID is cleared on error to ensure the on_updated path is
+    taken regardless of whether the actual server ID changed.
+
+    Additionally, if `needs_location_refresh` returns True during normal polling (indicating
+    that the workspace entry is in an error state or has a stale version key), `on_disconnect`
+    and `on_reconnected` are fired even if the server ID hasn't changed. This enables recovery
+    when a prior refresh failed (e.g. during a Kubernetes rolling deployment where the gRPC
+    call routed to a dying pod).
+
+    `on_updated` is called each time a server ID change is detected and does not cause the
+    thread to exit. `on_error` is called at most once per disconnect cycle but does not cause
+    the thread to exit. `on_disconnect` and `on_reconnected` may be called multiple times to
+    properly handle intermittent network failures or workspace error recovery. The thread only
+    exits when shutdown_event is set.
     """
     check.str_param(location_name, "location_name")
     check.inst_param(client, "client", DagsterGrpcClient)
     check.callable_param(on_disconnect, "on_disconnect")
     check.callable_param(on_reconnected, "on_reconnected")
     check.callable_param(on_updated, "on_updated")
     check.callable_param(on_error, "on_error")
+    check.callable_param(needs_location_refresh, "needs_location_refresh")
     watch_interval = check.opt_numeric_param(watch_interval, "watch_interval", WATCH_INTERVAL)
     max_reconnect_attempts = check.opt_int_param(
         max_reconnect_attempts, "max_reconnect_attempts", MAX_RECONNECT_ATTEMPTS
     )
 
+    needs_location_refresh_count = 0
     server_id = {"current": None, "error": False}
 
-    def current_server_id():
+    def current_server_id() -> str | None:
         return server_id["current"]
 
-    def has_error():
+    def has_error() -> bool:
         return server_id["error"]
 
-    def set_server_id(new_id):
+    def _needs_location_refresh() -> bool:
+        current_id = current_server_id()
+        result = current_id is not None and needs_location_refresh(location_name, current_id)
+        if result:
+            nonlocal needs_location_refresh_count
+            needs_location_refresh_count += 1
+        return result
+
+    def set_server_id(new_id: str) -> None:
         server_id["current"] = new_id
         server_id["error"] = False
+        nonlocal needs_location_refresh_count
+        needs_location_refresh_count = 0
 
-    def set_error():
+    def set_error() -> None:
+        # Clearing current server ID ensures that post-error recovery always takes
+        # the on_updated path in reconnect_loop, which is needed to trigger a
+        # refresh that clears the error state in subscribers.
         server_id["current"] = None
         server_id["error"] = True
 
     def watch_for_changes():
+        nonlocal needs_location_refresh_count
+        needs_location_refresh_count = 0
         while True:
             if shutdown_event.is_set():
                 break
@@ -86,6 +118,20 @@ def watch_for_changes():
             elif curr != new_server_id:
                 set_server_id(new_server_id)
                 on_updated(location_name, new_server_id)
+            elif _needs_location_refresh():
+                # Wait 2 cycles to confirm the error persists (giving get_server_id() a
+                # chance to throw DagsterUserCodeUnreachableError and trigger reconnect_loop
+                # instead). After that, retry every 10 cycles to avoid log spam and
+                # unnecessary gRPC calls when the location is permanently stuck.
+                if (
+                    needs_location_refresh_count >= 2
+                    and (needs_location_refresh_count - 2) % 10 == 0
+                ):
+                    # get_server_id() keeps succeeding but _needs_location_refresh() still
+                    # returns True. The workspace entry is errored/stale. Fire disconnect +
+                    # reconnect callbacks to trigger a recovery refresh.
+                    on_disconnect(location_name)
+                    on_reconnected(location_name)
 
             shutdown_event.wait(watch_interval)
 
@@ -100,9 +146,12 @@ def reconnect_loop():
                 new_server_id = client.get_server_id(timeout=REQUEST_TIMEOUT)
                 if current_server_id() == new_server_id and not has_error():
                     # Intermittent failure, was able to reconnect to the same server
+                    # before max_reconnect_attempts was exhausted.
                     on_reconnected(location_name)
                     return
                 else:
+                    # Either the server ID changed, or we're recovering after on_error
+                    # was already called. Either way, on_updated triggers a refresh.
                     on_updated(location_name, new_server_id)
                     set_server_id(new_server_id)
                     return
@@ -112,6 +161,8 @@ def reconnect_loop():
             if attempts >= max_reconnect_attempts and not has_error():
                 on_error(location_name)
                 set_error()
+                # Intentionally does not return — the loop continues so that if the
+                # server eventually comes back, it will be detected via on_updated.
 
     while True:
         if shutdown_event.is_set():
@@ -126,21 +177,22 @@ def reconnect_loop():
 def create_grpc_watch_thread(
     location_name: str,
     client: DagsterGrpcClient,
-    on_disconnect: Callable[[str], None] | None = None,
-    on_reconnected: Callable[[str], None] | None = None,
-    on_updated: Callable[[str, str], None] | None = None,
-    on_error: Callable[[str], None] | None = None,
+    on_disconnect: Callable[[str], None],
+    on_reconnected: Callable[[str], None],
+    on_updated: Callable[[str, str], None],
+    on_error: Callable[[str], None],
+    needs_location_refresh: Callable[[str, str], bool],
     watch_interval: float | None = None,
     max_reconnect_attempts: int | None = None,
 ) -> tuple[threading.Event, threading.Thread]:
     check.str_param(location_name, "location_name")
     check.inst_param(client, "client", DagsterGrpcClient)
 
-    noop = lambda *a: None
-    on_disconnect = check.opt_callable_param(on_disconnect, "on_disconnect", noop)
-    on_reconnected = check.opt_callable_param(on_reconnected, "on_reconnected", noop)
-    on_updated = check.opt_callable_param(on_updated, "on_updated", noop)
-    on_error = check.opt_callable_param(on_error, "on_error", noop)
+    check.callable_param(on_disconnect, "on_disconnect")
+    check.callable_param(on_reconnected, "on_reconnected")
+    check.callable_param(on_updated, "on_updated")
+    check.callable_param(on_error, "on_error")
+    check.callable_param(needs_location_refresh, "needs_location_refresh")
 
     shutdown_event = threading.Event()
     thread = threading.Thread(
@@ -152,6 +204,7 @@ def create_grpc_watch_thread(
             on_reconnected,
             on_updated,
             on_error,
+            needs_location_refresh,
             shutdown_event,
             watch_interval,
             max_reconnect_attempts,