From f1cb3f21ad909230d28d1491dbc29e3c8736ac96 Mon Sep 17 00:00:00 2001
From: Isaac Hellendag <2823852+hellendag@users.noreply.github.com>
Date: Thu, 12 Oct 2023 10:25:56 -0500
Subject: [PATCH] [ui] Improve global search perf (#16965)

## Summary & Motivation

Improve render and interaction performance on global search for users
with very large workspaces.

A key part of the problem here is that writing/reading with the Apollo
cache becomes prohibitively slow at scale. Therefore:

- Use `cache-first` policy on search queries. Currently, we always hit
the backend when opening search, which can be expensive both in terms of
performing the query and updating the cache. Instead, we can leverage
the workspace and asset catalog queries to read data from the Apollo
cache, if available.
- Only query or read the cache once per search query. If we already have
WebWorkers set up for search, don't repeat the query at all, even to do
a cache lookup -- the data is already available on the workers. Just go
straight to the existing workers, thus making search instantly
available.
- Skip the Apollo cache entirely for the run timeline query. When the
user has a huge number of runs, the Apollo cache read/write can be very
slow, blocking the main thread and affecting everything else on the
page. Since this data needs to be fresh in all cases anyway, just skip
the cache.

## How I Tested These Changes

Viewing the app as a user with tens of thousands of objects, open search
after the workspace and asset catalog have loaded. Verify that new
queries are not performed, and that the data is available to be searched
as soon as the cache has populated the values. Close search, reopen it.
Verify that search is available right away, without having to wait for
the cache to populate the values.
---
 .../ui-core/src/runs/useRunsForTimeline.tsx   |  4 ++++
 .../ui-core/src/search/useGlobalSearch.tsx    | 22 +++++++++++++++----
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/runs/useRunsForTimeline.tsx b/js_modules/dagster-ui/packages/ui-core/src/runs/useRunsForTimeline.tsx
index 1bc5815606201..c491028db0815 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/runs/useRunsForTimeline.tsx
+++ b/js_modules/dagster-ui/packages/ui-core/src/runs/useRunsForTimeline.tsx
@@ -23,6 +23,10 @@ export const useRunsForTimeline = (range: [number, number], runsFilter: RunsFilt
 
   const queryData = useQuery<RunTimelineQuery, RunTimelineQueryVariables>(RUN_TIMELINE_QUERY, {
     notifyOnNetworkStatusChange: true,
+    // With a very large number of runs, operating on the Apollo cache is too expensive and
+    // can block the main thread. This data has to be up-to-the-second fresh anyway, so just
+    // skip the cache entirely.
+    fetchPolicy: 'no-cache',
     variables: {
       inProgressFilter: {
         ...runsFilter,
diff --git a/js_modules/dagster-ui/packages/ui-core/src/search/useGlobalSearch.tsx b/js_modules/dagster-ui/packages/ui-core/src/search/useGlobalSearch.tsx
index 5bd03b11e2fb9..f3a94f1bbaf71 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/search/useGlobalSearch.tsx
+++ b/js_modules/dagster-ui/packages/ui-core/src/search/useGlobalSearch.tsx
@@ -176,10 +176,12 @@ const EMPTY_RESPONSE = {queryString: '', results: []};
  * A `terminate` function is provided, but it's probably not necessary to use it.
  */
 export const useGlobalSearch = () => {
-  const primarySearch = React.useRef<WorkerSearchResult>();
-  const secondarySearch = React.useRef<WorkerSearchResult>();
+  const primarySearch = React.useRef<WorkerSearchResult | null>(null);
+  const secondarySearch = React.useRef<WorkerSearchResult | null>(null);
 
   const primary = useLazyQuery<SearchPrimaryQuery>(SEARCH_PRIMARY_QUERY, {
+    // Try to make aggressive use of workspace values from the Apollo cache.
+    fetchPolicy: 'cache-first',
     onCompleted: (data: SearchPrimaryQuery) => {
       const results = primaryDataToSearchResults({data});
       if (!primarySearch.current) {
@@ -190,6 +192,8 @@ export const useGlobalSearch = () => {
   });
 
   const secondary = useLazyQuery<SearchSecondaryQuery>(SEARCH_SECONDARY_QUERY, {
+    // As above, try to aggressively use asset information from Apollo cache if possible.
+    fetchPolicy: 'cache-first',
     onCompleted: (data: SearchSecondaryQuery) => {
       const results = secondaryDataToSearchResults({data});
       if (!secondarySearch.current) {
@@ -202,9 +206,14 @@ export const useGlobalSearch = () => {
   const [performPrimaryLazyQuery, primaryResult] = primary;
   const [performSecondaryLazyQuery, secondaryResult] = secondary;
 
+  // If we already have WebWorkers set up, initialization is complete and this will be a no-op.
   const initialize = React.useCallback(async () => {
-    performPrimaryLazyQuery();
-    performSecondaryLazyQuery();
+    if (!primarySearch.current) {
+      performPrimaryLazyQuery();
+    }
+    if (!secondarySearch.current) {
+      performSecondaryLazyQuery();
+    }
   }, [performPrimaryLazyQuery, performSecondaryLazyQuery]);
 
   const searchPrimary = React.useCallback(async (queryString: string) => {
@@ -215,9 +224,14 @@ export const useGlobalSearch = () => {
     return secondarySearch.current ? secondarySearch.current.search(queryString) : EMPTY_RESPONSE;
   }, []);
 
+  // Terminate the workers. Be careful with this: for users with very large workspaces, we should
+  // avoid constantly re-querying and restarting the threads. It should only be used when we know
+  // that there is fresh data to repopulate search.
   const terminate = React.useCallback(() => {
     primarySearch.current?.terminate();
+    primarySearch.current = null;
     secondarySearch.current?.terminate();
+    secondarySearch.current = null;
   }, []);
 
   return {