Skip to content

Commit

Permalink
feat: implement Runner proxy
Browse files Browse the repository at this point in the history
  • Loading branch information
Dimava committed Oct 29, 2024
1 parent e8f6143 commit 30333bd
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 0 deletions.
23 changes: 23 additions & 0 deletions src/common.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,29 @@ export interface TinypoolWorker {
threadId: number
}

type SliceFirst<T extends any[]> = //
T extends [] | [any?] | [any]
? T
: T extends [any, any, ...any[]]
? never
: Required<T> extends [infer A, ...any[]]
? [A?]
: never

export type TinypoolRunner<T> = {
[K in keyof T as T[K] extends (...args: any[]) => any
? K
: never]: T[K] extends (...args: infer A) => infer R
? [R] extends [Promise<any>]
? true extends R & 'any'
? (...args: SliceFirst<A>) => Promise<any>
: A['length'] extends 0 | 1
? T[K]
: (...args: SliceFirst<A>) => Promise<R>
: (...args: SliceFirst<A>) => Promise<R>
: never
}

/**
* Tinypool's internal messaging between main thread and workers.
* - Utilizers can use `__tinypool_worker_message__` property to identify
Expand Down
17 changes: 17 additions & 0 deletions src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import {
type TinypoolData,
type TinypoolWorker,
type TinypoolChannel,
type TinypoolRunner,
} from './common'
import ThreadWorker from './runtime/thread-worker'
import ProcessWorker from './runtime/process-worker'
Expand Down Expand Up @@ -1159,6 +1160,22 @@ class Tinypool extends EventEmitterAsyncResource {
})
}

withRunner<Module>(): this & { runner: TinypoolRunner<Module> } {
return this as any
}

get runner(): TinypoolRunner<{ [K in never]: never }> {
return new Proxy({} as TinypoolRunner<{ [K in never]: never }>, {
get:
(_target: any, p: string) =>
(...a: any[]) => {
if (a.length > 1)
throw new Error('TinypoolRunner doesn’t support args array')
return this.run(a[0], { name: p })
},
})
}

async destroy() {
await this.#pool.destroy()
this.emitDestroy()
Expand Down
21 changes: 21 additions & 0 deletions test/fixtures/multiple.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
export function a(): 'a'

export function b(): 'b'

export default a

export function identity<V>(v: V): V
export function identityAsync<V>(v: V): Promise<V>

export function foobar<V>(o: { foobar: V }): V
export function foobarAsync<V>(o: { foobar: V }): Promise<V>

export function args<A>(...args: A[]): A
export function argsAsync<A extends any[]>(...args: A): Promise<A>

export function firstArg(a: 1, b?: 2): 1
export function firstArgAsync(a: 2, b?: 3): Promise<2>

export const digit: 4

export function returnsAny(): any
35 changes: 35 additions & 0 deletions test/fixtures/multiple.js
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,38 @@ export function b() {
}

export default a

export function identity(v) {
return v
}
export function identityAsync(v) {
return v
}

export function foobar({ foobar }) {
return foobar
}
export function foobarAsync({ foobar }) {
return foobar
}

export function args(...args) {
return args
}
// eslint-disable-next-line @typescript-eslint/require-await
export async function argsAsync(...args) {
return args
}
export function firstArg(a) {
return a
}
// eslint-disable-next-line @typescript-eslint/require-await
export async function firstArgAsync(a) {
return a
}

export const digit = 4

export function returnsAny() {
return 'any'
}
61 changes: 61 additions & 0 deletions test/runner.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import { dirname, resolve } from 'node:path'
import Tinypool from 'tinypool'
import { fileURLToPath } from 'node:url'

const __dirname = dirname(fileURLToPath(import.meta.url))

test('runner worker_threads test', async () => {
const { runner } = new Tinypool({
filename: resolve(__dirname, 'fixtures/multiple.js'),
runtime: 'worker_threads',
}).withRunner<typeof import('./fixtures/multiple.js')>()

expect((await runner.a()) satisfies 'a').toBe('a')
expect((await runner.b()) satisfies 'b').toBe('b')

expect(await runner.identity(123)).toBe(123)
expect((await runner.identityAsync(123)) satisfies 123).toBe(123)

expect(await runner.foobar({ foobar: 1 })).toBe(1)
expect((await runner.foobarAsync({ foobar: 1 })) satisfies 1).toBe(1)

expect((await runner.firstArg(1)) satisfies 1).toBe(1)
expect((await runner.firstArgAsync(2)) satisfies 2).toBe(2)

// @ts-expect-error should throw
expect(() => runner.firstArg(1, 2)).toThrow('doesn’t support args array')
// @ts-expect-error should throw
expect(() => runner.firstArgAsync(2, 3)).toThrow('doesn’t support args array')
// @ts-expect-error should throw
expect(() => runner.args(1, 2, 3)).toThrow('doesn’t support args array')
// @ts-expect-error should throw
expect(() => runner.argsAsync(1, 2, 3)).toThrow('doesn’t support args array')
})

test('runner child_process test', async () => {
const { runner } = new Tinypool({
filename: resolve(__dirname, 'fixtures/multiple.js'),
runtime: 'child_process',
}).withRunner<typeof import('./fixtures/multiple.js')>()

expect((await runner.a()) satisfies 'a').toBe('a')
expect((await runner.b()) satisfies 'b').toBe('b')

expect(await runner.identity(123)).toBe(123)
expect((await runner.identityAsync(123)) satisfies 123).toBe(123)

expect(await runner.foobar({ foobar: 1 })).toBe(1)
expect((await runner.foobarAsync({ foobar: 1 })) satisfies 1).toBe(1)

expect((await runner.firstArg(1)) satisfies 1).toBe(1)
expect((await runner.firstArgAsync(2)) satisfies 2).toBe(2)

// @ts-expect-error should throw
expect(() => runner.firstArg(1, 2)).toThrow('doesn’t support args array')
// @ts-expect-error should throw
expect(() => runner.firstArgAsync(2, 3)).toThrow('doesn’t support args array')
// @ts-expect-error should throw
expect(() => runner.args(1, 2, 3)).toThrow('doesn’t support args array')
// @ts-expect-error should throw
expect(() => runner.argsAsync(1, 2, 3)).toThrow('doesn’t support args array')
})

0 comments on commit 30333bd

Please sign in to comment.