From 0ade33526565379e0297be5996375cd073b3fe66 Mon Sep 17 00:00:00 2001
From: yoho <panyoho@gmail.com>
Date: Tue, 8 Mar 2022 18:24:29 +0800
Subject: [PATCH] feat(worker): bundle worker emit asset file (#6599)

---
 .../playground/wasm/__tests__/wasm.spec.ts    |   4 ++++
 packages/playground/wasm/add.wasm             | Bin 0 -> 89 bytes
 packages/playground/wasm/index.html           |  15 ++++++++++++
 packages/playground/wasm/vite.config.ts       |   8 +++++++
 packages/playground/wasm/worker.js            |   5 ++++
 .../worker/__tests__/worker.spec.ts           |  22 +++++++-----------
 packages/playground/worker/index.html         |  15 ++++++++++++
 packages/playground/worker/sub-worker.js      |  13 +++++++++++
 packages/playground/worker/vite.config.ts     |   4 +---
 .../playground/worker/worker-nested-worker.js |  13 +++++++++++
 packages/vite/src/node/plugins/worker.ts      |  21 ++++++++++++++---
 .../src/node/plugins/workerImportMetaUrl.ts   |   2 +-
 12 files changed, 102 insertions(+), 20 deletions(-)
 create mode 100644 packages/playground/wasm/add.wasm
 create mode 100644 packages/playground/wasm/vite.config.ts
 create mode 100644 packages/playground/wasm/worker.js
 create mode 100644 packages/playground/worker/sub-worker.js
 create mode 100644 packages/playground/worker/worker-nested-worker.js

diff --git a/packages/playground/wasm/__tests__/wasm.spec.ts b/packages/playground/wasm/__tests__/wasm.spec.ts
index b292f0ddc6e7f8..112617212251fa 100644
--- a/packages/playground/wasm/__tests__/wasm.spec.ts
+++ b/packages/playground/wasm/__tests__/wasm.spec.ts
@@ -9,3 +9,7 @@ test('should work when output', async () => {
   await page.click('.output-wasm .run')
   await untilUpdated(() => page.textContent('.output-wasm .result'), '24')
 })
+
+test('should work when wasm in worker', async () => {
+  await untilUpdated(() => page.textContent('.worker-wasm .result'), '3')
+})
diff --git a/packages/playground/wasm/add.wasm b/packages/playground/wasm/add.wasm
new file mode 100644
index 0000000000000000000000000000000000000000..76e95b227902b0423936d9e949419cb0a97c18af
GIT binary patch
literal 89
zcmZQbEY4+QU|?WmXG~zKuV<`hW@Q0#8Q5i5m=jY{7#LW9q!9xnh%#ni0#POmOk9dA
O><kJF3XECYxM%=H?+eoa

literal 0
HcmV?d00001

diff --git a/packages/playground/wasm/index.html b/packages/playground/wasm/index.html
index 27bdc8167e064e..ecb0b66e913fbb 100644
--- a/packages/playground/wasm/index.html
+++ b/packages/playground/wasm/index.html
@@ -12,9 +12,20 @@ <h3>When wasm is output, result should be 24</h3>
   <span class="result"></span>
 </div>
 
+<div class="worker-wasm">
+  <h3>worker wasm</h3>
+  <span class="result"></span>
+</div>
+
 <script type="module">
   import light from './light.wasm'
   import heavy from './heavy.wasm'
+  import myWorker from './worker?worker'
+
+  const w = new myWorker()
+  w.addEventListener('message', (ev) => {
+    text('.worker-wasm .result', ev.data.result)
+  })
 
   async function testWasm(init, resultElement) {
     const { exported_func } = await init({
@@ -25,6 +36,10 @@ <h3>When wasm is output, result should be 24</h3>
     exported_func()
   }
 
+  function text(el, text) {
+    document.querySelector(el).textContent = text
+  }
+
   document
     .querySelector('.inline-wasm .run')
     .addEventListener('click', async () =>
diff --git a/packages/playground/wasm/vite.config.ts b/packages/playground/wasm/vite.config.ts
new file mode 100644
index 00000000000000..43833d2f95d302
--- /dev/null
+++ b/packages/playground/wasm/vite.config.ts
@@ -0,0 +1,8 @@
+import { defineConfig } from 'vite'
+export default defineConfig({
+  build: {
+    // make can no emit light.wasm
+    // and emit add.wasm
+    assetsInlineLimit: 80
+  }
+})
diff --git a/packages/playground/wasm/worker.js b/packages/playground/wasm/worker.js
new file mode 100644
index 00000000000000..a483865bd42bff
--- /dev/null
+++ b/packages/playground/wasm/worker.js
@@ -0,0 +1,5 @@
+import init from './add.wasm'
+init().then((exports) => {
+  // eslint-disable-next-line no-undef
+  self.postMessage({ result: exports.add(1, 2) })
+})
diff --git a/packages/playground/worker/__tests__/worker.spec.ts b/packages/playground/worker/__tests__/worker.spec.ts
index 1c33e6ec68daf9..6d93e810c0c510 100644
--- a/packages/playground/worker/__tests__/worker.spec.ts
+++ b/packages/playground/worker/__tests__/worker.spec.ts
@@ -51,12 +51,20 @@ test.concurrent.each([[true], [false]])('shared worker', async (doTick) => {
   await waitSharedWorkerTick(page)
 })
 
+test('worker emitted', async () => {
+  await untilUpdated(() => page.textContent('.nested-worker'), 'pong')
+  await untilUpdated(
+    () => page.textContent('.nested-worker-dynamic-import'),
+    '"msg":"pong"'
+  )
+})
+
 if (isBuild) {
   const assetsDir = path.resolve(testDir, 'dist/assets')
   // assert correct files
   test('inlined code generation', async () => {
     const files = fs.readdirSync(assetsDir)
-    expect(files.length).toBe(8)
+    expect(files.length).toBe(11)
     const index = files.find((f) => f.includes('index'))
     const content = fs.readFileSync(path.resolve(assetsDir, index), 'utf-8')
     const worker = files.find((f) => f.includes('my-worker'))
@@ -75,18 +83,6 @@ if (isBuild) {
     expect(content).toMatch(`(window.URL||window.webkitURL).createObjectURL`)
     expect(content).toMatch(`window.Blob`)
   })
-
-  test('worker need bundle', () => {
-    fs.readdirSync(assetsDir)
-      .filter(
-        (file) =>
-          file.includes('url-worker') || file.includes('url-shared-worker')
-      )
-      .forEach((file) => {
-        const content = fs.readFileSync(path.resolve(assetsDir, file), 'utf-8')
-        expect(content.startsWith('(function(){')).toBe(true)
-      })
-  })
 }
 
 test('classic worker is run', async () => {
diff --git a/packages/playground/worker/index.html b/packages/playground/worker/index.html
index 10207453712aa2..b3525da299ff5a 100644
--- a/packages/playground/worker/index.html
+++ b/packages/playground/worker/index.html
@@ -26,6 +26,9 @@
 <p>new SharedWorker(new Url('path', import.meta.url), { type: 'module' })</p>
 <div class="shared-worker-import-meta-url"></div>
 
+<p>nested worker</p>
+<div class="nested-worker"></div>
+<div class="nested-worker-dynamic-import"></div>
 <p>new Worker(new Url('path', import.meta.url))</p>
 <div class="classic-worker"></div>
 
@@ -37,6 +40,7 @@
   import InlineWorker from './my-worker?worker&inline'
   import mySharedWorker from './my-shared-worker?sharedworker&name=shared'
   import TSOutputWorker from './possible-ts-output-worker?worker'
+  import NestedWorker from './worker-nested-worker?worker'
   import { mode } from './workerImport'
   import './classic-worker'
 
@@ -114,4 +118,15 @@
     )
   })
   w2.port.start()
+
+  const nestedWorker = new NestedWorker()
+  nestedWorker.addEventListener('message', (ev) => {
+    console.log(ev)
+    if (typeof ev.data === 'string') {
+      text('.nested-worker', ev.data)
+    } else {
+      text('.nested-worker-dynamic-import', JSON.stringify(ev.data))
+    }
+  })
+  nestedWorker.postMessage('ping')
 </script>
diff --git a/packages/playground/worker/sub-worker.js b/packages/playground/worker/sub-worker.js
new file mode 100644
index 00000000000000..ab64b3667099bb
--- /dev/null
+++ b/packages/playground/worker/sub-worker.js
@@ -0,0 +1,13 @@
+self.onmessage = (event) => {
+  if (event.data === 'ping') {
+    self.postMessage('pong')
+  }
+}
+const data = import('./workerImport')
+data.then((data) => {
+  const { mode, msg } = data
+  self.postMessage({
+    mode,
+    msg
+  })
+})
diff --git a/packages/playground/worker/vite.config.ts b/packages/playground/worker/vite.config.ts
index 6bf980f46e2362..6cef7d9cea0bed 100644
--- a/packages/playground/worker/vite.config.ts
+++ b/packages/playground/worker/vite.config.ts
@@ -2,10 +2,8 @@ import vueJsx from '@vitejs/plugin-vue-jsx'
 import { defineConfig } from 'vite'
 
 export default defineConfig({
-  build: {
-    target: process.env.NODE_ENV === 'production' ? 'chrome60' : 'esnext'
-  },
   worker: {
+    format: 'es',
     plugins: [vueJsx()]
   }
 })
diff --git a/packages/playground/worker/worker-nested-worker.js b/packages/playground/worker/worker-nested-worker.js
new file mode 100644
index 00000000000000..6d4d1e4969005f
--- /dev/null
+++ b/packages/playground/worker/worker-nested-worker.js
@@ -0,0 +1,13 @@
+import SubWorker from './sub-worker?worker'
+
+const subWorker = new SubWorker()
+
+self.onmessage = (event) => {
+  if (event.data === 'ping') {
+    subWorker.postMessage('ping')
+  }
+}
+
+subWorker.onmessage = (event) => {
+  self.postMessage(event.data)
+}
diff --git a/packages/vite/src/node/plugins/worker.ts b/packages/vite/src/node/plugins/worker.ts
index 5ce0936748ea05..70bfd916e21cab 100644
--- a/packages/vite/src/node/plugins/worker.ts
+++ b/packages/vite/src/node/plugins/worker.ts
@@ -10,6 +10,7 @@ import { onRollupWarning } from '../build'
 const WorkerFileId = 'worker_file'
 
 export async function bundleWorkerEntry(
+  ctx: Rollup.TransformPluginContext,
   config: ResolvedConfig,
   id: string
 ): Promise<Buffer> {
@@ -27,11 +28,25 @@ export async function bundleWorkerEntry(
   })
   let code: string
   try {
-    const { output } = await bundle.generate({
+    const {
+      output: [outputCode, ...outputChunks]
+    } = await bundle.generate({
       format,
       sourcemap: config.build.sourcemap
     })
-    code = output[0].code
+    code = outputCode.code
+    outputChunks.forEach((outputChunk) => {
+      if (outputChunk.type === 'asset') {
+        ctx.emitFile(outputChunk)
+      }
+      if (outputChunk.type === 'chunk') {
+        ctx.emitFile({
+          fileName: `${config.build.assetsDir}/${outputChunk.fileName}`,
+          source: outputChunk.code,
+          type: 'asset'
+        })
+      }
+    })
   } finally {
     await bundle.close()
   }
@@ -72,7 +87,7 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
 
       let url: string
       if (isBuild) {
-        const code = await bundleWorkerEntry(config, id)
+        const code = await bundleWorkerEntry(this, config, id)
         if (query.inline != null) {
           const { format } = config.worker
           const workerOptions = format === 'es' ? '{type: "module"}' : '{}'
diff --git a/packages/vite/src/node/plugins/workerImportMetaUrl.ts b/packages/vite/src/node/plugins/workerImportMetaUrl.ts
index 023b567ab9fb45..837bbd45aa697d 100644
--- a/packages/vite/src/node/plugins/workerImportMetaUrl.ts
+++ b/packages/vite/src/node/plugins/workerImportMetaUrl.ts
@@ -138,7 +138,7 @@ export function workerImportMetaUrlPlugin(config: ResolvedConfig): Plugin {
           const file = path.resolve(path.dirname(id), rawUrl.slice(1, -1))
           let url: string
           if (isBuild) {
-            const content = await bundleWorkerEntry(config, file)
+            const content = await bundleWorkerEntry(this, config, file)
             const basename = path.parse(cleanUrl(file)).name
             const contentHash = getAssetHash(content)
             const fileName = path.posix.join(