[PROD-92594] Gracefully shutdown kafka consumers v2

__snapshots__/diamorphosis.test.ts.js

@@ -154,7 +154,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when nothing
       }
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "queue": {
     "url": "",
@@ -389,7 +390,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when console
       }
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "queue": {
     "url": "",
@@ -624,7 +626,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when console
       }
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "queue": {
     "url": "",
@@ -859,7 +862,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when console
       }
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "queue": {
     "url": "",
@@ -1094,7 +1098,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when console
       }
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "queue": {
     "url": "",
@@ -1218,7 +1223,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when console
       "password": ""
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "healthCheck": {
     "kafka": true,
@@ -1458,7 +1464,8 @@ exports['Diamorphosis Test should set json/console loggingvariables when console
       "password": ""
     },
     "connectionTimeout": 5000,
-    "authenticationTimeout": 10000
+    "authenticationTimeout": 10000,
+    "shutdownTimeoutMs": 25000
   },
   "healthCheck": {
     "kafka": true,

docs/integrations/kafka.md

@@ -57,6 +57,7 @@ module.exports = {
       "username": "",
       "password": ""
     },
+    "shutdownTimeoutMs": 25000, // How long to wait for consumers and the producer to disconnect on SIGTERM. Overridable with KAFKA_SHUTDOWN_TIMEOUT_MS
     "producer": { // Can be given if the producer is using a different kafka broker than the consumer. Might be useful for migrating brokers.
       "brokers": [],
       "ssl": true,
@@ -172,6 +173,25 @@ await kafka.producer.send({
 
 See [kafkaJS producer](https://kafka.js.org/docs/producing) for more.
 
+### Graceful shutdown
+
+On `SIGTERM`, orka stops fetching, waits for in-flight messages to finish, commits their offsets, leaves the consumer group, and then disconnects the producer. The wait is bounded by `kafka.shutdownTimeoutMs` (default `25000`). Override it with the `KAFKA_SHUTDOWN_TIMEOUT_MS` environment variable, and keep it under the pod's `terminationGracePeriodSeconds`.
+
+The listener stays registered for the whole drain, so Honeybadger will not force-exit the process while a message is still running.
+
+orka does not exit the process. With that listener registered, Node also will not exit on its own while other handles (a database pool, Redis) are open, and the process would otherwise sit until Kubernetes sends `SIGKILL`. Consumer apps should drain and then exit:
+
+```js
+const { shutdownKafka } = require('@workablehr/orka');
+
+process.on('SIGTERM', async () => {
+  await shutdownKafka();
+  process.exit(0);
+});
+```
+
+`shutdownKafka()` resolves with `{ timedOut: false }` when the drain finishes, or `{ timedOut: true }` if `shutdownTimeoutMs` elapses first. It is a no-op when Kafka was not initialized, and calling it more than once returns the same drain.
+
 ### Creating topics programmatically
 
 You can easily create topics programmatically with:

src/initializers/diamorphosis.ts

@@ -215,7 +215,8 @@ function addKafkaConfig(config) {
       }
     },
     connectionTimeout: 5000,
-    authenticationTimeout: 10000
+    authenticationTimeout: 10000,
+    shutdownTimeoutMs: config.kafka?.shutdownTimeoutMs ?? 25000
   };
 }
 

src/initializers/kafka/base-kafka-handler.ts

@@ -53,15 +53,16 @@ abstract class Base<Input, Output> {
     this.jsonParseValue = options.jsonParseValue ?? true;
     this.runOptions = options.runOptions;
     this.logger = logger || getLogger(`kafka.consumer.${this.topic}`);
-    kafka.createConsumer(options.consumerOptions).then(c => {
-      this.consumer = c;
-      if (onConsumerCreated) {
-        onConsumerCreated(c);
-      }
-      this.consume()
-        .then(() => this.logger.info(`Kafka consumer connected to topic: ${this.topic}`))
-        .catch(err => this.logger.error(err));
-    });
+    kafka
+      .createConsumer(options.consumerOptions)
+      .then(c => {
+        this.consumer = c;
+        if (onConsumerCreated) {
+          onConsumerCreated(c);
+        }
+        return this.consume().then(() => this.logger.info(`Kafka consumer connected to topic: ${this.topic}`));
+      })
+      .catch(err => this.logger.error(err));
   }
 
   parseValue(buffer: Buffer) {

src/initializers/kafka/index.ts

@@ -14,6 +14,7 @@ export {
   BaseKafkaHandlerOptions,
   BaseKafkaHandlerBatch
 } from './base-kafka-handler';
+export type { ShutdownResult } from './kafka';
 
 let kafka: Kafka;
 export default async (
@@ -64,6 +65,11 @@ export const disconnectProducer = async () => {
   await kafka.disconnect();
 };
 
+export const shutdownKafka = async () => {
+  if (!kafka) return;
+  return kafka.shutdown();
+};
+
 export const isHealthy = () => {
   if (!kafka) return;
   return kafka.isHealthy();

src/initializers/kafka/kafka.ts

@@ -6,20 +6,26 @@ import { flatten, isEmpty } from 'lodash';
 const { Kafka, Partitioners }: typeof KafkajsType = requireInjected('kafkajs');
 const logger = getLogger('orka.kafka');
 
+const DEFAULT_SHUTDOWN_TIMEOUT_MS = 25000;
+
+export type ShutdownResult = { timedOut: boolean };
+
 export default class OrkaKafka {
   private options: KafkaConfig;
   private healthy: boolean = true;
   public consumeClient: KafkajsType.Kafka;
   public produceClient: KafkajsType.Kafka;
   public producer: KafkajsType.Producer;
   private readonly sigtermHandler: () => void;
-  private readonly consumers: Map<string, KafkajsType.Consumer> = new Map();
+  private readonly consumers = new Set<KafkajsType.Consumer>();
+  private readonly consumerGroupIds = new WeakMap<KafkajsType.Consumer, string>();
+  private shutdownPromise?: Promise<ShutdownResult>;
 
   constructor(options: KafkaConfig) {
     this.options = options;
     this.sigtermHandler = this.handleSigterm.bind(this);
 
-    process.once('SIGTERM', this.sigtermHandler);
+    process.on('SIGTERM', this.sigtermHandler);
   }
 
   public async connect(options?: KafkajsType.ProducerConfig) {
@@ -67,6 +73,9 @@ export default class OrkaKafka {
   }
 
   public async createConsumer({ groupId, ...rest }: KafkajsType.ConsumerConfig = {} as any) {
+    if (this.shutdownPromise) {
+      throw new Error('Kafka shutdown in progress. Refusing to create a new consumer.');
+    }
     groupId ??= this.options.groupId;
     const { brokers, clientId, connectionTimeout, authenticationTimeout } = this.options;
     const kafkaLogger = getLogger('orka.kafka.consumer');
@@ -83,7 +92,8 @@ export default class OrkaKafka {
       authenticationTimeout
     });
     const consumer = this.consumeClient.consumer({ groupId, ...rest });
-    this.consumers.set(groupId, consumer);
+    this.consumers.add(consumer);
+    this.consumerGroupIds.set(consumer, groupId);
     await consumer.connect();
     logger.info(`Kafka consumer(${groupId}) connected ${brokers.join(', ')}`);
     return consumer;
@@ -119,7 +129,11 @@ export default class OrkaKafka {
           const [offsets] = await admin.fetchOffsets({ groupId, topics: [topic], resolveOffsets: false });
           if (offsets.partitions.every(({ offset }) => offset === '-1')) {
             // groupId is not configured
-            const [oldOffsets] = (await admin.fetchOffsets({ groupId: oldGroupId, topics: [topic], resolveOffsets: false }));
+            const [oldOffsets] = await admin.fetchOffsets({
+              groupId: oldGroupId,
+              topics: [topic],
+              resolveOffsets: false
+            });
             const knownOffsets = oldOffsets.partitions.filter(o => o.offset !== '-1');
             if (!isEmpty(knownOffsets)) await admin.setOffsets({ groupId, topic, partitions: knownOffsets });
             return { groupId, renamedFrom: oldGroupId, topic, oldOffsets: knownOffsets };
@@ -178,18 +192,68 @@ export default class OrkaKafka {
     return level;
   }
 
-  private handleSigterm() {
-    logger.info('Received SIGTERM. Cancelling kafka consumers..');
+  /**
+   * Stops every consumer, then the producer. In-flight messages finish and offsets are committed
+   * before the returned promise resolves. Repeated calls share one drain. Never rejects:
+   * a timeout resolves with `{ timedOut: true }`.
+   */
+  public shutdown(): Promise<ShutdownResult> {
+    if (!this.shutdownPromise) {
+      this.shutdownPromise = this.drain();
+    }
+    return this.shutdownPromise;
+  }
+
+  private drain(): Promise<ShutdownResult> {
+    logger.info('Received SIGTERM. Draining kafka consumers..');
+    const timeoutMs = this.options.shutdownTimeoutMs ?? DEFAULT_SHUTDOWN_TIMEOUT_MS;
+
+    return new Promise(resolve => {
+      let settled = false;
+      const timer = setTimeout(() => {
+        if (settled) return;
+        settled = true;
+        logger.warn(`Kafka shutdown timed out after ${timeoutMs}ms. Consumers may still be draining.`);
+        resolve({ timedOut: true });
+      }, timeoutMs);
 
-    for (const [groupId, consumer] of this.consumers) {
-      consumer.disconnect()
+      this.disconnectConsumersThenProducer()
+        .catch(err => {
+          logger.error('Kafka shutdown failed', err);
+        })
         .then(() => {
-          logger.debug(`Consumer for kafka group ${groupId} disconnected.`);
-        }).catch(e => {
-          logger.error(`Consumer for kafka group ${groupId} failed to disconnect`, e);
+          if (settled) return;
+          settled = true;
+          clearTimeout(timer);
+          resolve({ timedOut: false });
         });
+    });
+  }
+
+  private async disconnectConsumersThenProducer() {
+    await Promise.allSettled(Array.from(this.consumers).map(consumer => this.disconnectConsumer(consumer)));
+    try {
+      await this.disconnect();
+    } catch (err) {
+      logger.error('Kafka producer failed to disconnect', err);
+    }
+  }
+
+  private async disconnectConsumer(consumer: KafkajsType.Consumer) {
+    const groupId = this.consumerGroupIds.get(consumer);
+    try {
+      await consumer.disconnect();
+      logger.debug(`Consumer for kafka group ${groupId} disconnected.`);
+    } catch (err) {
+      logger.error(`Consumer for kafka group ${groupId} failed to disconnect`, err);
     }
   }
+
+  private handleSigterm() {
+    this.shutdown().catch(err => {
+      logger.error('Kafka shutdown failed', err);
+    });
+  }
 }
 
 function getAuthOptions(options: {

src/typings/kafka.d.ts

@@ -5,19 +5,22 @@ export interface KafkaConfig {
     ca: string;
     rejectUnauthorized: boolean;
   };
-  sasl?: {
-    mechanism: 'plain';
-    username: string;
-    password: string;
-  } | {
-    mechanism: 'scram-sha-256';
-    username: string;
-    password: string;
-  } | {
-    mechanism: 'scram-sha-512';
-    username: string;
-    password: string;
-  };
+  sasl?:
+    | {
+        mechanism: 'plain';
+        username: string;
+        password: string;
+      }
+    | {
+        mechanism: 'scram-sha-256';
+        username: string;
+        password: string;
+      }
+    | {
+        mechanism: 'scram-sha-512';
+        username: string;
+        password: string;
+      };
   ssl?: boolean;
   groupId: string;
   clientId: string;
@@ -31,24 +34,29 @@ export interface KafkaConfig {
       ca: string;
     };
     ssl?: boolean;
-    sasl?: {
-      mechanism: 'plain';
-      username: string;
-      password: string;
-    } | {
-      mechanism: 'scram-sha-256';
-      username: string;
-      password: string;
-    } | {
-      mechanism: 'scram-sha-512';
-      username: string;
-      password: string;
-    };
+    sasl?:
+      | {
+          mechanism: 'plain';
+          username: string;
+          password: string;
+        }
+      | {
+          mechanism: 'scram-sha-256';
+          username: string;
+          password: string;
+        }
+      | {
+          mechanism: 'scram-sha-512';
+          username: string;
+          password: string;
+        };
   };
   log?: {
     level: 'info' | 'debug' | 'error' | 'warn' | 'nothing';
     errorToWarn: string[];
   };
   connectionTimeout?: number;
   authenticationTimeout?: number;
+  /** How long to wait for consumers and the producer to disconnect on SIGTERM. Defaults to 25000. */
+  shutdownTimeoutMs?: number;
 }

test/diamorphosis/diamorphosis.test.ts

@@ -92,7 +92,8 @@ describe('Diamorphosis Test', () => {
             sasl: { mechanism: '', username: '', password: '' }
           },
           connectionTimeout: 5000,
-          authenticationTimeout: 10000
+          authenticationTimeout: 10000,
+          shutdownTimeoutMs: 25000
         });
       });