From 349eb0f0ad78f9cb491278eb4c7f9fe0c2e78b79 Mon Sep 17 00:00:00 2001
From: Guillaume Chau <guillaume.b.chau@gmail.com>
Date: Fri, 7 May 2021 18:00:52 +0200
Subject: [PATCH] feat: onServerPrefetch (#3070)

Support equivalent of `serverPrefetch` option via Composition API.
---
 packages/runtime-core/src/apiLifecycle.ts     |  16 +-
 packages/runtime-core/src/component.ts        |  12 +-
 packages/runtime-core/src/componentOptions.ts |   7 +-
 packages/runtime-core/src/index.ts            |   3 +-
 .../server-renderer/__tests__/render.spec.ts  | 210 +++++++++++++++++-
 packages/server-renderer/src/render.ts        |  20 +-
 6 files changed, 247 insertions(+), 21 deletions(-)

diff --git a/packages/runtime-core/src/apiLifecycle.ts b/packages/runtime-core/src/apiLifecycle.ts
index 4d7b53d36a7..7decc103a16 100644
--- a/packages/runtime-core/src/apiLifecycle.ts
+++ b/packages/runtime-core/src/apiLifecycle.ts
@@ -65,8 +65,9 @@ export function injectHook(
 export const createHook = <T extends Function = () => any>(
   lifecycle: LifecycleHooks
 ) => (hook: T, target: ComponentInternalInstance | null = currentInstance) =>
-  // post-create lifecycle registrations are noops during SSR
-  !isInSSRComponentSetup && injectHook(lifecycle, hook, target)
+  // post-create lifecycle registrations are noops during SSR (except for serverPrefetch)
+  (!isInSSRComponentSetup || lifecycle === LifecycleHooks.SERVER_PREFETCH) &&
+  injectHook(lifecycle, hook, target)
 
 export const onBeforeMount = createHook(LifecycleHooks.BEFORE_MOUNT)
 export const onMounted = createHook(LifecycleHooks.MOUNTED)
@@ -74,6 +75,7 @@ export const onBeforeUpdate = createHook(LifecycleHooks.BEFORE_UPDATE)
 export const onUpdated = createHook(LifecycleHooks.UPDATED)
 export const onBeforeUnmount = createHook(LifecycleHooks.BEFORE_UNMOUNT)
 export const onUnmounted = createHook(LifecycleHooks.UNMOUNTED)
+export const onServerPrefetch = createHook(LifecycleHooks.SERVER_PREFETCH)
 
 export type DebuggerHook = (e: DebuggerEvent) => void
 export const onRenderTriggered = createHook<DebuggerHook>(
@@ -83,15 +85,15 @@ export const onRenderTracked = createHook<DebuggerHook>(
   LifecycleHooks.RENDER_TRACKED
 )
 
-export type ErrorCapturedHook = (
-  err: unknown,
+export type ErrorCapturedHook<TError = unknown> = (
+  err: TError,
   instance: ComponentPublicInstance | null,
   info: string
 ) => boolean | void
 
-export const onErrorCaptured = (
-  hook: ErrorCapturedHook,
+export function onErrorCaptured<TError = Error>(
+  hook: ErrorCapturedHook<TError>,
   target: ComponentInternalInstance | null = currentInstance
-) => {
+) {
   injectHook(LifecycleHooks.ERROR_CAPTURED, hook, target)
 }
diff --git a/packages/runtime-core/src/component.ts b/packages/runtime-core/src/component.ts
index 54cf487bb45..aa92b3ed215 100644
--- a/packages/runtime-core/src/component.ts
+++ b/packages/runtime-core/src/component.ts
@@ -153,7 +153,7 @@ export type Component<
 
 export { ComponentOptions }
 
-type LifecycleHook = Function[] | null
+type LifecycleHook<TFn = Function> = TFn[] | null
 
 export const enum LifecycleHooks {
   BEFORE_CREATE = 'bc',
@@ -168,7 +168,8 @@ export const enum LifecycleHooks {
   ACTIVATED = 'a',
   RENDER_TRIGGERED = 'rtg',
   RENDER_TRACKED = 'rtc',
-  ERROR_CAPTURED = 'ec'
+  ERROR_CAPTURED = 'ec',
+  SERVER_PREFETCH = 'sp'
 }
 
 export interface SetupContext<E = EmitsOptions> {
@@ -414,6 +415,10 @@ export interface ComponentInternalInstance {
    * @internal
    */
   [LifecycleHooks.ERROR_CAPTURED]: LifecycleHook
+  /**
+   * @internal
+   */
+  [LifecycleHooks.SERVER_PREFETCH]: LifecycleHook<() => Promise<unknown>>
 }
 
 const emptyAppContext = createAppContext()
@@ -497,7 +502,8 @@ export function createComponentInstance(
     a: null,
     rtg: null,
     rtc: null,
-    ec: null
+    ec: null,
+    sp: null
   }
   if (__DEV__) {
     instance.ctx = createRenderContext(instance)
diff --git a/packages/runtime-core/src/componentOptions.ts b/packages/runtime-core/src/componentOptions.ts
index 366ac379dcf..40f5669b508 100644
--- a/packages/runtime-core/src/componentOptions.ts
+++ b/packages/runtime-core/src/componentOptions.ts
@@ -40,7 +40,8 @@ import {
   onDeactivated,
   onRenderTriggered,
   DebuggerHook,
-  ErrorCapturedHook
+  ErrorCapturedHook,
+  onServerPrefetch
 } from './apiLifecycle'
 import {
   reactive,
@@ -555,6 +556,7 @@ export function applyOptions(
     renderTracked,
     renderTriggered,
     errorCaptured,
+    serverPrefetch,
     // public API
     expose
   } = options
@@ -798,6 +800,9 @@ export function applyOptions(
   if (unmounted) {
     onUnmounted(unmounted.bind(publicThis))
   }
+  if (serverPrefetch) {
+    onServerPrefetch(serverPrefetch.bind(publicThis))
+  }
 
   if (__COMPAT__) {
     if (
diff --git a/packages/runtime-core/src/index.ts b/packages/runtime-core/src/index.ts
index b4b63545364..bb90574d071 100644
--- a/packages/runtime-core/src/index.ts
+++ b/packages/runtime-core/src/index.ts
@@ -37,7 +37,8 @@ export {
   onDeactivated,
   onRenderTracked,
   onRenderTriggered,
-  onErrorCaptured
+  onErrorCaptured,
+  onServerPrefetch
 } from './apiLifecycle'
 export { provide, inject } from './apiInject'
 export { nextTick } from './scheduler'
diff --git a/packages/server-renderer/__tests__/render.spec.ts b/packages/server-renderer/__tests__/render.spec.ts
index 54333e8f524..b3bb4980542 100644
--- a/packages/server-renderer/__tests__/render.spec.ts
+++ b/packages/server-renderer/__tests__/render.spec.ts
@@ -14,7 +14,9 @@ import {
   watchEffect,
   createVNode,
   resolveDynamicComponent,
-  renderSlot
+  renderSlot,
+  onErrorCaptured,
+  onServerPrefetch
 } from 'vue'
 import { escapeHtml } from '@vue/shared'
 import { renderToString } from '../src/renderToString'
@@ -859,5 +861,211 @@ function testRender(type: string, render: typeof renderToString) {
         )
       ).toBe(`<div>A</div><div>B</div>`)
     })
+
+    test('onServerPrefetch', async () => {
+      const msg = Promise.resolve('hello')
+      const app = createApp({
+        setup() {
+          const message = ref('')
+          onServerPrefetch(async () => {
+            message.value = await msg
+          })
+          return {
+            message
+          }
+        },
+        render() {
+          return h('div', this.message)
+        }
+      })
+      const html = await render(app)
+      expect(html).toBe(`<div>hello</div>`)
+    })
+
+    test('multiple onServerPrefetch', async () => {
+      const msg = Promise.resolve('hello')
+      const msg2 = Promise.resolve('hi')
+      const msg3 = Promise.resolve('bonjour')
+      const app = createApp({
+        setup() {
+          const message = ref('')
+          const message2 = ref('')
+          const message3 = ref('')
+          onServerPrefetch(async () => {
+            message.value = await msg
+          })
+          onServerPrefetch(async () => {
+            message2.value = await msg2
+          })
+          onServerPrefetch(async () => {
+            message3.value = await msg3
+          })
+          return {
+            message,
+            message2,
+            message3
+          }
+        },
+        render() {
+          return h('div', `${this.message} ${this.message2} ${this.message3}`)
+        }
+      })
+      const html = await render(app)
+      expect(html).toBe(`<div>hello hi bonjour</div>`)
+    })
+
+    test('onServerPrefetch are run in parallel', async () => {
+      const first = jest.fn(() => Promise.resolve())
+      const second = jest.fn(() => Promise.resolve())
+      let checkOther = [false, false]
+      let done = [false, false]
+      const app = createApp({
+        setup() {
+          onServerPrefetch(async () => {
+            checkOther[0] = done[1]
+            await first()
+            done[0] = true
+          })
+          onServerPrefetch(async () => {
+            checkOther[1] = done[0]
+            await second()
+            done[1] = true
+          })
+        },
+        render() {
+          return h('div', '')
+        }
+      })
+      await render(app)
+      expect(first).toHaveBeenCalled()
+      expect(second).toHaveBeenCalled()
+      expect(checkOther).toEqual([false, false])
+      expect(done).toEqual([true, true])
+    })
+
+    test('onServerPrefetch with serverPrefetch option', async () => {
+      const msg = Promise.resolve('hello')
+      const msg2 = Promise.resolve('hi')
+      const app = createApp({
+        data() {
+          return {
+            message: ''
+          }
+        },
+
+        async serverPrefetch() {
+          this.message = await msg
+        },
+
+        setup() {
+          const message2 = ref('')
+          onServerPrefetch(async () => {
+            message2.value = await msg2
+          })
+          return {
+            message2
+          }
+        },
+        render() {
+          return h('div', `${this.message} ${this.message2}`)
+        }
+      })
+      const html = await render(app)
+      expect(html).toBe(`<div>hello hi</div>`)
+    })
+
+    test('mixed in serverPrefetch', async () => {
+      const msg = Promise.resolve('hello')
+      const app = createApp({
+        data() {
+          return {
+            msg: ''
+          }
+        },
+        mixins: [
+          {
+            async serverPrefetch() {
+              this.msg = await msg
+            }
+          }
+        ],
+        render() {
+          return h('div', this.msg)
+        }
+      })
+      const html = await render(app)
+      expect(html).toBe(`<div>hello</div>`)
+    })
+
+    test('many serverPrefetch', async () => {
+      const foo = Promise.resolve('foo')
+      const bar = Promise.resolve('bar')
+      const baz = Promise.resolve('baz')
+      const app = createApp({
+        data() {
+          return {
+            foo: '',
+            bar: '',
+            baz: ''
+          }
+        },
+        mixins: [
+          {
+            async serverPrefetch() {
+              this.foo = await foo
+            }
+          },
+          {
+            async serverPrefetch() {
+              this.bar = await bar
+            }
+          }
+        ],
+        async serverPrefetch() {
+          this.baz = await baz
+        },
+        render() {
+          return h('div', `${this.foo}${this.bar}${this.baz}`)
+        }
+      })
+      const html = await render(app)
+      expect(html).toBe(`<div>foobarbaz</div>`)
+    })
+
+    test('onServerPrefetch throwing error', async () => {
+      let renderError: Error | null = null
+      let capturedError: Error | null = null
+
+      const Child = {
+        setup() {
+          onServerPrefetch(async () => {
+            throw new Error('An error')
+          })
+        },
+        render() {
+          return h('span')
+        }
+      }
+
+      const app = createApp({
+        setup() {
+          onErrorCaptured(e => {
+            capturedError = e
+            return false
+          })
+        },
+        render() {
+          return h('div', h(Child))
+        }
+      })
+
+      try {
+        await render(app)
+      } catch (e) {
+        renderError = e
+      }
+      expect(renderError).toBe(null)
+      expect(((capturedError as unknown) as Error).message).toBe('An error')
+    })
   })
 }
diff --git a/packages/server-renderer/src/render.ts b/packages/server-renderer/src/render.ts
index f820fa305af..eaba85605e5 100644
--- a/packages/server-renderer/src/render.ts
+++ b/packages/server-renderer/src/render.ts
@@ -2,7 +2,6 @@ import {
   Comment,
   Component,
   ComponentInternalInstance,
-  ComponentOptions,
   DirectiveBinding,
   Fragment,
   mergeProps,
@@ -87,13 +86,18 @@ export function renderComponentVNode(
   const instance = createComponentInstance(vnode, parentComponent, null)
   const res = setupComponent(instance, true /* isSSR */)
   const hasAsyncSetup = isPromise(res)
-  const prefetch = (vnode.type as ComponentOptions).serverPrefetch
-  if (hasAsyncSetup || prefetch) {
-    let p = hasAsyncSetup ? (res as Promise<void>) : Promise.resolve()
-    if (prefetch) {
-      p = p.then(() => prefetch.call(instance.proxy)).catch(err => {
-        warn(`[@vue/server-renderer]: Uncaught error in serverPrefetch:\n`, err)
-      })
+  const prefetches = instance.sp
+  if (hasAsyncSetup || prefetches) {
+    let p: Promise<unknown> = hasAsyncSetup
+      ? (res as Promise<void>)
+      : Promise.resolve()
+    if (prefetches) {
+      p = p
+        .then(() =>
+          Promise.all(prefetches.map(prefetch => prefetch.call(instance.proxy)))
+        )
+        // Note: error display is already done by the wrapped lifecycle hook function.
+        .catch(() => {})
     }
     return p.then(() => renderComponentSubTree(instance, slotScopeId))
   } else {