refactor(gatsby): get-page-data util (#27939)

* refactor(gatsby): clear pending page-data writes per-page

* refactor(gatsby): get-page-data util

* tmp: use new implementation in websocket for tests

* tmp: fix ts error

* fix tests

* Restore behind-the-flag behavior

* Add a sanity check for page-data file existence

* reset query running state on post bootstrap

* refactor(get-page-data): don't wrap pageData in websocket specific structure, conditionally use new utility in express page-data handler

* QUERY_RUN_REQUESTED doesn't need to be redux action

* handle time period between finishing query run and flushing page-data

* feat(state-machine): add query on demand transitions (#28111)

* feat(state-machine): add query on demand transitions

* single action doesn't need to be in array

* drop code that is not used yet

Co-authored-by: Michal Piechowiak <misiek.piechowiak@gmail.com>

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

@@ -71,6 +71,7 @@ Object {
     "trackedQueries": Map {
       "/my-sweet-new-page/" => Object {
         "dirty": 1,
+        "running": 0,
       },
     },
   },

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

@@ -107,8 +107,8 @@ describe(`create page`, () => {
 
     expect(state).toMatchObject({
       trackedQueries: new Map([
-        [`/foo`, { dirty: FLAG_DIRTY_PAGE }],
-        [`/bar`, { dirty: FLAG_DIRTY_PAGE }],
+        [`/foo`, { dirty: FLAG_DIRTY_PAGE, running: 0 }],
+        [`/bar`, { dirty: FLAG_DIRTY_PAGE, running: 0 }],
       ]),
     })
   })
@@ -119,7 +119,7 @@ describe(`create page`, () => {
     expect(state.trackedQueries.get(`/foo`)?.dirty).toEqual(0) // sanity-check
     state = createPage(state, Pages.foo)
 
-    expect(state.trackedQueries.get(`/foo`)).toEqual({
+    expect(state.trackedQueries.get(`/foo`)).toMatchObject({
       dirty: 0,
     })
   })
@@ -130,7 +130,7 @@ describe(`create page`, () => {
     expect(state.trackedQueries.get(`/foo`)?.dirty).toEqual(0) // sanity-check
     state = createPage(state, Pages.foo, { contextModified: true })
 
-    expect(state.trackedQueries.get(`/foo`)).toEqual({
+    expect(state.trackedQueries.get(`/foo`)).toMatchObject({
       dirty: FLAG_DIRTY_PAGE,
     })
   })
@@ -259,7 +259,9 @@ describe(`replace static query`, () => {
     state = reducer(state, replaceStaticQuery(StaticQueries.q1))
 
     expect(state).toMatchObject({
-      trackedQueries: new Map([[`sq--q1`, { dirty: FLAG_DIRTY_TEXT }]]),
+      trackedQueries: new Map([
+        [`sq--q1`, { dirty: FLAG_DIRTY_TEXT, running: 0 }],
+      ]),
     })
   })
 
@@ -276,7 +278,7 @@ describe(`replace static query`, () => {
     expect(state.trackedQueries.get(`sq--q1`)?.dirty).toEqual(0) // sanity-check
 
     state = reducer(state, replaceStaticQuery(StaticQueries.q1))
-    expect(state.trackedQueries.get(`sq--q1`)).toEqual({
+    expect(state.trackedQueries.get(`sq--q1`)).toMatchObject({
       dirty: FLAG_DIRTY_TEXT,
     })
   })
@@ -457,22 +459,26 @@ describe(`query extraction`, () => {
   it(`marks all page queries associated with the component as dirty on the first run`, () => {
     state = reducer(state, queryExtracted(ComponentQueries.bar, {} as any))
 
-    expect(state.trackedQueries.get(`/bar`)).toEqual({
+    expect(state.trackedQueries.get(`/bar`)).toMatchObject({
       dirty: FLAG_DIRTY_PAGE | FLAG_DIRTY_TEXT,
     })
-    expect(state.trackedQueries.get(`/bar2`)).toEqual({
+    expect(state.trackedQueries.get(`/bar2`)).toMatchObject({
       dirty: FLAG_DIRTY_PAGE | FLAG_DIRTY_TEXT,
     })
     // Sanity check
-    expect(state.trackedQueries.get(`/foo`)).toEqual({ dirty: FLAG_DIRTY_PAGE })
+    expect(state.trackedQueries.get(`/foo`)).toMatchObject({
+      dirty: FLAG_DIRTY_PAGE,
+    })
   })
 
   it(`doesn't mark page query as dirty if query text didn't change`, () => {
     state = editFooQuery(state, ComponentQueries.foo)
 
-    expect(state.trackedQueries.get(`/foo`)).toEqual({ dirty: 0 })
+    expect(state.trackedQueries.get(`/foo`)).toMatchObject({ dirty: 0 })
     // sanity-check (we didn't run or extract /bar)
-    expect(state.trackedQueries.get(`/bar`)).toEqual({ dirty: FLAG_DIRTY_PAGE })
+    expect(state.trackedQueries.get(`/bar`)).toMatchObject({
+      dirty: FLAG_DIRTY_PAGE,
+    })
   })
 
   it(`doesn't mark page query as dirty if component has query extraction errors`, () => {
@@ -522,7 +528,9 @@ describe(`query extraction`, () => {
   it(`marks all page queries associated with the component as dirty when query text changes`, () => {
     state = editFooQuery(state, ComponentQueries.fooEdited)
 
-    expect(state.trackedQueries.get(`/foo`)).toEqual({ dirty: FLAG_DIRTY_TEXT })
+    expect(state.trackedQueries.get(`/foo`)).toMatchObject({
+      dirty: FLAG_DIRTY_TEXT,
+    })
   })
 
   it.skip(`marks all static queries associated with this component as dirty`, () => {

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

@@ -16,6 +16,8 @@ export const FLAG_DIRTY_DATA = 0b0100
 
 export const FLAG_ERROR_EXTRACTION = 0b0001
 
+export const FLAG_RUNNING_INFLIGHT = 0b0001
+
 const initialState = (): IGatsbyState["queries"] => {
   return {
     byNode: new Map<NodeId, Set<QueryId>>(),
@@ -30,6 +32,7 @@ const initialState = (): IGatsbyState["queries"] => {
 const initialQueryState = (): IQueryState => {
   return {
     dirty: -1, // unknown, must be set right after init
+    running: 0,
   }
 }
 
@@ -153,6 +156,10 @@ export function queriesReducer(
       const { path } = action.payload
       state = clearNodeDependencies(state, path)
       state = clearConnectionDependencies(state, path)
+      const query = state.trackedQueries.get(path)
+      if (query) {
+        query.running = setFlag(query.running, FLAG_RUNNING_INFLIGHT)
+      }
       return state
     }
     case `CREATE_NODE`:
@@ -183,6 +190,16 @@ export function queriesReducer(
       const { path } = action.payload
       const query = registerQuery(state, path)
       query.dirty = 0
+      query.running = 0 // TODO: also
+      return state
+    }
+    case `SET_PROGRAM_STATUS`: {
+      if (action.payload === `BOOTSTRAP_FINISHED`) {
+        // Reset the running state (as it could've been persisted)
+        for (const [, query] of state.trackedQueries) {
+          query.running = 0
+        }
+      }
       return state
     }
     default:

packages/gatsby/src/redux/types.ts

@@ -154,6 +154,7 @@ export interface IStateProgram extends IProgram {
 
 export interface IQueryState {
   dirty: number
+  running: number
 }
 
 export interface IComponentState {

packages/gatsby/src/services/listen-for-mutations.ts

@@ -14,13 +14,19 @@ export const listenForMutations: InvokeCallback = (callback: Sender<any>) => {
     callback({ type: `WEBHOOK_RECEIVED`, payload: event })
   }
 
+  const emitQueryRunRequest = (event: unknown): void => {
+    callback({ type: `QUERY_RUN_REQUESTED`, payload: event })
+  }
+
   emitter.on(`ENQUEUE_NODE_MUTATION`, emitMutation)
   emitter.on(`WEBHOOK_RECEIVED`, emitWebhook)
   emitter.on(`SOURCE_FILE_CHANGED`, emitSourceChange)
+  emitter.on(`QUERY_RUN_REQUESTED`, emitQueryRunRequest)
 
   return function unsubscribeFromMutationListening(): void {
     emitter.off(`ENQUEUE_NODE_MUTATION`, emitMutation)
     emitter.off(`WEBHOOK_RECEIVED`, emitWebhook)
     emitter.off(`SOURCE_FILE_CHANGED`, emitSourceChange)
+    emitter.off(`QUERY_RUN_REQUESTED`, emitQueryRunRequest)
   }
 }

packages/gatsby/src/services/types.ts

@@ -41,4 +41,5 @@ export interface IBuildContext {
   webpackListener?: Actor<unknown, AnyEventObject>
   queryFilesDirty?: boolean
   sourceFilesDirty?: boolean
+  pendingQueryRuns?: Set<string>
 }

packages/gatsby/src/state-machines/develop/actions.ts

@@ -157,6 +157,22 @@ export const panicBecauseOfInfiniteLoop: ActionFunction<
   )
 }
 
+export const trackRequestedQueryRun = assign<IBuildContext, AnyEventObject>({
+  pendingQueryRuns: (context, { payload }) => {
+    const pendingQueryRuns = context.pendingQueryRuns || new Set<string>()
+    if (payload?.pagePath) {
+      pendingQueryRuns.add(payload.pagePath)
+    }
+    return pendingQueryRuns
+  },
+})
+
+export const clearPendingQueryRuns = assign<IBuildContext>(() => {
+  return {
+    pendingQueryRuns: new Set<string>(),
+  }
+})
+
 export const buildActions: ActionFunctionMap<IBuildContext, AnyEventObject> = {
   callApi,
   markNodesDirty,
@@ -180,4 +196,6 @@ export const buildActions: ActionFunctionMap<IBuildContext, AnyEventObject> = {
   setQueryRunningFinished,
   panic,
   logError,
+  trackRequestedQueryRun,
+  clearPendingQueryRuns,
 }

packages/gatsby/src/state-machines/develop/index.ts

@@ -31,6 +31,9 @@ const developConfig: MachineConfig<IBuildContext, any, AnyEventObject> = {
       target: `reloadingData`,
       actions: `assignWebhookBody`,
     },
+    QUERY_RUN_REQUESTED: {
+      actions: `trackRequestedQueryRun`,
+    },
   },
   states: {
     // Here we handle the initial bootstrap
@@ -106,6 +109,9 @@ const developConfig: MachineConfig<IBuildContext, any, AnyEventObject> = {
         ADD_NODE_MUTATION: {
           actions: [`markNodesDirty`, `callApi`],
         },
+        QUERY_RUN_REQUESTED: {
+          actions: forwardTo(`run-queries`),
+        },
       },
       invoke: {
         id: `run-queries`,
@@ -118,6 +124,7 @@ const developConfig: MachineConfig<IBuildContext, any, AnyEventObject> = {
           gatsbyNodeGraphQLFunction,
           graphqlRunner,
           websocketManager,
+          pendingQueryRuns,
         }: IBuildContext): IQueryRunningContext => {
           return {
             program,
@@ -126,6 +133,7 @@ const developConfig: MachineConfig<IBuildContext, any, AnyEventObject> = {
             gatsbyNodeGraphQLFunction,
             graphqlRunner,
             websocketManager,
+            pendingQueryRuns,
           }
         },
         onDone: [
@@ -145,28 +153,34 @@ const developConfig: MachineConfig<IBuildContext, any, AnyEventObject> = {
             target: `recreatingPages`,
             cond: ({ nodesMutatedDuringQueryRun }: IBuildContext): boolean =>
               !!nodesMutatedDuringQueryRun,
-            actions: [`markNodesClean`, `incrementRecompileCount`],
+            actions: [
+              `markNodesClean`,
+              `incrementRecompileCount`,
+              `clearPendingQueryRuns`,
+            ],
           },
           {
             // If we have no compiler (i.e. it's first run), then spin up the
             // webpack and socket.io servers
             target: `startingDevServers`,
-            actions: `setQueryRunningFinished`,
+            actions: [`setQueryRunningFinished`, `clearPendingQueryRuns`],
             cond: ({ compiler }: IBuildContext): boolean => !compiler,
           },
           {
             // If source files have changed, then recompile the JS bundle
             target: `recompiling`,
             cond: ({ sourceFilesDirty }: IBuildContext): boolean =>
               !!sourceFilesDirty,
+            actions: [`clearPendingQueryRuns`],
           },
           {
             // ...otherwise just wait.
             target: `waiting`,
+            actions: [`clearPendingQueryRuns`],
           },
         ],
         onError: {
-          actions: `logError`,
+          actions: [`logError`, `clearPendingQueryRuns`],
           target: `waiting`,
         },
       },
@@ -205,6 +219,13 @@ const developConfig: MachineConfig<IBuildContext, any, AnyEventObject> = {
     },
     // Idle, waiting for events that make us rebuild
     waiting: {
+      always: [
+        {
+          target: `runningQueries`,
+          cond: ({ pendingQueryRuns }: IBuildContext): boolean =>
+            !!pendingQueryRuns && pendingQueryRuns.size > 0,
+        },
+      ],
       entry: [`saveDbState`, `resetRecompileCount`],
       on: {
         // Forward these events to the child machine, so it can handle batching

packages/gatsby/src/state-machines/query-running/actions.ts

@@ -1,5 +1,10 @@
 import { IQueryRunningContext } from "./types"
-import { DoneInvokeEvent, assign, ActionFunctionMap } from "xstate"
+import {
+  DoneInvokeEvent,
+  assign,
+  ActionFunctionMap,
+  AnyEventObject,
+} from "xstate"
 import { enqueueFlush } from "../../utils/page-data"
 
 export const flushPageData = (): void => {
@@ -24,9 +29,30 @@ export const markSourceFilesClean = assign<IQueryRunningContext>({
   filesDirty: false,
 })
 
+export const trackRequestedQueryRun = assign<
+  IQueryRunningContext,
+  AnyEventObject
+>({
+  pendingQueryRuns: (context, { payload }) => {
+    const pendingQueryRuns = context.pendingQueryRuns || new Set<string>()
+    if (payload?.pagePath) {
+      pendingQueryRuns.add(payload.pagePath)
+    }
+    return pendingQueryRuns
+  },
+})
+
+export const clearCurrentlyHandledPendingQueryRuns = assign<
+  IQueryRunningContext
+>({
+  currentlyHandledPendingQueryRuns: undefined,
+})
+
 export const queryActions: ActionFunctionMap<IQueryRunningContext, any> = {
   assignDirtyQueries,
   flushPageData,
   markSourceFilesDirty,
   markSourceFilesClean,
+  trackRequestedQueryRun,
+  clearCurrentlyHandledPendingQueryRuns,
 }