Skip to content

Commit

Permalink
feat(web-worker): refactor into small peaces, add SharedWorker support
Browse files Browse the repository at this point in the history
  • Loading branch information
sheremet-va committed Dec 8, 2022
1 parent 54bf045 commit 8bb6b80
Show file tree
Hide file tree
Showing 14 changed files with 540 additions and 262 deletions.
22 changes: 14 additions & 8 deletions packages/vite-node/src/client.ts
Expand Up @@ -180,10 +180,6 @@ export class ViteNodeRunner {
const promise = this.directRequest(id, fsPath, callstack)
Object.assign(mod, { promise, evaluated: false })

promise.finally(() => {
mod.evaluated = true
})

return await promise
}

Expand Down Expand Up @@ -263,9 +259,14 @@ export class ViteNodeRunner {

if (externalize) {
debugNative(externalize)
const exports = await this.interopedImport(externalize)
mod.exports = exports
return exports
try {
const exports = await this.interopedImport(externalize)
mod.exports = exports
return exports
}
finally {
mod.evaluated = true
}
}

if (transformed == null)
Expand Down Expand Up @@ -363,7 +364,12 @@ export class ViteNodeRunner {
columnOffset: -codeDefinition.length,
})

await fn(...Object.values(context))
try {
await fn(...Object.values(context))
}
finally {
mod.evaluated = true
}

return exports
}
Expand Down
31 changes: 27 additions & 4 deletions packages/web-worker/README.md
Expand Up @@ -2,7 +2,14 @@

> Web Worker support for Vitest testing. Doesn't require JSDom.
Simulates Web Worker, but in the same thread. Supports both `new Worker(url)` and `import from './worker?worker`.
Simulates Web Worker, but in the same thread.

Supported:

- `new Worker(path)`
- `new SharedWorker(path)`
- `import MyWorker from './worker?worker'`
- `import MySharedWorker from './worker?sharedworker'`

## Installing

Expand Down Expand Up @@ -33,6 +40,22 @@ export default defineConfig({
})
```

You can also import `defineWebWorkers` from `@vitest/web-worker/pure` to defined workers, whenever you need:

```js
import { defineWebWorkers } from '@vitest/web-worker/pure'

