feat: basic mdc support, fix correlationId reset

Author: mitjacotic

src/integration.ts

@@ -3,6 +3,7 @@ import {
   IrisChannels,
 } from './lib/asyncapi/schema/channels'
 import { type CustomErrorI, registerRejectableErrors } from './lib/errors'
+import { MdcI, registerMdcProvider } from './lib/mdc'
 import Logger, { type LoggerI } from './logger'
 
 export interface IrisIntegrationI {
@@ -11,6 +12,19 @@ export interface IrisIntegrationI {
   asyncapi?: {
     customChannelClasses?: CustomChannelClassesI
   }
+
+  /**
+   * Since MDC is not integrated in same way as with Java IRIS
+   * counterpart it's up to integration whether to use it or not.
+   * In case it's used, this method can be used to obtain MDC instance.
+   *
+   * Current usage within node IRIS is to:
+   * - use reqiestId value as correlationId (for HTTP requests that
+   *   produce events)
+   * - be able to omit userId agument with the `publish.publishToUser()`
+   *   method
+   */
+  mdcProvider?: () => MdcI | undefined | Promise<MdcI | undefined>
 }
 
 /**
@@ -26,4 +40,7 @@ export function update(config: IrisIntegrationI): void {
   if (config.asyncapi?.customChannelClasses !== undefined) {
     IrisChannels.setCustomChannelClasses(config.asyncapi.customChannelClasses)
   }
+  if (config.mdcProvider !== undefined) {
+    registerMdcProvider(config.mdcProvider)
+  }
 }

src/lib/mdc.ts

@@ -11,13 +11,24 @@ import { MDC_CLASS, SetMetadata } from './storage'
 
 const { MESSAGE: MH } = MESSAGE_HEADERS
 
+let mdcProvider: undefined | (() => MDC | undefined | Promise<MDC | undefined>)
+
 export type MdcI = {
   sessionId?: string
   userId?: string
   clientTraceId?: string
   correlationId?: string
   eventType?: string
   clientVersion?: string
+
+  /**
+   * HTTP related
+   * if MDC provider funciton is provided then
+   * when requestId is available within the MDC
+   * it is used as a correlationId instaed of
+   * generating a new one.
+   */
+  requestId?: string
 }
 
 export class MDC implements MdcI {
@@ -27,6 +38,7 @@ export class MDC implements MdcI {
   correlationId?: string
   eventType?: string
   clientVersion?: string
+  requestId?: string
 }
 
 SetMetadata(MDC_CLASS, true)(MDC)
@@ -60,6 +72,16 @@ export const isMDCClass = (target: unknown): boolean => {
   return false
 }
 
+export async function getMdc(): Promise<MdcI | undefined> {
+  if (mdcProvider !== undefined) {
+    return await mdcProvider()
+  }
+}
+
+export function registerMdcProvider(provider: typeof mdcProvider) {
+  mdcProvider = provider
+}
+
 function setFromHeader(
   amqpMessage: Pick<AmqpMessage, 'properties'>,
   ref: MdcI,

src/lib/message_handler.param.decorator.ts

@@ -12,7 +12,7 @@ export class AmqpMessage implements amqplib.ConsumeMessage {
 
 SetMetadata(AMQP_MESSAGE_CLASS, true)(AmqpMessage)
 
-export const isAmqpMessageClass = (target: unknown): boolean => {
+export const isAmqpMessageClass = (target: unknown): target is AmqpMessage => {
   if (typeof target === 'object' && target !== null) {
     if (target instanceof AmqpMessage) {
       return true

src/lib/publish.ts

@@ -9,6 +9,7 @@ import * as constants from './constants'
 import { asError } from './errors'
 import flags from './flags'
 import * as helper from './helper'
+import { amqpToMDC, getMdc } from './mdc'
 import * as message from './message'
 import { getPublishExchangeProps } from './message.process'
 import {
@@ -42,7 +43,7 @@ export function getUserPublisher<T>(messageClass: ClassConstructor<T>) {
 
   return async (
     msg: T,
-    user: string | AmqpMessage,
+    user?: string | AmqpMessage,
     pubOpts?: publishI.PublishOptionsI,
   ): Promise<boolean> => publishToUser(messageClass, msg, user, pubOpts)
 }
@@ -57,21 +58,21 @@ export const publish = async <T>(
  * Send the mesaage to `user` exchange,
  * ignoring the exchange set on event being published.
  *
- * @{param user} - user id or AmqpMessage
+ * @param {string} [user] - user id or AmqpMessage
  *  when processing a message from user, AmqpMessage can
  *  be used to get user id from consumed message.
+ *
+ *  When omitted it is obtained from the MDC via mdcProvider.
+ *
+ *
  */
 export const publishToUser = async <T>(
   messageClass: ClassConstructor<T>,
   msg: T,
-  user: string | AmqpMessage,
+  user?: string | AmqpMessage,
   pubOpts?: publishI.PublishOptionsI,
 ): Promise<boolean> => {
-  const userIdString = typeof user === 'string'
-  const hasOriginalMsg = !userIdString && isAmqpMessageClass(user)
-  const userId = userIdString
-    ? user
-    : _.get(user, `properties.headers[${MESSAGE_HEADERS.MESSAGE.USER_ID}]`)
+  const userId = await obtainUserId(user || pubOpts?.amqpPublishOpts)
 
   if (!_.isString(userId) || _.isEmpty(userId)) {
     throw new Error('ERR_IRIS_PUBLISHER_USER_ID_NOT_RESOLVED')
@@ -81,7 +82,7 @@ export const publishToUser = async <T>(
     messageClass,
     msg,
     Object.assign({}, pubOpts, { userId }),
-    hasOriginalMsg ? <AmqpMessage>user : undefined,
+    isAmqpMessageClass(user) ? <AmqpMessage>user : undefined,
     message.Scope.USER,
   )
 }
@@ -122,7 +123,12 @@ async function internalPublish<T>(
     msgMeta,
     msgString,
     routingKey,
-    getAmqpBasicProperties(exchangeName, msgMeta, originalMessage, pubOpts),
+    await getAmqpBasicProperties(
+      exchangeName,
+      msgMeta,
+      originalMessage,
+      pubOpts,
+    ),
     overrideScope,
   )
 }
@@ -133,6 +139,7 @@ export async function doPublish(
   routingKeyArg: string,
   options?: amqplib.Options.Publish,
   overrideScope?: message.Scope,
+  originalMessage?: Pick<amqplib.Message, 'properties'>,
 ): Promise<boolean> {
   validateBeforePublish(msgMeta, options, overrideScope)
 
@@ -146,12 +153,25 @@ export async function doPublish(
 
   const routingKey = publishingExchangeRoutingKey ?? routingKeyArg
 
-  logger.debug(TAG, `Publishing message to "${publishingExchangeName}"`, {
+  const logInfo: any = {
     evt: msg,
     routingKey,
     publishingExchangeName,
     options: amqpHelper.safeAmqpObjectForLogging(options),
-  })
+  }
+
+  const mdc =
+    originalMessage !== undefined ? amqpToMDC(originalMessage) : await getMdc()
+
+  if (mdc !== undefined) {
+    logInfo.mdc = mdc
+  }
+
+  logger.debug(
+    TAG,
+    `Publishing message to "${publishingExchangeName}"`,
+    logInfo,
+  )
 
   const channel = await connection.assureDefaultChannel()
 
@@ -190,13 +210,13 @@ async function msg2String<T>(
   return JSON.stringify(msg)
 }
 
-function getAmqpBasicProperties(
+async function getAmqpBasicProperties(
   exchangeName: string,
   msgMeta: message.ProcessedMessageMetadataI,
   originalMsg?: Pick<amqplib.Message, 'properties'>,
   pubOpts?: publishI.PublishOptionsI,
-): Partial<amqplib.MessageProperties> {
-  const amqpProperties = getAmqpPropsWithoutHeaders(originalMsg, pubOpts)
+): Promise<Partial<amqplib.MessageProperties>> {
+  const amqpProperties = await getAmqpPropsWithoutHeaders(originalMsg, pubOpts)
   const amqpHeaders = getAmqpHeaders(exchangeName, originalMsg, pubOpts)
 
   // TODO: subscription updates, set cache ttl, subcription id etc.
@@ -208,11 +228,7 @@ function getAmqpBasicProperties(
 
   if (pubOpts?.userId !== undefined) {
     const serviceId = helper.getServiceName()
-    const correlationId = randomUUID()
 
-    // when overriding user header make sure
-    // to clean possible existing event context properties
-    amqpProperties.correlationId = correlationId
     amqpHeaders[MESSAGE_HEADERS.MESSAGE.ORIGIN_SERVICE_ID] = serviceId
     amqpHeaders[MESSAGE_HEADERS.MESSAGE.USER_ID] = pubOpts.userId
     delete amqpHeaders[MESSAGE_HEADERS.MESSAGE.ROUTER]
@@ -225,12 +241,12 @@ function getAmqpBasicProperties(
   }
 }
 
-function getAmqpPropsWithoutHeaders(
+async function getAmqpPropsWithoutHeaders(
   originalMsg?: Pick<amqplib.Message, 'properties'>,
   pubOpts?: publishI.PublishOptionsI,
-): Partial<Omit<amqplib.MessageProperties, 'headers'>> {
+): Promise<Partial<Omit<amqplib.MessageProperties, 'headers'>>> {
   const forceOptions = pubOpts?.amqpPublishOpts
-  const correlationId = randomUUID()
+  const correlationId = (await getMdc())?.requestId ?? randomUUID()
 
   const inheritedProperties:
     | Omit<amqplib.MessageProperties, 'headers'>
@@ -295,3 +311,35 @@ function validateBeforePublish(
     }
   }
 }
+
+async function obtainUserId(
+  user?:
+    | string
+    | AmqpMessage
+    | Pick<amqplib.Options.Publish, 'headers' | 'userId'>,
+): Promise<string | undefined> {
+  if (typeof user === 'string') {
+    return user
+  }
+
+  if (isAmqpMessageClass(user)) {
+    return (
+      _.get(user, 'properties.userId') ||
+      _.get(user, `properties.headers[${MESSAGE_HEADERS.MESSAGE.USER_ID}]`)
+    )
+  }
+
+  // MessageProperties
+  if (typeof user === 'object') {
+    return (
+      _.get(user, 'userId') ||
+      _.get(user, `headers[${MESSAGE_HEADERS.MESSAGE.USER_ID}]`)
+    )
+  }
+
+  const mdc = await getMdc()
+
+  if (mdc !== undefined) {
+    return mdc.userId
+  }
+}

test/e2e/publish_to_user.spec.ts

@@ -1,10 +1,12 @@
+import { randomUUID } from 'node:crypto'
 import { IsString } from 'class-validator'
 import {
   constants,
   AmqpMessage,
   Message,
   MessageHandler,
   Scope,
+  mdc,
   publish,
 } from '../../src'
 import { irisTesting } from '../setup'
@@ -54,13 +56,14 @@ describe('Publishing non Scope.USER messages to user via publish.[publishToUser/
   })
 
   test('Sending SESSION event to USER should work', async () => {
+    const userId = randomUUID()
     const spyHandler = vi.spyOn(handler, 'publishInternalToUser')
     const headers = { foo: `bar_${Date.now()}` }
 
     await publish.publishToUser(
       EvtSession,
       { msg: 'do-publish-internal-to-user' },
-      'user-id',
+      userId,
       { amqpPublishOpts: { headers } },
     )
 
@@ -78,21 +81,28 @@ describe('Publishing non Scope.USER messages to user via publish.[publishToUser/
           fields: expect.any(Object),
           content: expect.any(Buffer),
           properties: expect.objectContaining({
-            headers: expect.objectContaining(headers),
+            headers: expect.objectContaining({
+              ...headers,
+              [constants.MESSAGE_HEADERS.MESSAGE.USER_ID]: userId,
+            }),
           }),
         }),
       )
     })
   })
 
   test('Sending INTERNAL event to USER should work', async () => {
+    const userId = randomUUID()
     const spyHandler = vi.spyOn(handler, 'publishInternalToUser')
-    const headers = { foo: `bar_${Date.now()}` }
+    const headers = {
+      foo: `bar_${Date.now()}`,
+      [constants.MESSAGE_HEADERS.MESSAGE.USER_ID]: userId,
+    }
 
     await publish.publishToUser(
       EvtInternal,
       { msg: 'do-publish-internal-to-user' },
-      'user-id',
+      undefined,
       { amqpPublishOpts: { headers } },
     )
 
@@ -117,7 +127,42 @@ describe('Publishing non Scope.USER messages to user via publish.[publishToUser/
     })
   })
 
-  test('Should throw if user-id is missing', async () => {
+  test('Sending INTERNAL event to USER should work via registered mdcProvider', async () => {
+    const userId = randomUUID()
+    const spyHandler = vi.spyOn(handler, 'publishInternalToUser')
+
+    mdc.registerMdcProvider(() => ({ userId }))
+
+    await publish.publishToUser(EvtInternal, {
+      msg: 'do-publish-internal-to-user',
+    })
+
+    await vi.waitFor(() => {
+      expect(spyHandler).toHaveBeenCalledTimes(2)
+      expect(spyHandler).toHaveBeenNthCalledWith(
+        1,
+        { msg: 'do-publish-internal-to-user' },
+        expect.anything(),
+      )
+      expect(spyHandler).toHaveBeenNthCalledWith(
+        2,
+        { msg: 'hello' },
+        expect.objectContaining({
+          fields: expect.any(Object),
+          content: expect.any(Buffer),
+          properties: expect.objectContaining({
+            headers: expect.objectContaining({
+              [constants.MESSAGE_HEADERS.MESSAGE.USER_ID]: userId,
+            }),
+          }),
+        }),
+      )
+    })
+
+    mdc.registerMdcProvider(undefined)
+  })
+
+  test('Should throw if userId is missing', async () => {
     await expect(
       publish.publishToUser(
         EvtInternal,