[AssetLiveDataProvider] Retry asset keys that fail to fetch after the…

… poll interval. (#17393)

## Summary & Motivation

Currently if an asset key fails to fetch we never retry fetching it
again because it gets stuck in the requested state. This PR fixes that
and marks them as "fetched" when this happens. This has the consequence
of (1) letting us move on to the next set of asset keys which could
succeed and (2) it means we'll retry the failed asset keys again after
the polling interval threshold is crossed.

## How I Tested These Changes
jest. Locally threw an error and saw that we did end up retrying the
failed asset keys after the polling interval was met.

js_modules/dagster-ui/packages/ui-core/src/asset-data/AssetLiveDataProvider.tsx

@@ -347,8 +347,25 @@ async function _batchedQueryAssets(
     doNextFetch(pollRate);
   } catch (e) {
     console.error(e);
-    // Retry fetching in 5 seconds if theres a network error
-    setTimeout(doNextFetch, 5000);
+
+    if ((e as any)?.message?.includes('500')) {
+      // Mark these assets as fetched so that we don't retry them until after the poll interval rather than retrying them immediately.
+      // This is preferable because if the assets failed to fetch it's likely due to a timeout due to the query being too expensive and retrying it
+      // will not make it more likely to succeed and it would add more load to the database.
+      const fetchedTime = Date.now();
+      assetKeys.forEach((key) => {
+        lastFetchedOrRequested[tokenForAssetKey(key)] = {
+          fetched: fetchedTime,
+        };
+      });
+    } else {
+      // If it's not a timeout from the backend then lets keep retrying instead of moving on.
+      assetKeys.forEach((key) => {
+        delete lastFetchedOrRequested[tokenForAssetKey(key)];
+      });
+    }
+
+    setTimeout(doNextFetch, Math.min(pollRate, 5000));
   }
 }
 

js_modules/dagster-ui/packages/ui-core/src/asset-data/__tests__/AssetLiveDataProvider.test.tsx

@@ -2,6 +2,7 @@ jest.useFakeTimers();
 
 import {MockedProvider, MockedResponse} from '@apollo/client/testing';
 import {render, act, waitFor} from '@testing-library/react';
+import {GraphQLError} from 'graphql/error';
 import React from 'react';
 
 import {AssetKey, AssetKeyInput, buildAssetKey} from '../../graphql/types';
@@ -350,4 +351,69 @@ describe('AssetLiveDataProvider', () => {
       expect(resultFn3).toHaveBeenCalled();
     });
   });
+
+  it('Skips over asset keys that fail to fetch', async () => {
+    const assetKeys = [buildAssetKey({path: ['key1']})];
+    const mockedQuery = buildMockedAssetGraphLiveQuery(assetKeys, undefined, [
+      new GraphQLError('500'),
+    ]);
+    const mockedQuery2 = buildMockedAssetGraphLiveQuery(assetKeys, undefined, [
+      new GraphQLError('500'),
+    ]);
+
+    const resultFn = getMockResultFn(mockedQuery);
+    const resultFn2 = getMockResultFn(mockedQuery2);
+
+    const hookResult = jest.fn();
+    const hookResult2 = jest.fn();
+
+    const {rerender} = render(
+      <Test mocks={[mockedQuery, mockedQuery2]} hooks={[{keys: assetKeys, hookResult}]} />,
+    );
+
+    // Initially an empty object
+    expect(resultFn).toHaveBeenCalledTimes(0);
+    expect(hookResult.mock.calls[0]!.value).toEqual(undefined);
+
+    act(() => {
+      jest.runOnlyPendingTimers();
+    });
+
+    expect(resultFn).toHaveBeenCalled();
+    await waitFor(() => {
+      expect(hookResult.mock.calls[1]!.value).not.toEqual({});
+    });
+    expect(resultFn2).not.toHaveBeenCalled();
+
+    // Re-render with the same asset keys and expect no fetches to be made
+
+    rerender(
+      <Test
+        mocks={[mockedQuery, mockedQuery2]}
+        hooks={[{keys: assetKeys, hookResult: hookResult2}]}
+      />,
+    );
+
+    // Initially an empty object
+    expect(resultFn2).not.toHaveBeenCalled();
+    expect(hookResult2.mock.calls[0][0]).toEqual({});
+    act(() => {
+      jest.runOnlyPendingTimers();
+    });
+
+    // Not called because they failed to fetch recently
+    expect(resultFn2).not.toHaveBeenCalled();
+    expect(hookResult2.mock.calls[1]).toEqual(hookResult.mock.calls[1]);
+
+    await act(async () => {
+      await Promise.resolve();
+    });
+
+    // After the poll interval we retry the previously failed assets
+    act(() => {
+      jest.advanceTimersByTime(SUBSCRIPTION_IDLE_POLL_RATE + 1);
+    });
+    expect(resultFn2).toHaveBeenCalled();
+    expect(hookResult2.mock.calls[1]).toEqual(hookResult.mock.calls[1]);
+  });
 });

js_modules/dagster-ui/packages/ui-core/src/asset-data/__tests__/util.ts

@@ -1,3 +1,5 @@
+import {GraphQLError} from 'graphql/error';
+
 import {
   AssetKeyInput,
   buildAssetNode,
@@ -11,20 +13,35 @@ import {
   AssetGraphLiveQueryVariables,
 } from '../types/AssetLiveDataProvider.types';
 
-export function buildMockedAssetGraphLiveQuery(assetKeys: AssetKeyInput[]) {
+export function buildMockedAssetGraphLiveQuery(
+  assetKeys: AssetKeyInput[],
+  data:
+    | Parameters<
+        typeof buildQueryMock<AssetGraphLiveQuery, AssetGraphLiveQueryVariables>
+      >[0]['data']
+    | undefined = undefined,
+  errors: readonly GraphQLError[] | undefined = undefined,
+) {
   return buildQueryMock<AssetGraphLiveQuery, AssetGraphLiveQueryVariables>({
     query: ASSETS_GRAPH_LIVE_QUERY,
     variables: {
       // strip __typename
       assetKeys: assetKeys.map((assetKey) => ({path: assetKey.path})),
     },
-    data: {
-      assetNodes: assetKeys.map((assetKey) =>
-        buildAssetNode({assetKey: buildAssetKey(assetKey), id: JSON.stringify(assetKey)}),
-      ),
-      assetsLatestInfo: assetKeys.map((assetKey) =>
-        buildAssetLatestInfo({assetKey: buildAssetKey(assetKey), id: JSON.stringify(assetKey)}),
-      ),
-    },
+    data:
+      data === undefined && errors === undefined
+        ? {
+            assetNodes: assetKeys.map((assetKey) =>
+              buildAssetNode({assetKey: buildAssetKey(assetKey), id: JSON.stringify(assetKey)}),
+            ),
+            assetsLatestInfo: assetKeys.map((assetKey) =>
+              buildAssetLatestInfo({
+                assetKey: buildAssetKey(assetKey),
+                id: JSON.stringify(assetKey),
+              }),
+            ),
+          }
+        : data,
+    errors,
   });
 }