rxjs 实现一个 webworker pool

202 阅读1分钟
import {
    ReplaySubject,
    fromEvent,
    tap,
    Observable,
    merge,
    zipWith,
    map,
    takeUntil,
    EMPTY,
} from "rxjs";
import T from './worker.ts?worker'

export class WorkerPool<Task = any, Message = any> {
    private subject = new ReplaySubject<Task>(0, 0)
    private observables: Observable<Message>[] = []
    private postMessage$: Observable<[Task, Worker]> = EMPTY
    private subscribe$: Observable<Message> = EMPTY
    private readonly free$
    private readonly workerCtor
    private readonly cancel$ = new ReplaySubject(0)

    constructor(ctor: new (...args: any[]) => Worker, concurrent: number = window.navigator.hardwareConcurrency) {
       this.workerCtor = ctor
       this.free$ = new ReplaySubject<Worker>(concurrent)
       this.init(concurrent)
    }

    private init(concurrent: number) {
       for ( let i = 0; i < concurrent; i++ ) {
          const worker = new this.workerCtor()
          this.observables.push(
             // 这块可以在下载postmessage中转换为高阶Observable 
             // 但是每次订阅大概会浪费不到10ms的时间
             // 所以提前订阅能节约一定时间
             fromEvent<MessageEvent>(worker, 'message').pipe(
                takeUntil(this.cancel$),
                map((e: MessageEvent) => e.data),
                tap(() => {
                   this.free$.next( worker )
                }),
             )
          )
          this.free$.next(worker)
       }
       this.postMessage$ = this.subject.pipe(
          takeUntil(this.cancel$),
          zipWith(this.free$),
          tap(([task, wip]) => wip.postMessage(task)),
       )
       this.subscribe$ = merge(...this.observables).pipe(takeUntil(this.cancel$))
    }

    emit(payload: Task) {
       this.subject.next(payload)
    }

    subscribe(callback?: (message: Message) => void) {
       this.postMessage$.subscribe()
       this.subscribe$.subscribe(callback)
    }

    unsubscribe() {
       this.free$.pipe(tap(worker => worker.terminate()), takeUntil(this.cancel$)).subscribe()
       this.cancel$.next(null)
       this.cancel$.unsubscribe()
       this.free$.unsubscribe()
    }
}
// 测试case 
// 没有订阅之前 发送的消息无效
const pool = new WorkerPool(T)
pool.emit({
    action: "oddFilter",
    payload: {
       timestamp: Date.now(),
       data: new Array(10).fill(0).map((_, idx) => idx),
    },
});
pool.emit({
    action: "oddFilter",
    payload: {
       timestamp: Date.now(),
       data: new Array(100).fill(0).map((_, idx) => idx),
    },
});
pool.emit({
    action: "oddFilter",
    payload: {
       timestamp: Date.now(),
       data: new Array(1000).fill(0).map((_, idx) => idx),
    },
});
setTimeout(() => {
    pool.subscribe(console.log)
}, 3000)