if (process.env.SUPPORT_WORKERS)
defineWebWorkers({ clone: 'none' })
```

It accepts options:

- `clone`: `'native' | 'ponyfill' | 'none'`. Defines how should `Worker` clone message, when transferring data. Applies only to `Worker` communication. `SharedWorker` uses `MessageChannel` from Node's `worker_threads` module, and is not configurable.

> **Note**
> Requires Node 17, if you want to use native `structuredClone`. Otherwise, it fallbacks to [polyfill](https://github.com/ungap/structured-clone), if not specified as `none`. You can also configure this option with `VITEST_WEB_WORKER_CLONE` environmental variable.
## Examples

```ts
Expand All @@ -59,8 +82,8 @@ worker.onmessage = (e) => {

## Notes

- Does not support `onmessage = () => {}`. Please, use `self.onmessage = () => {}`.
- Worker does not support `onmessage = () => {}`. Please, use `self.onmessage = () => {}`.
- Shared worker does not support `onconnect = () => {}`. Please, use `self.onconnect = () => {}`.
- Transferring Buffer will not change its `byteLength`.
- You have access to shared global space as your tests.
- Requires Node 17, if you want to use native `structuredClone`. Otherwise, it fallbacks to [polyfill](https://github.com/ungap/structured-clone). You can configure this behavior by passing down `clone` option (`'native' | 'ponyfill' | 'none'`) to `defineWebWorker` or using `VITEST_WEB_WORKER_CLONE` environmental variable.
- If something is wrong, you can debug your worker, using `DEBUG=vitest:web-worker` environmental variable.
- You can debug your worker, using `DEBUG=vitest:web-worker` environmental variable.
5 changes: 3 additions & 2 deletions packages/web-worker/pure.d.ts
Expand Up @@ -2,6 +2,7 @@ declare type CloneOption = 'native' | 'ponyfill' | 'none';
interface DefineWorkerOptions {
clone: CloneOption;
}
declare function defineWebWorker(options?: DefineWorkerOptions): void;

export { defineWebWorker };
declare function defineWebWorkers(options?: DefineWorkerOptions): void;

export { defineWebWorkers };
4 changes: 2 additions & 2 deletions packages/web-worker/src/index.ts
@@ -1,3 +1,3 @@
import { defineWebWorker } from './pure'
import { defineWebWorkers } from './pure'

defineWebWorker()
defineWebWorkers()
258 changes: 12 additions & 246 deletions packages/web-worker/src/pure.ts
@@ -1,253 +1,19 @@
/* eslint-disable no-restricted-imports */
import { VitestRunner } from 'vitest/node'
import type { WorkerGlobalState } from 'vitest'
import ponyfillStructuredClone from '@ungap/structured-clone'
import { toFilePath } from 'vite-node/utils'
import createDebug from 'debug'
import { createWorkerConstructor } from './worker'
import type { DefineWorkerOptions } from './types'
import { assertGlobalExists } from './utils'
import { createSharedWorkerConstructor } from './shared-worker'

const debug = createDebug('vitest:web-worker')
export function defineWebWorkers(options?: DefineWorkerOptions) {
if (typeof Worker === 'undefined' || !('__VITEST_WEB_WORKER__' in globalThis.Worker)) {
assertGlobalExists('EventTarget')
assertGlobalExists('MessageEvent')

function getWorkerState(): WorkerGlobalState {
// @ts-expect-error untyped global
return globalThis.__vitest_worker__
}

type Procedure = (...args: any[]) => void
type CloneOption = 'native' | 'ponyfill' | 'none'

interface DefineWorkerOptions {
clone: CloneOption
}

interface InlineWorkerContext {
onmessage: Procedure | null
name?: string
close: () => void
dispatchEvent: (e: Event) => void
addEventListener: (e: string, fn: Procedure) => void
removeEventListener: (e: string, fn: Procedure) => void
postMessage: (data: any, transfer?: Transferable[] | StructuredSerializeOptions) => void
self: InlineWorkerContext
global: InlineWorkerContext
importScripts?: any
}

class InlineWorkerRunner extends VitestRunner {
constructor(options: any, private context: InlineWorkerContext) {
super(options)
}

prepareContext(context: Record<string, any>) {
const ctx = super.prepareContext(context)
// not supported for now, we can't synchronously load modules
const importScripts = () => {
throw new Error('[vitest] `importScripts` is not supported in Vite workers. Please, consider using `import` instead.')
}
return Object.assign(ctx, this.context, {
importScripts,
})
}
}

function assertGlobalExists(name: string) {
if (!(name in globalThis))
throw new Error(`[@vitest/web-worker] Cannot initiate a custom Web Worker. "${name}" is not supported in this environment. Please, consider using jsdom or happy-dom environment.`)
}

function createClonedMessageEvent(data: any, transferOrOptions: StructuredSerializeOptions | Transferable[] | undefined, clone: CloneOption) {
const transfer = Array.isArray(transferOrOptions) ? transferOrOptions : transferOrOptions?.transfer

debug('clone worker message %o', data)
const origin = typeof location === 'undefined' ? undefined : location.origin

if (typeof structuredClone === 'function' && clone === 'native') {
debug('create message event, using native structured clone')
return new MessageEvent('message', {
data: structuredClone(data, { transfer }),
origin,
})
}
if (clone !== 'none') {
debug('create message event, using polifylled structured clone')
transfer?.length && console.warn(
'[@vitest/web-worker] `structuredClone` is not supported in this environment. '
+ 'Falling back to polyfill, your transferable options will be lost. '
+ 'Set `VITEST_WEB_WORKER_CLONE` environmental variable to "none", if you don\'t want to loose it,'
+ 'or update to Node 17+.',
)
return new MessageEvent('message', {
data: ponyfillStructuredClone(data, { lossy: true }),
origin,
})
}
debug('create message event without cloning an object')
return new MessageEvent('message', {
data,
origin,
})
}

function createMessageEvent(data: any, transferOrOptions: StructuredSerializeOptions | Transferable[] | undefined, clone: CloneOption) {
try {
return createClonedMessageEvent(data, transferOrOptions, clone)
}
catch (error) {
debug('failed to clone message, dispatch "messageerror" event: %o', error)
return new MessageEvent('messageerror', {
data: error,
})
globalThis.Worker = createWorkerConstructor(options)
}
}

export function defineWebWorker(options?: DefineWorkerOptions) {
if (typeof Worker !== 'undefined' && '__VITEST_WEB_WORKER__' in globalThis.Worker)
return

assertGlobalExists('EventTarget')
assertGlobalExists('MessageEvent')

const { config, rpc, mockMap, moduleCache } = getWorkerState()

const runnerOptions = {
fetchModule(id: string) {
return rpc.fetch(id)
},
resolveId(id: string, importer?: string) {
return rpc.resolveId(id, importer)
},
moduleCache,
mockMap,
interopDefault: config.deps.interopDefault ?? true,
root: config.root,
base: config.base,
}

const cloneType = () => (options?.clone ?? process.env.VITEST_WEB_WORKER_CLONE ?? 'native') as CloneOption

globalThis.Worker = class Worker extends EventTarget {
static __VITEST_WEB_WORKER__ = true

private _vw_workerTarget = new EventTarget()
private _vw_insideListeners = new Map<string, EventListenerOrEventListenerObject>()
private _vw_outsideListeners = new Map<string, EventListenerOrEventListenerObject>()
private _vw_name: string
private _vw_messageQueue: any[] | null = []

public onmessage: null | Procedure = null
public onmessageerror: null | Procedure = null
public onerror: null | Procedure = null

constructor(url: URL | string, options?: WorkerOptions) {
super()

// should be equal to DedicatedWorkerGlobalScope
const context: InlineWorkerContext = {
onmessage: null,
name: options?.name,
close: () => this.terminate(),
dispatchEvent: (event: Event) => {
return this._vw_workerTarget.dispatchEvent(event)
},
addEventListener: (...args) => {
if (args[1])
this._vw_insideListeners.set(args[0], args[1])
return this._vw_workerTarget.addEventListener(...args)
},
removeEventListener: this._vw_workerTarget.removeEventListener,
postMessage: (...args) => {
if (!args.length)
throw new SyntaxError('"postMessage" requires at least one argument.')

debug('posting message %o from the worker %s to the main thread', args[0], this._vw_name)
const event = createMessageEvent(args[0], args[1], cloneType())
this.dispatchEvent(event)
},
get self() {
return context
},
get global() {
return context
},
}

this._vw_workerTarget.addEventListener('message', (e) => {
context.onmessage?.(e)
})

this.addEventListener('message', (e) => {
this.onmessage?.(e)
})

this.addEventListener('messageerror', (e) => {
this.onmessageerror?.(e)
})

const runner = new InlineWorkerRunner(runnerOptions, context)

const id = (url instanceof URL ? url.toString() : url).replace(/^file:\/+/, '/')

const fsPath = toFilePath(id, config.root)

this._vw_name = options?.name ?? fsPath

debug('initialize worker %s', this._vw_name)

runner.executeFile(fsPath)
.then(() => {
// worker should be new every time, invalidate its sub dependency
moduleCache.invalidateSubDepTree([fsPath, `mock:${fsPath}`])
const q = this._vw_messageQueue
this._vw_messageQueue = null
if (q)
q.forEach(([data, transfer]) => this.postMessage(data, transfer), this)
debug('worker %s successfully initialized', this._vw_name)
}).catch((e) => {
debug('worker %s failed to initialize: %o', this._vw_name, e)
const EventConstructor = globalThis.ErrorEvent || globalThis.Event
const error = new EventConstructor('error', {
error: e,
message: e.message,
})
this.dispatchEvent(error)
this.onerror?.(e)
console.error(e)
})
}

addEventListener(type: string, callback: EventListenerOrEventListenerObject | null, options?: boolean | AddEventListenerOptions): void {
if (callback)
this._vw_outsideListeners.set(type, callback)
return super.addEventListener(type, callback, options)
}

postMessage(...args: [any, StructuredSerializeOptions | Transferable[] | undefined]): void {
if (!args.length)
throw new SyntaxError('"postMessage" requires at least one argument.')

const [data, transferOrOptions] = args
if (this._vw_messageQueue != null) {
debug('worker %s is not yet initialized, queue message %s', this._vw_name, data)
this._vw_messageQueue.push([data, transferOrOptions])
return
}

debug('posting message %o from the main thread to the worker %s', data, this._vw_name)

const event = createMessageEvent(data, transferOrOptions, cloneType())
if (event.type === 'messageerror')
this.dispatchEvent(event)
else
this._vw_workerTarget.dispatchEvent(event)
}
if (typeof SharedWorker === 'undefined' || !('__VITEST_WEB_WORKER__' in globalThis.SharedWorker)) {
assertGlobalExists('EventTarget')

terminate() {
debug('terminating worker %s', this._vw_name)
this._vw_outsideListeners.forEach((fn, type) => {
this.removeEventListener(type, fn)
})
this._vw_insideListeners.forEach((fn, type) => {
this._vw_workerTarget.removeEventListener(type, fn)
})
}
globalThis.SharedWorker = createSharedWorkerConstructor()
}
}
18 changes: 18 additions & 0 deletions packages/web-worker/src/runner.ts
@@ -0,0 +1,18 @@
import { VitestRunner } from 'vitest/node'

export class InlineWorkerRunner extends VitestRunner {
constructor(options: any, private context: any) {
super(options)
}

prepareContext(context: Record<string, any>) {
const ctx = super.prepareContext(context)
// not supported for now, we can't synchronously load modules
const importScripts = () => {
throw new Error('[vitest] `importScripts` is not supported in Vite workers. Please, consider using `import` instead.')
}
return Object.assign(ctx, this.context, {
importScripts,
})
}
}

0 comments on commit 8bb6b80

Please sign in to comment.