swarpc 0.9.0 → 0.11.0

package/src/server.ts

@@ -3,6 +3,7 @@
  * @mergeModuleWith <project>
  */
 
+/// <reference lib="webworker" />
 import { type } from "arktype"
 import { createLogger, type LogLevel } from "./log.js"
 import {
@@ -10,6 +11,7 @@ import {
   Payload,
   PayloadCore,
   PayloadHeaderSchema,
+  PayloadInitializeSchema,
   PayloadSchema,
   ProcedureImplementation,
   zImplementations,
@@ -17,6 +19,20 @@ import {
   type ProceduresMap,
 } from "./types.js"
 import { findTransferables } from "./utils.js"
+import { FauxLocalStorage } from "./localstorage.js"
+
+class MockedWorkerGlobalScope {
+  constructor() {}
+}
+
+const SharedWorkerGlobalScope =
+  globalThis.SharedWorkerGlobalScope ?? MockedWorkerGlobalScope
+
+const DedicatedWorkerGlobalScope =
+  globalThis.DedicatedWorkerGlobalScope ?? MockedWorkerGlobalScope
+
+const ServiceWorkerGlobalScope =
+  globalThis.ServiceWorkerGlobalScope ?? MockedWorkerGlobalScope
 
 /**
  * The sw&rpc server instance, which provides methods to register {@link ProcedureImplementation | procedure implementations},
@@ -25,7 +41,7 @@ import { findTransferables } from "./utils.js"
 export type SwarpcServer<Procedures extends ProceduresMap> = {
   [zProcedures]: Procedures
   [zImplementations]: ImplementationsMap<Procedures>
-  start(self: Window | Worker): void
+  start(): Promise<void>
 } & {
   [F in keyof Procedures]: (
     impl: ProcedureImplementation<
@@ -43,7 +59,9 @@ const abortedRequests = new Set<string>()
  * Creates a sw&rpc server instance.
  * @param procedures procedures the server will implement, see {@link ProceduresMap}
  * @param options various options
- * @param options.worker if provided, the server will use this worker to post messages, instead of sending it to all clients
+ * @param options.scope The worker scope to use, defaults to the `self` of the file where Server() is called.
+ * @param options.loglevel Maximum log level to use, defaults to "debug" (shows everything). "info" will not show debug messages, "warn" will only show warnings and errors, "error" will only show errors.
+ * @param options._scopeType @internal Don't touch, this is used in testing environments because the mock is subpar. Manually overrides worker scope type detection.
  * @returns a SwarpcServer instance. Each property of the procedures map will be a method, that accepts a function implementing the procedure (see {@link ProcedureImplementation}). There is also .start(), to be called after implementing all procedures.
  *
  * An example of defining a server:
@@ -51,17 +69,49 @@ const abortedRequests = new Set<string>()
  */
 export function Server<Procedures extends ProceduresMap>(
   procedures: Procedures,
-  { worker, loglevel = "debug" }: { worker?: Worker; loglevel?: LogLevel } = {}
+  {
+    loglevel = "debug",
+    scope,
+    _scopeType,
+  }: {
+    scope?: WorkerGlobalScope
+    loglevel?: LogLevel
+    _scopeType?: "dedicated" | "shared" | "service"
+  } = {}
 ): SwarpcServer<Procedures> {
   const l = createLogger("server", loglevel)
 
+  // If scope is not provided, use the global scope
+  // This function is meant to be used in a worker, so `self` is a WorkerGlobalScope
+  scope ??= self as WorkerGlobalScope
+
+  function scopeIsShared(
+    scope: WorkerGlobalScope
+  ): scope is SharedWorkerGlobalScope {
+    return scope instanceof SharedWorkerGlobalScope || _scopeType === "shared"
+  }
+
+  function scopeIsDedicated(
+    scope: WorkerGlobalScope
+  ): scope is DedicatedWorkerGlobalScope {
+    return (
+      scope instanceof DedicatedWorkerGlobalScope || _scopeType === "dedicated"
+    )
+  }
+
+  function scopeIsService(
+    scope: WorkerGlobalScope
+  ): scope is ServiceWorkerGlobalScope {
+    return scope instanceof ServiceWorkerGlobalScope || _scopeType === "service"
+  }
+
   // Initialize the instance.
   // Procedures and implementations are stored on properties with symbol keys,
   // to avoid any conflicts with procedure names, and also discourage direct access to them.
   const instance = {
     [zProcedures]: procedures,
     [zImplementations]: {} as ImplementationsMap<Procedures>,
-    start: (self: Window) => {},
+    start: async () => {},
   } as SwarpcServer<Procedures>
 
   // Set all implementation-setter methods
@@ -85,7 +135,16 @@ export function Server<Procedures extends ProceduresMap>(
     }) as SwarpcServer<Procedures>[typeof functionName]
   }
 
-  instance.start = (self: Window) => {
+  instance.start = async () => {
+    const port = await new Promise<MessagePort | undefined>((resolve) => {
+      if (!scopeIsShared(scope)) return resolve(undefined)
+      l.debug(null, "Awaiting shared worker connection...")
+      scope.addEventListener("connect", ({ ports: [port] }) => {
+        l.debug(null, "Shared worker connected with port", port)
+        resolve(port)
+      })
+    })
+
     // Used to post messages back to the client
     const postMessage = async (
       autotransfer: boolean,
@@ -93,17 +152,27 @@ export function Server<Procedures extends ProceduresMap>(
     ) => {
       const transfer = autotransfer ? [] : findTransferables(data)
 
-      if (worker) {
-        self.postMessage(data, { transfer })
-      } else {
-        await (self as any).clients.matchAll().then((clients: any[]) => {
+      if (port) {
+        port.postMessage(data, { transfer })
+      } else if (scopeIsDedicated(scope)) {
+        scope.postMessage(data, { transfer })
+      } else if (scopeIsService(scope)) {
+        await scope.clients.matchAll().then((clients) => {
           clients.forEach((client) => client.postMessage(data, { transfer }))
         })
       }
     }
 
-    // Listen for messages from the client
-    self.addEventListener("message", async (event: MessageEvent) => {
+    const listener = async (
+      event: MessageEvent<any> | ExtendableMessageEvent
+    ): Promise<void> => {
+      if (PayloadInitializeSchema.allows(event.data)) {
+        const { localStorageData } = event.data
+        l.debug(null, "Setting up faux localStorage", localStorageData)
+        new FauxLocalStorage(localStorageData).register(scope)
+        return
+      }
+
       // Decode the payload
       const { requestId, functionName } = PayloadHeaderSchema(
         type.enumerated(...Object.keys(procedures))
@@ -151,6 +220,9 @@ export function Server<Procedures extends ProceduresMap>(
         schemas.success
       ).assert(event.data)
 
+      if ("localStorageData" in payload)
+        throw "Unreachable: #initialize request payload should've been handled already"
+
       // Handle abortion requests (pro-choice ftw!!)
       if (payload.abort) {
         const controller = abortControllers.get(requestId)
@@ -206,7 +278,21 @@ export function Server<Procedures extends ProceduresMap>(
       } finally {
         abortedRequests.delete(requestId)
       }
-    })
+    }
+
+    // Listen for messages from the client
+    if (scopeIsShared(scope)) {
+      if (!port) throw new Error("SharedWorker port not initialized")
+      console.log("Listening for shared worker messages on port", port)
+      port.addEventListener("message", listener)
+      port.start()
+    } else if (scopeIsDedicated(scope)) {
+      scope.addEventListener("message", listener)
+    } else if (scopeIsService(scope)) {
+      scope.addEventListener("message", listener)
+    } else {
+      throw new Error(`Unsupported worker scope ${scope}`)
+    }
   }
 
   return instance

package/src/types.ts

@@ -4,7 +4,7 @@
  */
 
 import { type, type Type } from "arktype"
-import { Logger, RequestBoundLogger } from "./log.js"
+import { RequestBoundLogger } from "./log.js"
 
 /**
  * A procedure declaration
@@ -51,7 +51,7 @@ export type CancelablePromise<T = unknown> = {
    * Abort the request.
    * @param reason The reason for cancelling the request.
    */
-  cancel: (reason: string) => Promise<void>
+  cancel: (reason: string) => void
 }
 
 /**
@@ -87,7 +87,7 @@ export type ProcedureImplementation<
 
 /**
  * Declarations of procedures by name.
- * 
+ *
  * An example of declaring procedures:
  * {@includeCode ../example/src/lib/procedures.ts}
  */
@@ -131,6 +131,14 @@ export type Hooks<Procedures extends ProceduresMap> = {
   ) => void
 }
 
+export const PayloadInitializeSchema = type({
+  by: '"sw&rpc"',
+  functionName: '"#initialize"',
+  localStorageData: "Record<string, unknown>",
+})
+
+export type PayloadInitialize = typeof PayloadInitializeSchema.infer
+
 /**
  * @source
  */
@@ -184,11 +192,11 @@ export type PayloadCore<
  * @source
  */
 export const PayloadSchema = type
-  .scope({ PayloadCoreSchema, PayloadHeaderSchema })
+  .scope({ PayloadCoreSchema, PayloadHeaderSchema, PayloadInitializeSchema })
   .type("<Name extends string, I, P, S>", [
-    "PayloadHeaderSchema<Name>",
-    "&",
-    "PayloadCoreSchema<I, P, S>",
+    ["PayloadHeaderSchema<Name>", "&", "PayloadCoreSchema<I, P, S>"],
+    "|",
+    "PayloadInitializeSchema",
   ])
 
 /**
@@ -197,7 +205,7 @@ export const PayloadSchema = type
 export type Payload<
   PM extends ProceduresMap,
   Name extends keyof PM = keyof PM,
-> = PayloadHeader<PM, Name> & PayloadCore<PM, Name>
+> = (PayloadHeader<PM, Name> & PayloadCore<PM, Name>) | PayloadInitialize
 
 /**
  * A procedure's corresponding method on the client instance -- used to call the procedure. If you want to be able to cancel the request, you can use the `cancelable` method instead of running the procedure directly.