Feat: kafka coonsumer observability (#287)

* Adding request context to handlers

* Adding request context support

* Adding transactionObservabilityManager

* Lint fix

* Release prepare

* Fixing minor issue

* Coverage fix

* Adding request id header config param

* Removing todo

* Adding request context to publisher

* Improving coverage

Author: CarlosGamero

packages/kafka/lib/AbstractKafkaConsumer.ts

@@ -1,5 +1,11 @@
+import { randomUUID } from 'node:crypto'
 import { setTimeout } from 'node:timers/promises'
-import { InternalError, stringValueSerializer } from '@lokalise/node-core'
+import {
+  InternalError,
+  type TransactionObservabilityManager,
+  stringValueSerializer,
+} from '@lokalise/node-core'
+import type { QueueConsumerDependencies } from '@message-queue-toolkit/core'
 import {
   type ConsumeOptions,
   Consumer,
@@ -10,17 +16,20 @@ import {
   stringDeserializer,
 } from '@platformatic/kafka'
 import { AbstractKafkaService, type BaseKafkaOptions } from './AbstractKafkaService.ts'
-import { KafkaHandlerContainer } from './handler-container/KafkaHandlerContainer.js'
-import type { KafkaHandlerRouting } from './handler-container/KafkaHandlerRoutingBuilder.js'
-import type { KafkaHandler } from './handler-container/index.js'
+import { KafkaHandlerContainer } from './handler-container/KafkaHandlerContainer.ts'
+import type { KafkaHandlerRouting } from './handler-container/KafkaHandlerRoutingBuilder.ts'
+import type { KafkaHandler, RequestContext } from './handler-container/index.ts'
 import type { KafkaDependencies, TopicConfig } from './types.ts'
 
+export type KafkaConsumerDependencies = KafkaDependencies &
+  Pick<QueueConsumerDependencies, 'transactionObservabilityManager'>
+
 export type KafkaConsumerOptions<TopicsConfig extends TopicConfig[]> = BaseKafkaOptions &
   Omit<
-    ConsumerOptions<string, object, string, object>,
+    ConsumerOptions<string, object, string, string>,
     'deserializers' | 'autocommit' | 'bootstrapBrokers'
   > &
-  Omit<ConsumeOptions<string, object, string, object>, 'topics'> & {
+  Omit<ConsumeOptions<string, object, string, string>, 'topics'> & {
     handlers: KafkaHandlerRouting<TopicsConfig>
   }
 
@@ -33,14 +42,19 @@ const MAX_IN_MEMORY_RETRIES = 3
 export abstract class AbstractKafkaConsumer<
   TopicsConfig extends TopicConfig[],
 > extends AbstractKafkaService<TopicsConfig, KafkaConsumerOptions<TopicsConfig>> {
-  private readonly consumer: Consumer<string, object, string, object>
-  private consumerStream?: MessagesStream<string, object, string, object>
+  private readonly consumer: Consumer<string, object, string, string>
+  private consumerStream?: MessagesStream<string, object, string, string>
 
+  private readonly transactionObservabilityManager: TransactionObservabilityManager
   private readonly handlerContainer: KafkaHandlerContainer<TopicsConfig>
 
-  constructor(dependencies: KafkaDependencies, options: KafkaConsumerOptions<TopicsConfig>) {
+  constructor(
+    dependencies: KafkaConsumerDependencies,
+    options: KafkaConsumerOptions<TopicsConfig>,
+  ) {
     super(dependencies, options)
 
+    this.transactionObservabilityManager = dependencies.transactionObservabilityManager
     this.handlerContainer = new KafkaHandlerContainer<TopicsConfig>(
       options.handlers,
       options.messageTypeField,
@@ -54,7 +68,7 @@ export abstract class AbstractKafkaConsumer<
         key: stringDeserializer,
         value: jsonDeserializer,
         headerKey: stringDeserializer,
-        headerValue: jsonDeserializer,
+        headerValue: stringDeserializer,
       },
     })
   }
@@ -87,18 +101,15 @@ export abstract class AbstractKafkaConsumer<
     await this.consumer.close()
   }
 
-  /*
-  TODO: https://lokalise.atlassian.net/browse/EDEXP-493
-    - Improve logging with logger child on constructor + add request context?
-    - Message logging
-    - Observability
-   */
-
-  private async consume(message: Message<string, object, string, object>): Promise<void> {
+  private async consume(message: Message<string, object, string, string>): Promise<void> {
     const handler = this.handlerContainer.resolveHandler(message.topic, message.value)
     // if there is no handler for the message, we ignore it (simulating subscription)
     if (!handler) return message.commit()
 
+    /* v8 ignore next */
+    const transactionId = this.resolveMessageId(message.value) ?? randomUUID()
+    this.transactionObservabilityManager?.start(this.buildTransactionName(message), transactionId)
+
     const parseResult = handler.schema.safeParse(message.value)
     if (!parseResult.success) {
       this.handlerError(parseResult.error, {
@@ -116,13 +127,19 @@ export abstract class AbstractKafkaConsumer<
 
     const validatedMessage = parseResult.data
 
+    const requestContext = this.getRequestContext(message)
+
     let retries = 0
     let consumed = false
     do {
       // exponential backoff -> 2^(retry-1)
       if (retries > 0) await setTimeout(Math.pow(2, retries - 1))
 
-      consumed = await this.tryToConsume({ ...message, value: validatedMessage }, handler.handler)
+      consumed = await this.tryToConsume(
+        { ...message, value: validatedMessage },
+        handler.handler,
+        requestContext,
+      )
       if (consumed) break
 
       retries++
@@ -142,15 +159,18 @@ export abstract class AbstractKafkaConsumer<
       })
     }
 
+    this.transactionObservabilityManager?.stop(transactionId)
+
     return message.commit()
   }
 
   private async tryToConsume<MessageValue extends object>(
-    message: Message<string, MessageValue, string, object>,
+    message: Message<string, MessageValue, string, string>,
     handler: KafkaHandler<MessageValue>,
+    requestContext: RequestContext,
   ): Promise<boolean> {
     try {
-      await handler(message)
+      await handler(message, requestContext)
       return true
     } catch (error) {
       this.handlerError(error, {
@@ -161,4 +181,27 @@ export abstract class AbstractKafkaConsumer<
 
     return false
   }
+
+  private buildTransactionName(message: Message<string, object, string, string>) {
+    const messageType = this.resolveMessageType(message.value)
+
+    let name = `kafka:${message.topic}`
+    if (messageType?.trim().length) name += `:${messageType.trim()}`
+
+    return name
+  }
+
+  private getRequestContext(message: Message<string, object, string, string>): RequestContext {
+    let reqId = message.headers.get(this.resolveHeaderRequestIdField())
+    if (!reqId || reqId.trim().length === 0) reqId = randomUUID()
+
+    return {
+      reqId,
+      logger: this.logger.child({
+        'x-request-id': reqId,
+        topic: message.topic,
+        messageKey: message.key,
+      }),
+    }
+  }
 }

packages/kafka/lib/AbstractKafkaPublisher.ts

@@ -8,6 +8,7 @@ import {
   stringSerializer,
 } from '@platformatic/kafka'
 import { AbstractKafkaService, type BaseKafkaOptions } from './AbstractKafkaService.ts'
+import type { RequestContext } from './handler-container/index.js'
 import type {
   KafkaDependencies,
   SupportedMessageValuesInput,
@@ -17,12 +18,12 @@ import type {
 } from './types.ts'
 
 export type KafkaPublisherOptions<TopicsConfig extends TopicConfig[]> = BaseKafkaOptions &
-  Omit<ProduceOptions<string, object, string, object>, 'serializers'> & {
+  Omit<ProduceOptions<string, object, string, string>, 'serializers'> & {
     topicsConfig: TopicsConfig
   }
 
 export type KafkaMessageOptions = Omit<
-  MessageToProduce<string, object, string, object>,
+  MessageToProduce<string, object, string, string>,
   'topic' | 'value'
 >
 
@@ -35,7 +36,7 @@ export abstract class AbstractKafkaPublisher<
     MessageSchemaContainer<SupportedMessageValuesInput<TopicsConfig>>
   >
 
-  private readonly producer: Producer<string, object, string, object>
+  private readonly producer: Producer<string, object, string, string>
   private isInitiated: boolean
 
   constructor(dependencies: KafkaDependencies, options: KafkaPublisherOptions<TopicsConfig>) {
@@ -61,7 +62,7 @@ export abstract class AbstractKafkaPublisher<
         key: stringSerializer,
         value: jsonSerializer,
         headerKey: stringSerializer,
-        headerValue: jsonSerializer,
+        headerValue: stringSerializer,
       },
     })
   }
@@ -91,6 +92,7 @@ export abstract class AbstractKafkaPublisher<
   async publish<Topic extends SupportedTopics<TopicsConfig>>(
     topic: Topic,
     message: SupportedMessageValuesInputForTopic<TopicsConfig, Topic>,
+    requestContext?: RequestContext,
     options?: KafkaMessageOptions,
   ): Promise<void> {
     const schemaResult = this.schemaContainers[topic]?.resolveSchema(message)
@@ -102,8 +104,15 @@ export abstract class AbstractKafkaPublisher<
     try {
       const parsedMessage = schemaResult.result.parse(message)
 
+      const headers = {
+        ...options?.headers,
+        [this.resolveHeaderRequestIdField()]: requestContext?.reqId ?? '',
+      }
+
       // biome-ignore lint/style/noNonNullAssertion: Should always exist due to lazy init
-      await this.producer!.send({ messages: [{ ...options, topic, value: parsedMessage }] })
+      await this.producer!.send({
+        messages: [{ ...options, topic, value: parsedMessage, headers }],
+      })
 
       this.handleMessageProcessed({
         message: parsedMessage,

packages/kafka/lib/AbstractKafkaService.ts

@@ -24,6 +24,12 @@ export type BaseKafkaOptions = {
   kafka: KafkaConfig
   messageTypeField?: string
   messageIdField?: string
+  /**
+   * The field in the message headers that contains the request ID.
+   * This is used to correlate logs and transactions with the request.
+   * Defaults to 'x-request-id'.
+   */
+  headerRequestIdField?: string
   handlerSpy?: HandlerSpy<object> | HandlerSpyParams | boolean
   logMessages?: boolean
 } & Omit<BaseOptions, keyof KafkaConfig> // Exclude properties that are already in KafkaConfig
@@ -67,6 +73,10 @@ export abstract class AbstractKafkaService<
     return message[this.options.messageIdField]
   }
 
+  protected resolveHeaderRequestIdField(): string {
+    return this.options.headerRequestIdField ?? 'x-request-id'
+  }
+
   protected handleMessageProcessed(params: {
     message: SupportedMessageValues<TopicsConfig>
     processingResult: MessageProcessingResult

packages/kafka/lib/handler-container/KafkaHandlerConfig.ts

@@ -1,8 +1,15 @@
+import type { CommonLogger } from '@lokalise/node-core'
 import type { Message } from '@platformatic/kafka'
 import type { ZodSchema } from 'zod'
 
+export interface RequestContext {
+  logger: CommonLogger
+  reqId: string
+}
+
 export type KafkaHandler<MessageValue extends object> = (
-  message: Message<string, MessageValue, string, object>,
+  message: Message<string, MessageValue, string, string>,
+  requestContext: RequestContext,
 ) => Promise<void> | void
 
 export class KafkaHandlerConfig<MessageValue extends object> {

packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.spec.ts

@@ -2,7 +2,7 @@ import type { Message } from '@platformatic/kafka'
 import { expectTypeOf } from 'vitest'
 import z from 'zod'
 import type { TopicConfig } from '../types.ts'
-import { KafkaHandlerConfig } from './KafkaHandlerConfig.ts'
+import { KafkaHandlerConfig, type RequestContext } from './KafkaHandlerConfig.ts'
 import { KafkaHandlerRoutingBuilder } from './KafkaHandlerRoutingBuilder.ts'
 
 const CREATE_SCHEMA = z.object({ type: z.literal('create') })
@@ -17,25 +17,29 @@ type TopicsConfig = typeof topicsConfig
 
 describe('KafkaHandlerRoutingBuilder', () => {
   it('should build routing config', () => {
-    type ExpectedMessage<MessageValue> = Message<string, MessageValue, string, object>
+    type ExpectedMessage<MessageValue> = Message<string, MessageValue, string, string>
+
     // Given
     const builder = new KafkaHandlerRoutingBuilder<TopicsConfig>()
       .addConfig(
         'all',
-        new KafkaHandlerConfig(CREATE_SCHEMA, (message) => {
+        new KafkaHandlerConfig(CREATE_SCHEMA, (message, requestContext) => {
           expectTypeOf(message).toEqualTypeOf<ExpectedMessage<z.infer<typeof CREATE_SCHEMA>>>()
+          expectTypeOf(requestContext).toEqualTypeOf<RequestContext>()
         }),
       )
       .addConfig(
         'all',
-        new KafkaHandlerConfig(UPDATE_SCHEMA, (message) => {
+        new KafkaHandlerConfig(UPDATE_SCHEMA, (message, requestContext) => {
           expectTypeOf(message).toEqualTypeOf<ExpectedMessage<z.infer<typeof UPDATE_SCHEMA>>>()
+          expectTypeOf(requestContext).toEqualTypeOf<RequestContext>()
         }),
       )
       .addConfig(
         'empty',
-        new KafkaHandlerConfig(EMPTY_SCHEMA, (message) => {
+        new KafkaHandlerConfig(EMPTY_SCHEMA, (message, requestContext) => {
           expectTypeOf(message).toEqualTypeOf<ExpectedMessage<z.infer<typeof EMPTY_SCHEMA>>>()
+          expectTypeOf(requestContext).toEqualTypeOf<RequestContext>()
         }),
       )
 

packages/kafka/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@message-queue-toolkit/kafka",
-    "version": "0.0.0",
+    "version": "0.1.0",
     "engines": {
         "node": ">= 22.14.0"
     },
@@ -54,7 +54,7 @@
         "zod": "^3.25.7"
     },
     "peerDependencies": {
-        "@message-queue-toolkit/core": ">=21.0.0",
+        "@message-queue-toolkit/core": ">=21.3.0",
         "@message-queue-toolkit/schemas": ">=6.0.0"
     },
     "devDependencies": {