[dagit] Feature flag for disabling WebSocket usage (#7557)

dagster-io · Apr 25, 2022 · 5480b04 · 5480b04
1 parent a141905
commit 5480b04
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 5 deletions.
diff --git a/js_modules/dagit/packages/core/src/app/Flags.tsx b/js_modules/dagit/packages/core/src/app/Flags.tsx
@@ -8,6 +8,7 @@ const DAGIT_FLAGS_KEY = 'DAGIT_FLAGS';
 export enum FeatureFlag {
   flagDebugConsoleLogging = 'flagDebugConsoleLogging',
   flagAlwaysCollapseNavigation = 'flagAlwaysCollapseNavigation',
+  flagDisableWebsockets = 'flagDisableWebsockets',
 }
 
 export const getFeatureFlags: () => FeatureFlag[] = memoize(

diff --git a/js_modules/dagit/packages/core/src/app/SettingsRoot.tsx b/js_modules/dagit/packages/core/src/app/SettingsRoot.tsx
@@ -111,6 +111,16 @@ const SettingsRoot = () => {
                 />
               ),
             },
+            {
+              key: 'Disable WebSockets',
+              value: (
+                <Checkbox
+                  format="switch"
+                  checked={flags.includes(FeatureFlag.flagDisableWebsockets)}
+                  onChange={() => toggleFlag(FeatureFlag.flagDisableWebsockets)}
+                />
+              ),
+            },
           ]}
         />
       </Box>

diff --git a/js_modules/dagit/packages/core/src/app/WebSocketProvider.tsx b/js_modules/dagit/packages/core/src/app/WebSocketProvider.tsx
@@ -4,17 +4,21 @@ import * as React from 'react';
 import styled from 'styled-components/macro';
 import {SubscriptionClient} from 'subscriptions-transport-ws';
 
+import {useFeatureFlags} from './Flags';
+
 type Availability = 'attempting-to-connect' | 'unavailable' | 'available';
 
 export type WebSocketContextType = {
   availability: Availability;
   status: number;
+  disabled?: boolean;
   websocketClient?: SubscriptionClient;
 };
 
 export const WebSocketContext = React.createContext<WebSocketContextType>({
   availability: 'attempting-to-connect',
   status: WebSocket.CONNECTING,
+  disabled: false,
 });
 
 const WS_EVENTS = [
@@ -36,6 +40,8 @@ interface Props {
 export const WebSocketProvider: React.FC<Props> = (props) => {
   const {children, websocketClient} = props;
   const [status, setStatus] = React.useState(websocketClient.status);
+  const {flagDisableWebsockets: disabled} = useFeatureFlags();
+
   const [availability, setAvailability] = React.useState<Availability>(
     websocketClient.status === WebSocket.OPEN
       ? 'available'
@@ -49,8 +55,9 @@ export const WebSocketProvider: React.FC<Props> = (props) => {
       availability,
       status,
       websocketClient,
+      disabled,
     }),
-    [availability, status, websocketClient],
+    [availability, disabled, status, websocketClient],
   );
 
   const debouncedSetter = React.useMemo(() => debounce(setStatus, DEBOUNCE_TIME), []);

diff --git a/js_modules/dagit/packages/core/src/gantt/GanttChart.tsx b/js_modules/dagit/packages/core/src/gantt/GanttChart.tsx
@@ -230,8 +230,8 @@ const GanttChartInner = (props: GanttChartInnerProps) => {
 
   const {rootServerURI} = React.useContext(AppContext);
 
-  const {availability, status} = React.useContext(WebSocketContext);
-  const lostWebsocket = availability === 'available' && status === WebSocket.CLOSED;
+  const {availability, disabled, status} = React.useContext(WebSocketContext);
+  const lostWebsocket = !disabled && availability === 'available' && status === WebSocket.CLOSED;
 
   // The slider in the UI updates `options.zoom` from 1-100. We convert that value
   // into a px-per-ms "scale", where the minimum is the value required to zoom-to-fit.

diff --git a/js_modules/dagit/packages/core/src/runs/LogsProvider.tsx b/js_modules/dagit/packages/core/src/runs/LogsProvider.tsx
@@ -264,13 +264,13 @@ const LogsProviderWithQuery = (props: LogsProviderWithQueryProps) => {
 
 export const LogsProvider: React.FC<LogsProviderProps> = (props) => {
   const {children, runId} = props;
-  const {availability} = React.useContext(WebSocketContext);
+  const {availability, disabled} = React.useContext(WebSocketContext);
 
   if (availability === 'attempting-to-connect') {
     return <>{children({allNodes: [], loading: true})}</>;
   }
 
-  if (availability === 'unavailable') {
+  if (availability === 'unavailable' || disabled) {
     return <LogsProviderWithQuery runId={runId}>{children}</LogsProviderWithQuery>;
   }
 

diff --git a/js_modules/dagit/packages/core/src/testing/TestProvider.tsx b/js_modules/dagit/packages/core/src/testing/TestProvider.tsx
@@ -37,6 +37,7 @@ const testValue: AppContextValue = {
 const websocketValue: WebSocketContextType = {
   availability: 'available',
   status: WebSocket.OPEN,
+  disabled: false,
 };
 
 interface Props {