js-rpc2 1.1.1 → 1.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "js-rpc2",
-  "version": "1.1.1",
+  "version": "1.2.0",
   "description": "js web websocket http rpc",
   "main": "index.js",
   "type": "module",

package/src/lib.js

@@ -1,3 +1,7 @@
+/**
+ * @import { ExtensionApi } from "./types.js"
+ */
+
 const JS_RPC_WITH_CRYPTO = true
 
 export const sleep = (/** @type {number} */ timeout) => new Promise((resolve) => setTimeout(resolve, timeout))
@@ -535,12 +539,14 @@ export function parseRpcItemData(array) {
 }
 
 /**
- * @param {object} extension
+ * @template T
+ * @param {ExtensionApi<T>} extension
  * @param {WritableStreamDefaultWriter<Uint8Array>} writer
  * @param {Uint8Array} buffer
- * @param {(msg:string)=>void} [logger]
+ * @param {(msg:string)=>void} logger
+ * @param {any} context
  */
-export async function rpcRunServerDecodeBuffer(extension, writer, buffer, logger) {
+export async function rpcRunServerDecodeBuffer(extension, writer, buffer, logger, context) {
     /** @type{RPC_DATA} */
     let box = null
     let dataId = 0
@@ -570,7 +576,17 @@ export async function rpcRunServerDecodeBuffer(extension, writer, buffer, logger
                 params.push(p.data)
             }
         }
-        let ret = await extension[fnName](...params)
+        let store = extension.asyncLocalStorage
+        let ret = await new Promise((resolve, reject) => {
+            store.run(context, async () => {
+                try {
+                    let ret = await extension[fnName](...params)
+                    resolve(ret)
+                } catch (error) {
+                    reject(error)
+                }
+            })
+        })
         if (Array.isArray(ret)) {
             box = { id: o.id, type: RPC_TYPE_RETURN_ARRAY, data: buildRpcItemData(ret), }
         } else {
@@ -632,11 +648,13 @@ export function createRPCProxy(apiInvoke) {
  */
 
 /**
+ * @template T
  * @param {{ 
  * rpcKey: string; 
- * extension: object; 
+ * extension: ExtensionApi<T>;
  * logger?: (msg:string)=>void; 
  * async?: boolean; 
+ * context?:any;
  * }} param
  */
 export function createRpcServerHelper(param) {
@@ -644,12 +662,13 @@ export function createRpcServerHelper(param) {
     const encode = createEncodeStream(rpc_key_iv)
     const decode = createDecodeStream(rpc_key_iv)
     let writer = encode.writable.getWriter()
+    let context = param.context
     decode.readable.pipeTo(new WritableStream({
         async write(buffer) {
             if (param.async) {
-                rpcRunServerDecodeBuffer(param.extension, writer, buffer, param.logger).catch(console.error)
+                rpcRunServerDecodeBuffer(param.extension, writer, buffer, param.logger, context).catch(console.error)
             } else {
-                await rpcRunServerDecodeBuffer(param.extension, writer, buffer, param.logger)
+                await rpcRunServerDecodeBuffer(param.extension, writer, buffer, param.logger, context)
             }
         },
         async close() {

package/src/lib.test.js

@@ -6,6 +6,9 @@ import { WebSocketServer } from 'ws'
 import Koa from 'koa'
 import Router from 'koa-router'
 import { createRpcServerKoaRouter, createRpcServerWebSocket } from './server.js'
+import { AsyncLocalStorage } from 'node:async_hooks'
+import { Readable, Transform } from 'node:stream'
+import { pipeline } from 'node:stream/promises'
 
 test('basic', async () => {
     // node --test-name-pattern="^basic$" src/lib.test.js
@@ -18,6 +21,8 @@ test('basic', async () => {
     let server = createServer(app.callback()).listen(9000)
 
     class RpcApi {
+        asyncLocalStorage = new AsyncLocalStorage()
+
         /**
          * 
          * @param {string} string 
@@ -62,7 +67,7 @@ test('basic', async () => {
         url: `http://127.0.0.1:9000/rpc/abc`, // others
     })
 
-    let ret = await rpc.hello('123', new Uint8Array(3), { q: 2, w: 3, e: 4 }, null, undefined, [1, 2, 3, 4], (message, num) => {
+    let ret = await rpc.hello('123', new Uint8Array(3), { q: 2, w: 3, e: 4 }, null, undefined, [1, 2, 3, 4], async (message, num) => {
         console.info('callback', message, num)
     })
     console.info(ret)
@@ -74,6 +79,7 @@ test('测试RPC调用-WebSocket', async () => {
     // node --test-name-pattern="^测试RPC调用-WebSocket$" src/lib.test.js
 
     const extension = {
+        asyncLocalStorage: new AsyncLocalStorage(),
         hello: async function (/** @type {string} */ name) {
             return `hello ${name}`
         },
@@ -168,6 +174,7 @@ test('测试RPC调用-WebSocket', async () => {
 test('测试RPC调用-KoaRouter', async () => {
     // node --test-name-pattern="^测试RPC调用-KoaRouter$" src/lib.test.js
     const extension = {
+        asyncLocalStorage: new AsyncLocalStorage(),
         hello: async function (/** @type {string} */ name) {
             return `hello ${name}`
         },
@@ -276,6 +283,54 @@ test('测试RPC调用-KoaRouter', async () => {
     console.info('over!')
 })
 
+test('测试RPC调用-KoaRouter-AsyncLocalStorage', async () => {
+    // node --test-name-pattern="^测试RPC调用-KoaRouter-AsyncLocalStorage$" src/lib.test.js
+    const extension = {
+        asyncLocalStorage: new AsyncLocalStorage(),
+        hello: async function (/** @type {string} */ name) {
+            console.info('asyncLocalStorage.getStore', this.asyncLocalStorage.getStore())
+            return `hello ${name}`
+        },
+    }
+
+    await runWithAbortController(async (ac) => {
+        let server = createServer()
+        let app = new Koa()
+        let router = new Router()
+
+        ac.signal.addEventListener('abort', () => { server.close() })
+        createRpcServerKoaRouter({
+            path: '/abc',
+            router: router,
+            rpcKey: 'abc',
+            extension: extension,
+            logger(msg) {
+                console.info(msg)
+            },
+        })
+
+        server.addListener('request', app.callback())
+        app.use(router.routes())
+        app.use(router.allowedMethods())
+
+        server.listen(9000)
+        await sleep(100)
+
+        /** @type{typeof extension} */
+        let client = createRpcClientHttp({
+            url: `http://127.0.0.1:9000/abc`,
+            rpcKey: 'abc',
+            signal: ac.signal,
+        })
+        await sleep(100)
+
+        let string = await client.hello('asdfghjkl')
+        console.info(string)
+        strictEqual(string, 'hello asdfghjkl')
+    })
+    console.info('over!')
+})
+
 /**
  * @param {(ac: AbortController) => Promise<void>} func
  */
@@ -299,6 +354,7 @@ test('error-stack', async () => {
     let server = createServer(app.callback()).listen(9000)
 
     class RpcApi {
+        asyncLocalStorage = new AsyncLocalStorage()
         async hello() {
             new URL('/')
         }
@@ -329,4 +385,59 @@ test('error-stack', async () => {
     }
     server.close()
 
+})
+
+test('async-local-storage', async () => {
+    // node --test-name-pattern="^async-local-storage$" src/lib.test.js
+    let store = new AsyncLocalStorage()
+
+    let p1 = Promise.withResolvers()
+    store.run({ a: 1 }, async () => {
+        console.info('--- 1 1 ', store.getStore())
+        await sleep(100)
+        console.info('--- 1 2 ', store.getStore())
+        p1.resolve()
+    })
+    let p2 = Promise.withResolvers()
+    store.run({ a: 2 }, async () => {
+        console.info('--- 2 1 ', store.getStore())
+        await sleep(100)
+        console.info('--- 2 2 ', store.getStore())
+        p2.resolve()
+    })
+    await Promise.all([p1.promise, p2.promise])
+})
+
+test('async-local-storage-stream', async () => {
+    // node --test-name-pattern="^async-local-storage-stream$" src/lib.test.js
+    let store = new AsyncLocalStorage()
+    let p1 = Promise.withResolvers()
+    let readable = new Readable({ read() { } })
+
+    store.run({ a: 1 }, async () => {
+        console.info('--- step 1 ', store.getStore())
+        !(async () => {
+            for (let i = 0; i < 3; i++) {
+                readable.push(`inner data:${i}`)
+                await sleep(10)
+            }
+        })()
+        await pipeline(readable, new Transform({
+            transform(chunk, _, callback) {
+                console.info(`--- step transform ${chunk} :`, store.getStore())
+                callback()
+            }
+        }))
+        console.info('--- step 2 ', store.getStore())
+        p1.resolve()
+    })
+
+    for (let i = 0; i < 3; i++) {
+        readable.push(`data:${i}`)
+        await sleep(10)
+    }
+    readable.push(null)
+    await sleep(100)
+
+    await Promise.all([p1.promise])
 })

package/src/server.js

@@ -1,8 +1,10 @@
 import { Readable } from "node:stream"
 import { createRpcServerHelper } from "./lib.js"
+import { AsyncLocalStorage } from "node:async_hooks"
 
 /**
  * @import {WebSocketServer} from 'ws'
+ * @import {ExtensionApi} from './types.js'
  */
 
 export { createRpcServerHelper }
@@ -12,15 +14,19 @@ export { createRpcServerHelper }
  */
 
 /**
+ * @template T
  * @param {{
  * path: string; 
  * wss: WebSocketServer; 
  * rpcKey:string;
- * extension: Object; 
+ * extension: ExtensionApi<T>; 
  * logger?:(msg:string)=>void;
  * }} param
  */
 export function createRpcServerWebSocket(param) {
+    if (!param.extension.asyncLocalStorage) {
+        param.extension.asyncLocalStorage = new AsyncLocalStorage()
+    }
     param.wss.on('connection', (ws, request) => {
         let url = request.url
         if (url != param.path) {
@@ -54,18 +60,21 @@ export function createRpcServerWebSocket(param) {
 }
 
 /**
- * 
+ * @template T
  * @param {{
  * path: string; 
  * router: Router<any, {}>; 
  * rpcKey?:string;
  * logger?:(msg:string)=>void;
- * extension: Object; 
+ * extension: ExtensionApi<T>; 
  * }} param 
  */
 export function createRpcServerKoaRouter(param) {
+    if (!param.extension.asyncLocalStorage) {
+        param.extension.asyncLocalStorage = new AsyncLocalStorage()
+    }
     param.router.post(param.path, async (ctx) => {
-        let helper = createRpcServerHelper({ rpcKey: param.rpcKey, extension: param.extension, logger: param.logger })
+        let helper = createRpcServerHelper({ rpcKey: param.rpcKey, extension: param.extension, logger: param.logger, context: ctx })
         /** @type{ReadableStream} */
         let a = Readable.toWeb(ctx.req)
         await a.pipeTo(helper.writable)

package/src/types.ts

@@ -0,0 +1,3 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+
+export type ExtensionApi<T> = { asyncLocalStorage: AsyncLocalStorage<T> } & object;