diff --git a/src/IMQRPCOptions.ts b/src/IMQRPCOptions.ts index c6d7ad2..47999f4 100644 --- a/src/IMQRPCOptions.ts +++ b/src/IMQRPCOptions.ts @@ -49,6 +49,28 @@ export interface IMQAfterCall<_T> { (req: IMQRPCRequest, res?: IMQRPCResponse): Promise; } +/** + * Around hook wrapping the actual service method invocation. It receives the + * request/response and a `next` callback that runs the method and resolves to + * its return value; it MUST call `next()` (returning its resolved value) to + * produce the response data. Unlike `beforeCall`/`afterCall`, this lets a hook + * run the method inside its own scope — e.g. establishing an OpenTelemetry + * context so any spans the method (and its downstream calls) create nest under + * the request span. When unset, the method is invoked directly. + * + * @param {IMQRPCRequest} req - the incoming request + * @param {IMQRPCResponse} res - the response being prepared + * @param {() => Promise} next - runs the method, resolves to its result + * @return {Promise} - the value to use as the response data + */ +export interface IMQWrapCall<_T> { + ( + req: IMQRPCRequest, + res: IMQRPCResponse, + next: () => Promise, + ): Promise; +} + /** * Options for the built-in metrics server. */ @@ -67,6 +89,7 @@ export interface IMQServiceOptions extends IMQOptions { metricsServer?: IMQMetricsServerOptions; beforeCall?: IMQBeforeCall; afterCall?: IMQAfterCall; + wrapCall?: IMQWrapCall; } /** diff --git a/src/IMQService.ts b/src/IMQService.ts index 616a0b7..7a07b97 100644 --- a/src/IMQService.ts +++ b/src/IMQService.ts @@ -48,7 +48,11 @@ import { SIGNALS } from './helpers/index.js'; import { cpus } from 'node:os'; import cluster, { type Worker } from 'node:cluster'; import { type ArgDescription } from './IMQRPCDescription.js'; -import { type IMQBeforeCall, type IMQAfterCall } from './IMQRPCOptions.js'; +import { + type IMQBeforeCall, + type IMQAfterCall, + type IMQWrapCall, +} from './IMQRPCOptions.js'; import { runWithRequest } from './IMQRequestContext.js'; import { createServer, type Server } from 'node:http'; @@ -270,11 +274,24 @@ export abstract class IMQService { } try { - response.data = this[method].apply(this, args); + // Run the method through the optional `wrapCall` around-hook so a + // hook can execute it within its own scope (e.g. an OpenTelemetry + // context). Without a hook the method is invoked directly, exactly + // as before. + const invoke = async (): Promise => { + const result = this[method].apply(this, args); + + return result && result.then ? await result : result; + }; - if (response.data && response.data.then) { - response.data = await response.data; - } + const wrapCall = this.options.wrapCall as + | IMQWrapCall + | undefined; + + response.data = + typeof wrapCall === 'function' + ? await wrapCall.call(this, request, response, invoke) + : await invoke(); } catch (err: any) { response.error = IMQError( err.code || 'IMQ_RPC_CALL_ERROR',