Merge pull request nestjs#9954 from kosh-b/feature/kafka-heartbeat-passed-to-endpoint

feat(microservices): add Kafka heartbeat callback to KafkaContext

Author: kamilmysliwiec

packages/microservices/ctx-host/kafka.context.ts

@@ -6,6 +6,7 @@ type KafkaContextArgs = [
   partition: number,
   topic: string,
   consumer: Consumer,
+  heartbeat: () => Promise<void>,
 ];
 
 export class KafkaContext extends BaseRpcContext<KafkaContextArgs> {
@@ -40,4 +41,11 @@ export class KafkaContext extends BaseRpcContext<KafkaContextArgs> {
   getConsumer() {
     return this.args[3];
   }
+
+  /**
+   * Returns the Kafka heartbeat callback.
+   */
+   getHeartbeat() {
+    return this.args[4];
+  } 
 }

packages/microservices/external/kafka.interface.ts

@@ -870,6 +870,8 @@ export interface EachMessagePayload {
   topic: string;
   partition: number;
   message: KafkaMessage;
+  heartbeat(): Promise<void>;
+  pause(): () => void;
 }
 
 export interface EachBatchPayload {

packages/microservices/server/server-kafka.ts

@@ -167,6 +167,7 @@ export class ServerKafka extends Server implements CustomTransportStrategy {
       payload.partition,
       payload.topic,
       this.consumer,
+      payload.heartbeat,
     ]);
     const handler = this.getHandlerByPattern(packet.pattern);
     // if the correlation id or reply topic is not set

packages/microservices/test/client/client-kafka.spec.ts

@@ -22,6 +22,8 @@ describe('ClientKafka', () => {
   const timestamp = new Date().toISOString();
   const attributes = 1;
   const messageValue = 'test-message';
+  const heartbeat = async () => {};
+  const pause = () => () => {};
 
   // message
   const message: KafkaMessage = {
@@ -57,6 +59,8 @@ describe('ClientKafka', () => {
       },
       message,
     ),
+    heartbeat,
+    pause,
   };
 
   const payloadDisposed: EachMessagePayload = {
@@ -75,6 +79,8 @@ describe('ClientKafka', () => {
         value: { test: true },
       },
     ),
+    heartbeat,
+    pause,
   };
 
   const payloadError: EachMessagePayload = {
@@ -93,6 +99,8 @@ describe('ClientKafka', () => {
         value: null,
       },
     ),
+    heartbeat,
+    pause,
   };
 
   const payloadWithoutCorrelation: EachMessagePayload = {
@@ -104,6 +112,8 @@ describe('ClientKafka', () => {
       },
       message,
     ),
+    heartbeat,
+    pause,
   };
 
   // deserialized payload
@@ -118,6 +128,8 @@ describe('ClientKafka', () => {
       },
       deserializedMessage,
     ),
+    heartbeat,
+    pause,
   };
 
   const deserializedPayloadDisposed: EachMessagePayload = {
@@ -136,6 +148,8 @@ describe('ClientKafka', () => {
         value: { test: true },
       },
     ),
+    heartbeat,
+    pause,
   };
 
   const deserializedPayloadError: EachMessagePayload = {
@@ -154,6 +168,8 @@ describe('ClientKafka', () => {
         value: null,
       },
     ),
+    heartbeat,
+    pause,
   };
 
   let client: ClientKafka;

packages/microservices/test/ctx-host/kafka.context.spec.ts

@@ -3,12 +3,12 @@ import { KafkaContext } from '../../ctx-host';
 import { Consumer, KafkaMessage } from '../../external/kafka.interface';
 
 describe('KafkaContext', () => {
-  const args = ['test', { test: true }, undefined, { test: 'consumer' }];
+  const args = ['test', { test: true }, undefined, { test: 'consumer' }, () => {}];
   let context: KafkaContext;
 
   beforeEach(() => {
     context = new KafkaContext(
-      args as [KafkaMessage, number, string, Consumer],
+      args as [KafkaMessage, number, string, Consumer, () => Promise<void>],
     );
   });
   describe('getTopic', () => {
@@ -31,4 +31,9 @@ describe('KafkaContext', () => {
       expect(context.getConsumer()).to.deep.eq({ test: 'consumer' });
     });
   });
+  describe('getHeartbeat', () => {
+    it('should return heartbeat callback', () => {
+      expect(context.getHeartbeat()).to.be.eql(args[4]);
+    });
+  });
 });

packages/microservices/test/server/server-kafka.spec.ts

@@ -26,6 +26,8 @@ describe('ServerKafka', () => {
   const key = '1';
   const timestamp = new Date().toISOString();
   const messageValue = 'test-message';
+  const heartbeat = async () => {};
+  const pause = () => () => {};
 
   const eventMessage: KafkaMessage = {
     key: Buffer.from(key),
@@ -44,6 +46,8 @@ describe('ServerKafka', () => {
       },
       eventMessage,
     ),
+    heartbeat,
+    pause,
   };
 
   const eventWithCorrelationIdPayload: EachMessagePayload = {
@@ -57,6 +61,8 @@ describe('ServerKafka', () => {
       },
       eventMessage,
     ),
+    heartbeat,
+    pause,
   };
 
   const message: KafkaMessage = Object.assign(
@@ -73,6 +79,8 @@ describe('ServerKafka', () => {
     topic,
     partition: 0,
     message,
+    heartbeat,
+    pause,
   };
 
   let server: ServerKafka;