perf(gatsby): fix performance regression with query dependency cleaning (#28032)

* perf(gatsby): fix performance regression with query dependency cleaning

* update snapshot

* More consistent codestyle

(cherry picked from commit de5517b)

Author: vladar

packages/gatsby/src/redux/__tests__/__snapshots__/index.js.snap

@@ -58,6 +58,7 @@ Object {
     "byConnection": Map {},
     "byNode": Map {},
     "deletedQueries": Set {},
+    "queryNodes": Map {},
     "trackedComponents": Map {
       "/Users/username/dev/site/src/templates/my-sweet-new-page.js" => Object {
         "componentPath": "/Users/username/dev/site/src/templates/my-sweet-new-page.js",

packages/gatsby/src/redux/reducers/__tests__/__snapshots__/queries.ts.snap

@@ -13,6 +13,11 @@ Object {
     },
   },
   "deletedQueries": Set {},
+  "queryNodes": Map {
+    "/hi/" => Set {
+      "SuperCoolNode",
+    },
+  },
   "trackedComponents": Map {},
   "trackedQueries": Map {},
 }

packages/gatsby/src/redux/reducers/__tests__/queries.ts

@@ -88,6 +88,7 @@ it(`has expected initial state`, () => {
       "byConnection": Map {},
       "byNode": Map {},
       "deletedQueries": Set {},
+      "queryNodes": Map {},
       "trackedComponents": Map {},
       "trackedQueries": Map {},
     }

packages/gatsby/src/redux/reducers/queries.ts

@@ -20,6 +20,7 @@ const initialState = (): IGatsbyState["queries"] => {
   return {
     byNode: new Map<NodeId, Set<QueryId>>(),
     byConnection: new Map<ConnectionName, Set<QueryId>>(),
+    queryNodes: new Map<QueryId, Set<NodeId>>(),
     trackedQueries: new Map<QueryId, IQueryState>(),
     trackedComponents: new Map<ComponentPath, IComponentState>(),
     deletedQueries: new Set<QueryId>(),
@@ -88,12 +89,8 @@ export function queriesReducer(
         for (const component of state.trackedComponents.values()) {
           component.pages.delete(queryId)
         }
-        for (const nodeQueries of state.byNode.values()) {
-          nodeQueries.delete(queryId)
-        }
-        for (const connectionQueries of state.byConnection.values()) {
-          connectionQueries.delete(queryId)
-        }
+        state = clearNodeDependencies(state, queryId)
+        state = clearConnectionDependencies(state, queryId)
         state.trackedQueries.delete(queryId)
       }
       state.deletedQueries.clear()
@@ -110,12 +107,12 @@ export function queriesReducer(
       }
       if (component.query !== query) {
         // Invalidate all pages associated with a component when query text changes
-        component.pages.forEach(queryId => {
+        for (const queryId of component.pages) {
           const query = state.trackedQueries.get(queryId)
           if (query) {
             query.dirty = setFlag(query.dirty, FLAG_DIRTY_TEXT)
           }
-        })
+        }
         component.query = query
       }
       return state
@@ -144,27 +141,18 @@ export function queriesReducer(
     case `CREATE_COMPONENT_DEPENDENCY`: {
       const { path: queryId, nodeId, connection } = action.payload
       if (nodeId) {
-        const queryIds = state.byNode.get(nodeId) ?? new Set<QueryId>()
-        queryIds.add(queryId)
-        state.byNode.set(nodeId, queryIds)
+        state = addNodeDependency(state, queryId, nodeId)
       }
       if (connection) {
-        const queryIds =
-          state.byConnection.get(connection) ?? new Set<QueryId>()
-        queryIds.add(queryId)
-        state.byConnection.set(connection, queryIds)
+        state = addConnectionDependency(state, queryId, connection)
       }
       return state
     }
     case `QUERY_START`: {
       // Reset data dependencies as they will be updated when running the query
       const { path } = action.payload
-      state.byNode.forEach(queryIds => {
-        queryIds.delete(path)
-      })
-      state.byConnection.forEach(queryIds => {
-        queryIds.delete(path)
-      })
+      state = clearNodeDependencies(state, path)
+      state = clearConnectionDependencies(state, path)
       return state
     }
     case `CREATE_NODE`:
@@ -177,18 +165,18 @@ export function queriesReducer(
       const queriesByConnection =
         state.byConnection.get(node.internal.type) ?? []
 
-      queriesByNode.forEach(queryId => {
+      for (const queryId of queriesByNode) {
         const query = state.trackedQueries.get(queryId)
         if (query) {
           query.dirty = setFlag(query.dirty, FLAG_DIRTY_DATA)
         }
-      })
-      queriesByConnection.forEach(queryId => {
+      }
+      for (const queryId of queriesByConnection) {
         const query = state.trackedQueries.get(queryId)
         if (query) {
           query.dirty = setFlag(query.dirty, FLAG_DIRTY_DATA)
         }
-      })
+      }
       return state
     }
     case `PAGE_QUERY_RUN`: {
@@ -213,6 +201,69 @@ export function hasFlag(allFlags: number, flag: number): boolean {
   return allFlags >= 0 && (allFlags & flag) > 0
 }
 
+function addNodeDependency(
+  state: IGatsbyState["queries"],
+  queryId: QueryId,
+  nodeId: NodeId
+): IGatsbyState["queries"] {
+  // Perf: using two-side maps.
+  //   Without additional `queryNodes` map we would have to loop through
+  //   all existing nodes in `clearNodeDependencies` to delete node <-> query dependency
+  let nodeQueries = state.byNode.get(nodeId)
+  if (!nodeQueries) {
+    nodeQueries = new Set<QueryId>()
+    state.byNode.set(nodeId, nodeQueries)
+  }
+  let queryNodes = state.queryNodes.get(queryId)
+  if (!queryNodes) {
+    queryNodes = new Set<NodeId>()
+    state.queryNodes.set(queryId, queryNodes)
+  }
+  nodeQueries.add(queryId)
+  queryNodes.add(nodeId)
+  return state
+}
+
+function addConnectionDependency(
+  state: IGatsbyState["queries"],
+  queryId: QueryId,
+  connection: ConnectionName
+): IGatsbyState["queries"] {
+  // Note: not using two-side maps for connections as associated overhead
+  //   for small number of elements is greater then benefits, so no perf. gains
+  let queryIds = state.byConnection.get(connection)
+  if (!queryIds) {
+    queryIds = new Set()
+    state.byConnection.set(connection, queryIds)
+  }
+  queryIds.add(queryId)
+  return state
+}
+
+function clearNodeDependencies(
+  state: IGatsbyState["queries"],
+  queryId: QueryId
+): IGatsbyState["queries"] {
+  const queryNodeIds = state.queryNodes.get(queryId) ?? new Set()
+  for (const nodeId of queryNodeIds) {
+    const nodeQueries = state.byNode.get(nodeId)
+    if (nodeQueries) {
+      nodeQueries.delete(queryId)
+    }
+  }
+  return state
+}
+
+function clearConnectionDependencies(
+  state: IGatsbyState["queries"],
+  queryId: QueryId
+): IGatsbyState["queries"] {
+  for (const [, queryIds] of state.byConnection) {
+    queryIds.delete(queryId)
+  }
+  return state
+}
+
 function registerQuery(
   state: IGatsbyState["queries"],
   queryId: QueryId

packages/gatsby/src/redux/types.ts

@@ -210,6 +210,7 @@ export interface IGatsbyState {
   queries: {
     byNode: Map<Identifier, Set<Identifier>>
     byConnection: Map<string, Set<Identifier>>
+    queryNodes: Map<Identifier, Set<Identifier>>
     trackedQueries: Map<Identifier, IQueryState>
     trackedComponents: Map<string, IComponentState>
     deletedQueries: Set<Identifier>