Expose isWorkerRuntime() function (#228)

Author: andywer

src/master/implementation.browser.ts

@@ -17,8 +17,7 @@ function createSourceBlobURL(code: string): string {
   return URL.createObjectURL(blob)
 }
 
-
-export function selectWorkerImplementation(): typeof WorkerImplementation {
+function selectWorkerImplementation(): typeof WorkerImplementation {
   if (typeof Worker === "undefined") {
     // Might happen on Safari, for instance
     // The idea is to only fail if the constructor is actually used
@@ -46,3 +45,17 @@ export function selectWorkerImplementation(): typeof WorkerImplementation {
     }
   }
 }
+
+let implementation: typeof WorkerImplementation
+
+export function getWorkerImplementation(): typeof WorkerImplementation {
+  if (!implementation) {
+    implementation = selectWorkerImplementation()
+  }
+  return implementation
+}
+
+export function isWorkerRuntime() {
+  const isWindowContext = typeof self !== "undefined" && typeof Window !== "undefined" && self instanceof Window
+  return typeof self !== "undefined" && self.postMessage && !isWindowContext ? true : false
+}

src/master/implementation.node.ts

@@ -1,3 +1,4 @@
+/// <reference lib="dom" />
 // tslint:disable function-constructor no-eval no-duplicate-super max-classes-per-file
 
 import getCallsites, { CallSite } from "callsites"
@@ -6,7 +7,14 @@ import { cpus } from 'os'
 import * as path from "path"
 import { ThreadsWorkerOptions, WorkerImplementation } from "../types/master"
 
+interface WorkerGlobalScope {
+  addEventListener(eventName: string, listener: (event: Event) => void): void
+  postMessage(message: any, transferables?: any[]): void
+  removeEventListener(eventName: string, listener: (event: Event) => void): void
+}
+
 declare const __non_webpack_require__: typeof require
+declare const self: WorkerGlobalScope
 
 type WorkerEventName = "error" | "message"
 
@@ -191,12 +199,36 @@ function initTinyWorker(): typeof WorkerImplementation {
   return Worker as any
 }
 
-export function selectWorkerImplementation(): typeof WorkerImplementation {
+let implementation: typeof WorkerImplementation
+let isTinyWorker: boolean
+
+function selectWorkerImplementation(): typeof WorkerImplementation {
   try {
+    isTinyWorker = false
     return initWorkerThreadsWorker()
   } catch(error) {
     // tslint:disable-next-line no-console
     console.debug("Node worker_threads not available. Trying to fall back to tiny-worker polyfill...")
+    isTinyWorker = true
     return initTinyWorker()
   }
 }
+
+export function getWorkerImplementation(): typeof WorkerImplementation {
+  if (!implementation) {
+    implementation = selectWorkerImplementation()
+  }
+  return implementation
+}
+
+export function isWorkerRuntime() {
+  if (isTinyWorker) {
+    return typeof self !== "undefined" && self.postMessage ? true : false
+  } else {
+    // Webpack hack
+    const isMainThread = typeof __non_webpack_require__ === "function"
+      ? __non_webpack_require__("worker_threads").isMainThread
+      : eval("require")("worker_threads").isMainThread
+    return !isMainThread
+  }
+}

src/master/implementation.ts

@@ -12,5 +12,10 @@ import * as NodeImplementation from "./implementation.node"
 const runningInNode = typeof process !== 'undefined' && process.arch !== 'browser' && 'pid' in process
 const implementation = runningInNode ? NodeImplementation : BrowserImplementation
 
+/** Default size of pools. Depending on the platform the value might vary from device to device. */
 export const defaultPoolSize = implementation.defaultPoolSize
-export const selectWorkerImplementation = implementation.selectWorkerImplementation
+
+export const getWorkerImplementation = implementation.getWorkerImplementation
+
+/** Returns `true` if this code is currently running in a worker. */
+export const isWorkerRuntime = implementation.isWorkerRuntime

src/master/index.ts

@@ -1,12 +1,13 @@
 import { Worker as WorkerType } from "../types/master"
-import { selectWorkerImplementation } from "./implementation"
+import { getWorkerImplementation, isWorkerRuntime } from "./implementation"
 
 export { FunctionThread, ModuleThread } from "../types/master"
 export { Pool } from "./pool"
 export { spawn } from "./spawn"
 export { Thread } from "./thread"
+export { isWorkerRuntime }
 
 export type Worker = WorkerType
 
 /** Worker implementation. Either web worker or a node.js Worker class. */
-export const Worker = selectWorkerImplementation()
+export const Worker = getWorkerImplementation()

src/worker/index.ts

@@ -19,6 +19,9 @@ import Implementation from "./implementation"
 export { registerSerializer } from "../common"
 export { Transfer } from "../transferable"
 
+/** Returns `true` if this code is currently running in a worker. */
+export const isWorkerRuntime = Implementation.isWorkerRuntime
+
 let exposeCalled = false
 
 const isMasterJobRunMessage = (thing: any): thing is MasterJobRunMessage => thing && thing.type === MasterMessageType.run

test/webpack/addition-worker.ts

@@ -1,4 +1,8 @@
-import { expose } from "../../src/worker"
+import { expose, isWorkerRuntime } from "../../src/worker"
+
+if (!isWorkerRuntime()) {
+  throw Error("isWorkerRuntime() says we are not in a worker.")
+}
 
 expose(function add(a: number, b: number) {
   return a + b

test/webpack/app.ts

@@ -1,4 +1,4 @@
-import { spawn, Pool, Worker } from "../../src/index"
+import { isWorkerRuntime, spawn, Pool, Worker } from "../../src/index"
 
 type AdditionWorker = (a: number, b: number) => number
 type HelloWorker = (text: string) => string
@@ -46,8 +46,15 @@ async function test3() {
   }
 }
 
+function test4() {
+  if (isWorkerRuntime() !== false) {
+    throw Error("Expected isWorkerRuntime() to return false. Got: " + isWorkerRuntime())
+  }
+}
+
 export default () => Promise.all([
   test(),
   test2(),
-  test3()
+  test3(),
+  test4()
 ])