feat(gatsby): Support for off-main-thread in develop (#35648)

* feat(gatsby): initial support for off-main-thread in develop Co-authored-by: tyhopp <[email protected]> * feat(gatsby): Refactor partytown local proxy logic * Use origin instead of host in serve/develop proxy Co-authored-by: Michal Piechowiak <[email protected]>
gatsbyjs · May 16, 2022 · a8ab994 · a8ab994
1 parent 4575a3b
commit a8ab994
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 14 deletions.
diff --git a/packages/gatsby-script/src/gatsby-script.tsx b/packages/gatsby-script/src/gatsby-script.tsx
@@ -40,6 +40,7 @@ export function Script(props: ScriptProps): ReactElement | null {
     onLoad,
     onError,
   } = props || {}
+  const { collectScript } = useContext(PartytownContext)
 
   useEffect(() => {
     let script: HTMLScriptElement | null
@@ -53,6 +54,12 @@ export function Script(props: ScriptProps): ReactElement | null {
           script = injectScript(props)
         })
         break
+      case ScriptStrategy.offMainThread:
+        if (typeof window !== `undefined` && collectScript) {
+          const attributes = resolveAttributes(props)
+          collectScript(attributes)
+        }
+        break
     }
 
     return (): void => {
@@ -70,8 +77,7 @@ export function Script(props: ScriptProps): ReactElement | null {
     const inlineScript = resolveInlineScript(props)
     const attributes = resolveAttributes(props)
 
-    const { collectScript } = useContext(PartytownContext)
-    if (collectScript) {
+    if (typeof window === `undefined` && collectScript) {
       collectScript(attributes)
     }
 

diff --git a/packages/gatsby/src/commands/serve.ts b/packages/gatsby/src/commands/serve.ts
@@ -25,7 +25,10 @@ import { initTracer } from "../utils/tracer"
 import { configureTrailingSlash } from "../utils/express-middlewares"
 import { getDataStore, detectLmdbStore } from "../datastore"
 import { functionMiddlewares } from "../internal-plugins/functions/middleware"
-import proxy from "express-http-proxy"
+import {
+  partytownProxyPath,
+  partytownProxy,
+} from "../internal-plugins/partytown/proxy"
 
 process.env.GATSBY_EXPERIMENTAL_LMDB_STORE = `1`
 detectLmdbStore()
@@ -124,16 +127,7 @@ module.exports = async (program: IServeProgram): Promise<void> => {
   // Proxy gatsby-script using off-main-thread strategy
   const { partytownProxiedURLs = [] } = config || {}
 
-  app.use(
-    `/__partytown-proxy`,
-    proxy(req => new URL(req.query.url as string).host as string, {
-      filter: req => partytownProxiedURLs.some(url => req.query?.url === url),
-      proxyReqPathResolver: req => {
-        const { pathname = ``, search = `` } = new URL(req.query?.url as string)
-        return pathname + search
-      },
-    })
-  )
+  app.use(partytownProxyPath, partytownProxy(partytownProxiedURLs))
 
   // eslint-disable-next-line new-cap
   const router = express.Router()

diff --git a/packages/gatsby/src/internal-plugins/partytown/gatsby-browser.tsx b/packages/gatsby/src/internal-plugins/partytown/gatsby-browser.tsx
@@ -0,0 +1,63 @@
+import React, { ReactElement, useState } from "react"
+import type { GatsbySSR } from "gatsby"
+import { Partytown } from "@builder.io/partytown/react"
+import { PartytownContext } from "gatsby-script"
+import type { PartytownProps } from "@builder.io/partytown/react"
+
+interface ICollectedForwardsState {
+  collectedForwards: Set<string>
+  collectedAnyScript: boolean
+}
+
+function PartytownProvider({ children }): ReactElement {
+  const [{ collectedForwards, collectedAnyScript }, setState] =
+    useState<ICollectedForwardsState>({
+      collectedForwards: new Set(),
+      collectedAnyScript: false,
+    })
+
+  return (
+    <PartytownContext.Provider
+      value={{
+        collectScript: (newScript: PartytownProps): void => {
+          let stateShouldChange = false
+          const potentialNewState = {
+            collectedAnyScript,
+            collectedForwards,
+          }
+
+          if (!collectedAnyScript) {
+            potentialNewState.collectedAnyScript = true
+            stateShouldChange = true
+          }
+
+          if (newScript?.forward) {
+            if (Array.isArray(newScript.forward)) {
+              for (const singleForward of newScript.forward) {
+                if (!potentialNewState.collectedForwards.has(singleForward)) {
+                  potentialNewState.collectedForwards.add(singleForward)
+                  stateShouldChange = true
+                }
+              }
+            } else {
+              console.log(`unexpected shape of forward`, newScript)
+            }
+          }
+
+          if (stateShouldChange) {
+            setState(potentialNewState)
+          }
+        },
+      }}
+    >
+      {children}
+      {collectedAnyScript && (
+        <Partytown key="partytown" forward={Array.from(collectedForwards)} />
+      )}
+    </PartytownContext.Provider>
+  )
+}
+
+export const wrapRootElement: GatsbySSR[`wrapRootElement`] = ({ element }) => (
+  <PartytownProvider>{element}</PartytownProvider>
+)
diff --git a/packages/gatsby/src/internal-plugins/partytown/gatsby-node.ts b/packages/gatsby/src/internal-plugins/partytown/gatsby-node.ts
@@ -1,5 +1,7 @@
 import path from "path"
 import { copyLibFiles } from "@builder.io/partytown/utils"
+import { CreateDevServerArgs } from "gatsby"
+import { partytownProxyPath, partytownProxy } from "./proxy"
 
 /**
  * Copy Partytown library files to public.
@@ -24,9 +26,19 @@ exports.createPages = ({ actions, store }): void => {
     const encodedURL: string = encodeURI(host)
 
     createRedirect({
-      fromPath: `/__partytown-proxy?url=${encodedURL}`,
+      fromPath: `${partytownProxyPath}?url=${encodedURL}`,
       toPath: encodedURL,
       statusCode: 200,
     })
   }
 }
+
+export async function onCreateDevServer({
+  app,
+  store,
+}: CreateDevServerArgs): Promise<void> {
+  const { config } = store.getState()
+  const { partytownProxiedURLs = [] } = config || {}
+
+  app.use(partytownProxyPath, partytownProxy(partytownProxiedURLs))
+}
diff --git a/packages/gatsby/src/internal-plugins/partytown/proxy.ts b/packages/gatsby/src/internal-plugins/partytown/proxy.ts
@@ -0,0 +1,16 @@
+import proxy from "express-http-proxy"
+import type { RequestHandler } from "express"
+
+export const partytownProxyPath = `/__partytown-proxy`
+
+export function partytownProxy(
+  partytownProxiedURLs: Array<string>
+): RequestHandler {
+  return proxy(req => new URL(req.query.url as string).origin as string, {
+    filter: req => partytownProxiedURLs.some(url => req.query?.url === url),
+    proxyReqPathResolver: req => {
+      const { pathname = ``, search = `` } = new URL(req.query?.url as string)
+      return pathname + search
+    },
+  })
+}