From 68add2d588c53b59350b4b3ec5cc5183ed8e1735 Mon Sep 17 00:00:00 2001 From: Kyriakos Lesgidis Date: Thu, 29 Oct 2026 13:16:21 +0200 Subject: [PATCH] [PROD-92594] Gracefully shutdown kafka consumers v2 --- __snapshots__/diamorphosis.test.ts.js | 21 ++- docs/integrations/kafka.md | 20 +++ src/initializers/diamorphosis.ts | 3 +- src/initializers/kafka/base-kafka-handler.ts | 19 ++- src/initializers/kafka/index.ts | 6 + src/initializers/kafka/kafka.ts | 86 ++++++++-- src/typings/kafka.d.ts | 60 ++++--- test/diamorphosis/diamorphosis.test.ts | 3 +- test/initializers/kafka/index.test.ts | 7 +- test/initializers/kafka/kafka.test.ts | 171 ++++++++++++++++++- 10 files changed, 336 insertions(+), 60 deletions(-) diff --git a/__snapshots__/diamorphosis.test.ts.js b/__snapshots__/diamorphosis.test.ts.js index 83691a0d..3934b534 100644 --- a/__snapshots__/diamorphosis.test.ts.js +++ b/__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, diff --git a/docs/integrations/kafka.md b/docs/integrations/kafka.md index ec2da331..d107f288 100644 --- a/docs/integrations/kafka.md +++ b/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: diff --git a/src/initializers/diamorphosis.ts b/src/initializers/diamorphosis.ts index 71b98934..60ab5d35 100644 --- a/src/initializers/diamorphosis.ts +++ b/src/initializers/diamorphosis.ts @@ -215,7 +215,8 @@ function addKafkaConfig(config) { } }, connectionTimeout: 5000, - authenticationTimeout: 10000 + authenticationTimeout: 10000, + shutdownTimeoutMs: config.kafka?.shutdownTimeoutMs ?? 25000 }; } diff --git a/src/initializers/kafka/base-kafka-handler.ts b/src/initializers/kafka/base-kafka-handler.ts index 5707c8e1..40140183 100644 --- a/src/initializers/kafka/base-kafka-handler.ts +++ b/src/initializers/kafka/base-kafka-handler.ts @@ -53,15 +53,16 @@ abstract class Base { 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) { diff --git a/src/initializers/kafka/index.ts b/src/initializers/kafka/index.ts index 1c6687ab..e1d99d0b 100644 --- a/src/initializers/kafka/index.ts +++ b/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(); diff --git a/src/initializers/kafka/kafka.ts b/src/initializers/kafka/kafka.ts index 18beb461..d981b3b5 100644 --- a/src/initializers/kafka/kafka.ts +++ b/src/initializers/kafka/kafka.ts @@ -6,6 +6,10 @@ 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; @@ -13,13 +17,15 @@ export default class OrkaKafka { public produceClient: KafkajsType.Kafka; public producer: KafkajsType.Producer; private readonly sigtermHandler: () => void; - private readonly consumers: Map = new Map(); + private readonly consumers = new Set(); + private readonly consumerGroupIds = new WeakMap(); + private shutdownPromise?: Promise; 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 { + if (!this.shutdownPromise) { + this.shutdownPromise = this.drain(); + } + return this.shutdownPromise; + } + + private drain(): Promise { + 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: { diff --git a/src/typings/kafka.d.ts b/src/typings/kafka.d.ts index 2599581e..9ef6bf24 100644 --- a/src/typings/kafka.d.ts +++ b/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,19 +34,22 @@ 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'; @@ -51,4 +57,6 @@ export interface KafkaConfig { }; connectionTimeout?: number; authenticationTimeout?: number; + /** How long to wait for consumers and the producer to disconnect on SIGTERM. Defaults to 25000. */ + shutdownTimeoutMs?: number; } diff --git a/test/diamorphosis/diamorphosis.test.ts b/test/diamorphosis/diamorphosis.test.ts index a97b349b..a6d58186 100644 --- a/test/diamorphosis/diamorphosis.test.ts +++ b/test/diamorphosis/diamorphosis.test.ts @@ -92,7 +92,8 @@ describe('Diamorphosis Test', () => { sasl: { mechanism: '', username: '', password: '' } }, connectionTimeout: 5000, - authenticationTimeout: 10000 + authenticationTimeout: 10000, + shutdownTimeoutMs: 25000 }); }); diff --git a/test/initializers/kafka/index.test.ts b/test/initializers/kafka/index.test.ts index df0dcd12..f9398b85 100644 --- a/test/initializers/kafka/index.test.ts +++ b/test/initializers/kafka/index.test.ts @@ -1,8 +1,13 @@ import 'should'; -import init from '../../../src/initializers/kafka/index'; +import init, { shutdownKafka } from '../../../src/initializers/kafka/index'; describe('kafka init', () => { it('should handle the absence of kafka config', async () => { await init({} as any); }); + + it('shutdownKafka is a no-op when kafka is not initialized', async () => { + const result = await shutdownKafka(); + (result === undefined).should.eql(true); + }); }); diff --git a/test/initializers/kafka/kafka.test.ts b/test/initializers/kafka/kafka.test.ts index c3a72041..5c82238e 100644 --- a/test/initializers/kafka/kafka.test.ts +++ b/test/initializers/kafka/kafka.test.ts @@ -54,7 +54,7 @@ describe('kafka class', () => { kafkaStubReturn = { producer: () => producerStub, consumer: consumerStub, admin: adminStub }; kafkaStub = sinon.stub().returns(kafkaStubReturn); delete require.cache[require.resolve('../../../src/initializers/kafka/kafka')]; - mock('kafkajs', { Kafka: kafkaStub, Partitioners }); + mock('kafkajs', { Kafka: kafkaStub, Partitioners }); ({ default: Kafka } = await import('../../../src/initializers/kafka/kafka')); }); @@ -278,7 +278,7 @@ describe('kafka class', () => { ]); }); - it('should register sigtermHandler in SIGTERM process listeners and remove on disconnect', async function() { + it('should register sigtermHandler in SIGTERM process listeners and remove on disconnect', async function () { kafka = new Kafka({ certificates: { key: 'key', cert: 'cert', ca: 'ca', rejectUnauthorized: false }, groupId: 'groupId', @@ -296,7 +296,7 @@ describe('kafka class', () => { process.listeners('SIGTERM').should.containEql(handler); - kafka.disconnect(); + await kafka.disconnect(); process.listeners('SIGTERM').should.not.containEql(handler); }); @@ -317,9 +317,10 @@ describe('kafka class', () => { await kafka.createConsumer(); (kafka as any).handleSigterm(); + await kafka.shutdown(); consumerStub().disconnect.calledOnce.should.eql(true); - }) + }); }); }); @@ -448,4 +449,166 @@ describe('kafka class', () => { ]); }); }); + + describe('shutdown', function () { + function kafkaOptions(overrides = {}) { + return { + certificates: { key: 'key', cert: 'cert', ca: 'ca', rejectUnauthorized: false }, + groupId: 'groupId', + clientId: 'clientId', + brokers: ['broker-consumer'], + producer: { + brokers: ['broker-producer'], + certificates: { key: 'key', cert: 'cert', ca: 'ca', rejectUnauthorized: false } + }, + connectionTimeout: 5000, + authenticationTimeout: 10000, + ...overrides + }; + } + + function consumerDouble() { + return { + connect: sandbox.stub().resolves(), + disconnect: sandbox.stub().resolves(), + on: sandbox.stub() + }; + } + + it('keeps the SIGTERM listener until shutdown finishes', async () => { + kafka = new Kafka(kafkaOptions()); + const hanging = consumerDouble(); + let release: () => void = () => undefined; + hanging.disconnect.returns( + new Promise(resolve => { + release = resolve; + }) + ); + consumerStub.returns(hanging); + await kafka.createConsumer(); + + const handler = (kafka as any).sigtermHandler; + const pending = kafka.shutdown(); + + process.listeners('SIGTERM').should.containEql(handler); + + release(); + await pending; + + process.listeners('SIGTERM').should.not.containEql(handler); + }); + + it('disconnects the producer only after every consumer has settled', async () => { + kafka = new Kafka(kafkaOptions()); + await kafka.connect(); + + const first = consumerDouble(); + const second = consumerDouble(); + let releaseFirst: () => void = () => undefined; + first.disconnect.returns( + new Promise(resolve => { + releaseFirst = resolve; + }) + ); + consumerStub.onCall(0).returns(first); + consumerStub.onCall(1).returns(second); + + await kafka.createConsumer({ groupId: 'group-a' }); + await kafka.createConsumer({ groupId: 'group-b' }); + + const pending = kafka.shutdown(); + + first.disconnect.calledOnce.should.eql(true); + second.disconnect.calledOnce.should.eql(true); + producerStub.disconnect.called.should.eql(false); + + releaseFirst(); + const result = await pending; + + result.should.eql({ timedOut: false }); + producerStub.disconnect.calledOnce.should.eql(true); + }); + + it('keeps draining when a consumer disconnect rejects', async () => { + kafka = new Kafka(kafkaOptions()); + await kafka.connect(); + + const first = consumerDouble(); + const second = consumerDouble(); + first.disconnect.rejects(new Error('disconnect failed')); + consumerStub.onCall(0).returns(first); + consumerStub.onCall(1).returns(second); + + await kafka.createConsumer({ groupId: 'group-a' }); + await kafka.createConsumer({ groupId: 'group-b' }); + + const result = await kafka.shutdown(); + + result.should.eql({ timedOut: false }); + first.disconnect.calledOnce.should.eql(true); + second.disconnect.calledOnce.should.eql(true); + producerStub.disconnect.calledOnce.should.eql(true); + }); + + it('resolves with timedOut when disconnect never settles', async () => { + kafka = new Kafka(kafkaOptions({ shutdownTimeoutMs: 50 })); + await kafka.connect(); + + const hanging = consumerDouble(); + hanging.disconnect.returns(new Promise(() => undefined)); + consumerStub.returns(hanging); + await kafka.createConsumer(); + + const clock = sandbox.useFakeTimers(); + const pending = kafka.shutdown(); + await clock.tickAsync(50); + const result = await pending; + + result.should.eql({ timedOut: true }); + producerStub.disconnect.called.should.eql(false); + }); + + it('shares one drain across shutdown and SIGTERM', async () => { + kafka = new Kafka(kafkaOptions()); + await kafka.createConsumer(); + + (kafka as any).handleSigterm(); + const pending = kafka.shutdown(); + const again = kafka.shutdown(); + + pending.should.equal(again); + await pending; + consumerStub().disconnect.calledOnce.should.eql(true); + }); + + it('disconnects every consumer that shares a group id', async () => { + kafka = new Kafka(kafkaOptions()); + const first = consumerDouble(); + const second = consumerDouble(); + consumerStub.onCall(0).returns(first); + consumerStub.onCall(1).returns(second); + + await kafka.createConsumer({ groupId: 'shared' }); + await kafka.createConsumer({ groupId: 'shared' }); + await kafka.shutdown(); + + first.disconnect.calledOnce.should.eql(true); + second.disconnect.calledOnce.should.eql(true); + }); + + it('does not connect a consumer after shutdown has started', async () => { + kafka = new Kafka(kafkaOptions()); + await kafka.shutdown(); + + let message = ''; + try { + await kafka.createConsumer(); + } catch (err) { + message = err.message; + } + + message.should.eql('Kafka shutdown in progress. Refusing to create a new consumer.'); + consumerStub.called.should.eql(false); + }); + }); });