From 2313c53e5d5f0c3da5e84bb31ae5d7066aa1c1be Mon Sep 17 00:00:00 2001 From: Thomas Flament Date: Wed, 2 Sep 2026 15:49:53 +0200 Subject: [PATCH 1/2] Migrate BackbeatConsumer.close to async/await Rewrite close() as an async _close() with the callback-based close(cb) kept as a thin shim for existing callers. The async.waterfall becomes sequential awaits and the event/interval waits become small Promise helpers; behavior is unchanged. Done as a dedicated commit ahead of the shutdown fix, per the repo's async/await migration guidance. Issue: BB-833 --- lib/BackbeatConsumer.js | 61 ++++++++++++++++++++--------------------- 1 file changed, 30 insertions(+), 31 deletions(-) diff --git a/lib/BackbeatConsumer.js b/lib/BackbeatConsumer.js index 8df60c538..a99d2f0d4 100644 --- a/lib/BackbeatConsumer.js +++ b/lib/BackbeatConsumer.js @@ -1,7 +1,6 @@ const { EventEmitter } = require('events'); const kafka = require('node-rdkafka'); const assert = require('assert'); -const async = require('async'); const joi = require('joi'); const jsutil = require('arsenal').jsutil; const Logger = require('werelogs').Logger; @@ -30,6 +29,13 @@ const { startLinkedSpanFromKafkaEntry } = require('arsenal/build/lib/tracing').k const CLIENT_ID = 'BackbeatConsumer'; const { withTopicPrefix } = require('./util/topic'); +// Resolve after `ms` (call-time setTimeout so test fake timers apply). +const delay = ms => new Promise(resolve => setTimeout(resolve, ms)); +// Resolve when `emitter` next emits `event` (errors are ignored, matching +// the callback-based waits this replaced). +const waitForEvent = (emitter, event) => + new Promise(resolve => emitter.once(event, resolve)); + /** * Stats on how we are consuming Kafka * @typedef {Object} ConsumerStats @@ -1221,42 +1227,35 @@ class BackbeatConsumer extends EventEmitter { * @return {undefined} */ close(cb) { + // Keep the callback signature for existing callers; errors are + // swallowed as before (cb takes no arguments). + this._close().then(() => cb(), () => cb()); + } + + async _close() { if (this._publishOffsetsCronTimer) { clearInterval(this._publishOffsetsCronTimer); this._publishOffsetsCronTimer = null; } - if (this._publishOffsetsCronActive) { - return setTimeout(() => this.close(cb), 1000); + while (this._publishOffsetsCronActive) { + await delay(1000); } this._circuitBreaker.stop(); - return async.waterfall([ - next => { - if (this._consumer?.isConnected()) { - const subscription = this._getSubscription(); - if (subscription !== null) { - this._consumer.unsubscribe(); - // Wait for partition unassign to complete before - // disconnecting, the rebalance callback will handle - // waiting for current jobs to complete as well as commit - // the latest offsets - this.once('unassign', () => next()); - return; - } - } - process.nextTick(next); - }, - next => { - if (this._zookeeper) { - this._zookeeper.close(); - } - if (this._consumer?.isConnected()) { - this._consumer.disconnect(); - this._consumer.once('disconnected', () => next()); - } else { - process.nextTick(next); - } - }, - ], () => cb()); + if (this._consumer?.isConnected() && + this._getSubscription() !== null) { + this._consumer.unsubscribe(); + // Wait for partition unassign to complete before disconnecting; + // the rebalance callback drains current jobs and commits the + // latest offsets. + await waitForEvent(this, 'unassign'); + } + if (this._zookeeper) { + this._zookeeper.close(); + } + if (this._consumer?.isConnected()) { + this._consumer.disconnect(); + await waitForEvent(this._consumer, 'disconnected'); + } } } From 2dfe4d5622fa3417109ca57dcfbb3b47d49b360a Mon Sep 17 00:00:00 2001 From: Thomas Flament Date: Wed, 2 Sep 2026 15:55:17 +0200 Subject: [PATCH 2/2] Leave the consumer group promptly on shutdown close() unsubscribed then waited for the rebalance callback to un-assign, but librdkafka delivers no such callback when the consumer holds no assignment, and postpones the unsubscribe outright while a rebalance is in progress. The wait was unbounded, so close() never returned, the SIGTERM handler never reached process.exit, and the pod was SIGKILLed with the member still registered -- the group then held zero partitions until session.timeout.ms (45s) evicted it. rd_kafka_consumer_close() (reached via disconnect()) performs the whole protocol exit from any state, including a postponed unsubscribe, and sends the LeaveGroup on its own; its one requirement is that any rebalance callback it delivers is answered. So instead of gating disconnect() on an un-assign that may never come: - wait for the drain only when a revoke is actually coming (one already draining, or triggered by the unsubscribe because we still hold partitions), bounded so a postponed unsubscribe gives up after a short grace and lets disconnect() leave; - answer every rebalance callback raised once shutting down (decline a grant, un-assign a revoke) so disconnect() cannot wedge on an unanswered one; - bound the disconnect wait itself; - stop fetching once the shutdown has started, so the pipeline the drain is waiting on cannot refill; - bound the wait for an in-flight offset publish, rather than rescheduling close() indefinitely. Issue: BB-833 --- lib/BackbeatConsumer.js | 186 +++++++++++++++++-- lib/constants.js | 10 ++ tests/functional/lib/BackbeatConsumer.js | 218 +++++++++++++++++++++-- tests/unit/backbeatConsumer.js | 218 +++++++++++++++++++++++ 4 files changed, 603 insertions(+), 29 deletions(-) diff --git a/lib/BackbeatConsumer.js b/lib/BackbeatConsumer.js index a99d2f0d4..9cb93693a 100644 --- a/lib/BackbeatConsumer.js +++ b/lib/BackbeatConsumer.js @@ -19,6 +19,8 @@ const { backbeatConsumer: { CONCURRENCY_DEFAULT, MAX_QUEUED_DEFAULT, + DISCONNECT_TIMEOUT_MS, + SHUTDOWN_DRAIN_GRACE_MS, } } = require('./constants'); @@ -31,10 +33,6 @@ const { withTopicPrefix } = require('./util/topic'); // Resolve after `ms` (call-time setTimeout so test fake timers apply). const delay = ms => new Promise(resolve => setTimeout(resolve, ms)); -// Resolve when `emitter` next emits `event` (errors are ignored, matching -// the callback-based waits this replaced). -const waitForEvent = (emitter, event) => - new Promise(resolve => emitter.once(event, resolve)); /** * Stats on how we are consuming Kafka @@ -201,6 +199,15 @@ class BackbeatConsumer extends EventEmitter { // since the last consume, and is used to trigger a new consume // immediately after the current one this._tasksCompletedSinceLastConsume = false; + // Set at close() entry: stops fetching, declines new grants + this._closing = false; + // Set at the disconnect step: rebalance callbacks are answered + // at once, since the client's close blocks on them + this._disconnecting = false; + // True between receiving a revoke and answering it (draining) + this._waitingDrain = false; + // Bounds close()'s wait for the drain + this._closeDrainTimeout = null; /** @type {ConsumerStats} */ this.consumerStats = { lag: {} }; @@ -438,6 +445,13 @@ class BackbeatConsumer extends EventEmitter { this._tryConsumedTimeout = null; } + // stop fetching once closing: new entries would refill the + // pipeline close() is draining, and this ends the + // self-rescheduling consume loop that would poll a closed client + if (this._closing) { + return undefined; + } + // use non-flowing mode of consumption to add some flow // control: explicit consumption of messages is required, // needs explicit polling to get new messages @@ -765,6 +779,19 @@ class BackbeatConsumer extends EventEmitter { */ _onRebalance(err, assignment) { if (err.code === kafka.CODES.ERRORS.ERR__ASSIGN_PARTITIONS) { + if (this._closing) { + // Decline the grant, but still answer the callback or the + // client stays parked. assign() is refused mid-disconnect, + // so answer with unassign(). + this._log.info('rdkafka.assign declined, shutting down', + { assignment }); + this._bestEffort('unassign', + () => this._consumer.unassign()); + if (!this._waitingDrain) { + this.emit('unassign', unassignStatus.SHUTDOWN); + } + return; + } this._log.info('rdkafka.assign', { assignment }); try { @@ -785,6 +812,19 @@ class BackbeatConsumer extends EventEmitter { ledger: this._offsetLedger.getProcessingCount(this._topic), }); + if (this._disconnecting) { + // The client's close delivers this revoke and blocks until + // answered; entries can no longer commit, so answer at + // once instead of draining. + KafkaBacklogMetrics.onRebalance(this._topic, this._groupId, + unassignStatus.SHUTDOWN); + this._bestEffort('unassign', + () => this._consumer.unassign()); + this.emit('unassign', unassignStatus.SHUTDOWN); + return; + } + + this._waitingDrain = true; const unassign = jsutil.once(status => { this._log.info(`processing queue ${status}, un-assigning`, { queueLen: this._processingQueue.length(), @@ -820,6 +860,7 @@ class BackbeatConsumer extends EventEmitter { logger.bind(this._log)('rdkafka.unassign failed', { e: e.toString(), assignment }); } + this._waitingDrain = false; this.emit('unassign', status); }; @@ -1150,6 +1191,47 @@ class BackbeatConsumer extends EventEmitter { return subscription.length === 0 ? null : subscription; } + /** + * Read the consumer's current partition assignment. + * + * @return {Object[]|null} assigned partitions, or null if unavailable + */ + _getAssignments() { + try { + return this._consumer.assignments(); + } catch (err) { + this._log.debug('could not read consumer assignments', { + method: 'BackbeatConsumer._getAssignments', + topic: this._topic, + groupId: this._groupId, + error: err.message, + }); + return null; + } + } + + /** + * Run a consumer operation that may fail while closing (state-dependent + * calls throw ERR__STATE mid-disconnect). + * + * @param {string} op - operation name, for logging + * @param {function} fn - operation to run + * @returns {undefined} + */ + _bestEffort(op, fn) { + try { + fn(); + } catch (e) { + const logger = this._consumer.isConnected() ? + this._log.error : this._log.debug; + logger.bind(this._log)(`rdkafka.${op} failed`, { + topic: this._topic, + groupId: this._groupId, + e: e.toString(), + }); + } + } + /** * Helper method to wait for consumers to be assigned to partitions before * proceeding to send canary messages. @@ -1233,30 +1315,106 @@ class BackbeatConsumer extends EventEmitter { } async _close() { + this._closing = true; if (this._publishOffsetsCronTimer) { clearInterval(this._publishOffsetsCronTimer); this._publishOffsetsCronTimer = null; } - while (this._publishOffsetsCronActive) { + // Wait briefly for an in-flight offset publish to finish, but never + // block shutdown on it: after the grace, proceed and let it error out. + const publishDeadline = Date.now() + SHUTDOWN_DRAIN_GRACE_MS; + while (this._publishOffsetsCronActive && Date.now() < publishDeadline) { await delay(1000); } this._circuitBreaker.stop(); - if (this._consumer?.isConnected() && - this._getSubscription() !== null) { - this._consumer.unsubscribe(); - // Wait for partition unassign to complete before disconnecting; - // the rebalance callback drains current jobs and commits the - // latest offsets. - await waitForEvent(this, 'unassign'); + + if (this._consumer?.isConnected()) { + // Best-effort unsubscribe: a no-op when not subscribed, and the + // trigger for the drain revoke when we hold partitions. + // Unconditional so a transient ERR__STATE on subscription() can't + // skip it. + this._bestEffort('unsubscribe', + () => this._consumer.unsubscribe()); + const assignments = this._getAssignments(); + const holdsPartitions = + assignments !== null && assignments.length > 0; + // Wait for the drain only when a revoke is coming: one already + // draining, or triggered by the unsubscribe because we still hold + // partitions. Otherwise no 'unassign' arrives and disconnect() + // leaves the group on its own, so waiting would hang. + if (this._waitingDrain || holdsPartitions) { + await this._waitForShutdownUnassign(); + } } + + // From here rebalance callbacks are answered at once: disconnect() + // blocks until every pending one is answered. + this._disconnecting = true; if (this._zookeeper) { this._zookeeper.close(); } if (this._consumer?.isConnected()) { - this._consumer.disconnect(); - await waitForEvent(this._consumer, 'disconnected'); + await this._disconnectWithTimeout(); } } + + /** + * Wait for the revoke handler to drain current jobs and emit 'unassign'. + * Once a revoke is draining it emits on its own -- on drain completion or + * its watchdog at maxPollIntervalMs - 1000 -- so just wait for it. Only + * bound the wait for a revoke to start: if the unsubscribe was postponed + * by an in-flight rebalance none arrives, so after a short grace give up + * and let disconnect() perform the leave. + * @return {Promise} resolves once un-assigned or the grace ends + */ + _waitForShutdownUnassign() { + return new Promise(resolve => { + const done = jsutil.once(() => { + this.removeListener('unassign', done); + clearTimeout(this._closeDrainTimeout); + this._closeDrainTimeout = null; + resolve(); + }); + this.once('unassign', done); + this._closeDrainTimeout = setTimeout(() => { + if (this._waitingDrain) { + return; + } + this._log.warn('closing: no un-assign after unsubscribe, ' + + 'proceeding with disconnect', { + topic: this._topic, + groupId: this._groupId, + }); + this._bestEffort('unassign', + () => this._consumer.unassign()); + done(); + }, SHUTDOWN_DRAIN_GRACE_MS); + }); + } + + /** + * Disconnect the client, bounded by DISCONNECT_TIMEOUT_MS. + * @return {Promise} resolves once disconnected or the bound ends + */ + _disconnectWithTimeout() { + return new Promise(resolve => { + const done = jsutil.once(() => { + clearTimeout(timeout); + resolve(); + }); + const timeout = setTimeout(() => { + this._log.error('closing: disconnect timed out', { + topic: this._topic, + groupId: this._groupId, + timeoutMs: DISCONNECT_TIMEOUT_MS, + }); + done(); + }, DISCONNECT_TIMEOUT_MS); + this._consumer.once('disconnected', done); + this._bestEffort('disconnect', + () => this._consumer.disconnect()); + }); + } } module.exports = BackbeatConsumer; diff --git a/lib/constants.js b/lib/constants.js index 4148feb2a..6aa5c2f71 100644 --- a/lib/constants.js +++ b/lib/constants.js @@ -25,6 +25,7 @@ const constants = { IDLE: 'idle', DRAINED: 'drained', TIMEOUT: 'timeout', + SHUTDOWN: 'shutdown', }, statusReady: 'READY', statusUndefined: 'UNDEFINED', @@ -87,6 +88,15 @@ const constants = { CONCURRENCY_DEFAULT: 1, // controls the max number of messages to queue for processing. MAX_QUEUED_DEFAULT: 1000, + // Bounds close()'s disconnect step: rd_kafka_consumer_close() only + // returns once its LeaveGroup is acknowledged, which can stall past + // the pod's grace period when that request is queued behind an + // in-flight JoinGroup. Return anyway after this so shutdown completes. + DISCONNECT_TIMEOUT_MS: 5000, + // Grace for an unsubscribe to produce its revoke before close() + // gives up and disconnects; only bounds the wait for a revoke to + // start, one in progress drains under the full budget. + SHUTDOWN_DRAIN_GRACE_MS: 5000, }, // Objects below this size skip the Expect: 100-continue handshake as the // extra round-trip is a bigger overhead than the bandwidth saved diff --git a/tests/functional/lib/BackbeatConsumer.js b/tests/functional/lib/BackbeatConsumer.js index 060cb3aee..a92aa2f90 100644 --- a/tests/functional/lib/BackbeatConsumer.js +++ b/tests/functional/lib/BackbeatConsumer.js @@ -2,10 +2,13 @@ const assert = require('assert'); const async = require('async'); const sinon = require('sinon'); const werelogs = require('werelogs'); +const { promisify } = require('util'); +const Kafka = require('node-rdkafka'); const { metrics } = require('arsenal'); const ZookeeperManager = require('../../../lib/clients/ZookeeperManager'); +const { withTopicPrefix } = require('../../../lib/util/topic'); const BackbeatProducer = require('../../../lib/BackbeatProducer'); const BackbeatConsumer = require('../../../lib/BackbeatConsumer'); const { BreakerState, CircuitBreaker } = require('breakbeat').CircuitBreaker; @@ -24,6 +27,22 @@ const consumerKafkaConf = { }; const log = new werelogs.Logger('BackbeatConsumer:test'); +function waitFor(predicate, timeoutMs, description) { + const deadline = Date.now() + timeoutMs; + return new Promise((resolve, reject) => { + const check = () => { + if (predicate()) { + return resolve(); + } + if (Date.now() > deadline) { + return reject(new Error(`timed out waiting for ${description}`)); + } + return setTimeout(check, 200); + }; + check(); + }); +} + describe('BackbeatConsumer main tests', () => { const topic = 'backbeat-consumer-spec'; const groupId = `replication-group-${Math.random()}`; @@ -1275,6 +1294,188 @@ describe('BackbeatConsumer shutdown tests', () => { }).timeout(60000); }); +describe('BackbeatConsumer group departure tests', () => { + const topic = 'backbeat-consumer-spec-departure'; + const groupId = `bucket-processor-${Math.random()}`; + // enough partitions for every member to hold a share, so the group is + // observably settled before a departure is timed + const wantedPartitions = 3; + // a member that never sends LeaveGroup is only evicted once + // session.timeout.ms (45s) expires; leaving cleanly takes a few seconds + const takeoverBudgetMs = 30000; + const adminTimeoutMs = 20000; + const settleTimeoutMs = 60000; + const messages = Array.from({ length: 12 }, (_, i) => + ({ key: `key-${i}`, message: `{"value":"${i}"}` })); + // read back rather than assumed: a topic left over from an earlier run + // may be wider than we asked for + let partitionCount; + let admin; + let producer; + let leaving; + let survivor; + let newcomer; + // the survivor holds its tasks so its un-assign stays deferred, which is + // what keeps a rebalance in progress for as long as a test needs + let inFlight; + let holdTasks; + + function createConsumer(clientId, queueProcessor) { + return new BackbeatConsumer({ + clientId, + zookeeper: zookeeperConf, + kafka: consumerKafkaConf, + groupId, + topic, + concurrency: 1, + queueProcessor, + }); + } + + function partitionsHeld(consumer) { + // a closed client throws rather than reporting an empty assignment + return consumer._consumer.isConnected() ? + consumer._consumer.assignments().length : 0; + } + + function totalPartitionsHeld(consumers) { + return consumers.reduce( + (total, consumer) => total + partitionsHeld(consumer), 0); + } + + function completeInFlight() { + holdTasks = false; + const pending = inFlight; + inFlight = []; + pending.forEach(cb => cb()); + } + + function close(consumer) { + return new Promise(resolve => consumer.close(resolve)); + } + + before(async function before() { + this.timeout(60000); + admin = Kafka.AdminClient.create({ + 'client.id': 'kafka-admin', + 'metadata.broker.list': consumerKafkaConf.hosts, + }); + const { ERR_TOPIC_ALREADY_EXISTS, ERR_INVALID_PARTITIONS } = + Kafka.CODES.ERRORS; + try { + await promisify(admin.createTopic).bind(admin)({ + topic: withTopicPrefix(topic), + num_partitions: wantedPartitions, // eslint-disable-line camelcase + replication_factor: 1, // eslint-disable-line camelcase + }, adminTimeoutMs); + } catch (err) { + if (err.code !== ERR_TOPIC_ALREADY_EXISTS) { + throw err; + } + // left over from an earlier run, possibly narrower than we need: + // widen it rather than waiting on partitions that never come + try { + await promisify(admin.createPartitions).bind(admin)( + withTopicPrefix(topic), wantedPartitions, adminTimeoutMs); + } catch (widenErr) { + if (widenErr.code !== ERR_INVALID_PARTITIONS) { + throw widenErr; + } + } + } + }); + + after(() => admin.disconnect()); + + beforeEach(async function beforeEach() { + this.timeout(120000); + inFlight = []; + holdTasks = true; + producer = new BackbeatProducer({ + kafka: producerKafkaConf, + topic, + pollIntervalMs: 100, + }); + leaving = createConsumer('BackbeatConsumer-leaving', + (message, cb) => cb()); + survivor = createConsumer('BackbeatConsumer-survivor', + (message, cb) => (holdTasks ? inFlight.push(cb) : cb())); + newcomer = createConsumer('BackbeatConsumer-newcomer', + (message, cb) => cb()); + await Promise.all([ + new Promise(resolve => producer.on('ready', resolve)), + ...[leaving, survivor, newcomer].map(consumer => + new Promise(resolve => consumer.on('ready', resolve))), + ]); + const metadata = await promisify(leaving._consumer.getMetadata) + .bind(leaving._consumer)({ topic: leaving._topic, timeout: adminTimeoutMs }); + partitionCount = metadata.topics + .find(entry => entry.name === leaving._topic).partitions.length; + assert(partitionCount >= wantedPartitions, + `topic has ${partitionCount} partitions, need ${wantedPartitions}`); + + leaving.subscribe(); + survivor.subscribe(); + await waitFor( + () => totalPartitionsHeld([leaving, survivor]) === partitionCount, + settleTimeoutMs, 'the group to settle'); + }); + + afterEach(async function afterEach() { + this.timeout(30000); + // start closing first so the consumers stop fetching, then release + // the held tasks so the drain they are waiting on can finish + const closed = Promise.all([leaving, survivor, newcomer].map(close)); + completeInFlight(); + await closed; + await new Promise(resolve => producer.close(resolve)); + }); + + it('should leave the group so the remaining member takes over promptly', + async function departure() { + this.timeout(90000); + const start = Date.now(); + await close(leaving); + await waitFor(() => partitionsHeld(survivor) === partitionCount, + takeoverBudgetMs, 'the survivor to take over every partition'); + const elapsed = Date.now() - start; + assert(elapsed < takeoverBudgetMs, + `takeover took ${elapsed}ms: the group was not left cleanly`); + }); + + it('should leave the group when closed during a rebalance', + async function departureWhileRebalancing() { + this.timeout(90000); + // a task in flight on the survivor defers its un-assign, which holds + // the rebalance triggered below open for as long as we need + producer.send(messages, assert.ifError); + await waitFor(() => inFlight.length > 0, settleTimeoutMs, + 'the survivor to start a task'); + newcomer.subscribe(); + // the join revokes both members: `leaving` releases its partitions and + // waits to rejoin while the survivor drains, so it is closed holding + // no assignment with the rebalance still in progress -- nothing will + // revoke back to it to complete the departure + await waitFor(() => partitionsHeld(leaving) === 0 + && partitionsHeld(survivor) > 0, + settleTimeoutMs, 'the member being closed to release its partitions'); + const start = Date.now(); + await close(leaving); + const closedIn = Date.now() - start; + assert(closedIn < takeoverBudgetMs, + `close() took ${closedIn}ms while rebalancing`); + completeInFlight(); + // the newcomer must be given a share, else the survivor merely still + // holds what it had and nothing was actually handed over + await waitFor(() => partitionsHeld(newcomer) > 0 + && totalPartitionsHeld([survivor, newcomer]) === partitionCount, + takeoverBudgetMs, 'the remaining members to share every partition'); + const elapsed = Date.now() - start; + assert(elapsed < takeoverBudgetMs, + `takeover took ${elapsed}ms: the group was not left cleanly`); + }); +}); + describe('BackbeatConsumer fromOffset tests', () => { const topic = 'backbeat-consumer-spec-from-offset'; let producer; @@ -1286,20 +1487,6 @@ describe('BackbeatConsumer fromOffset tests', () => { process.nextTick(cb); } - function waitFor(predicate, timeoutMs, description, cb) { - const deadline = Date.now() + timeoutMs; - const check = () => { - if (predicate()) { - return cb(); - } - if (Date.now() > deadline) { - return cb(new Error(`timed out waiting for ${description}`)); - } - return setTimeout(check, 200); - }; - check(); - } - function startConsumer(cb) { consumer = new BackbeatConsumer({ clientId: 'BackbeatConsumer-fromOffset', @@ -1369,7 +1556,8 @@ describe('BackbeatConsumer fromOffset tests', () => { // session timeout (45s) to settle when a rebalance hits // its joining window next => waitFor(() => consumedMessages.includes(marker), 75000, - 'the pre-existing message to be consumed', next), + 'the pre-existing message to be consumed') + .then(() => next(), next), ], done); }); }); diff --git a/tests/unit/backbeatConsumer.js b/tests/unit/backbeatConsumer.js index 951fad035..f97f7c635 100644 --- a/tests/unit/backbeatConsumer.js +++ b/tests/unit/backbeatConsumer.js @@ -1,10 +1,15 @@ const assert = require('assert'); const sinon = require('sinon'); +const { EventEmitter } = require('events'); const BackbeatConsumer = require('../../lib/BackbeatConsumer'); const { CODES } = require('node-rdkafka'); const { kafka } = require('../config.json'); +const { + unassignStatus, + backbeatConsumer: { SHUTDOWN_DRAIN_GRACE_MS }, +} = require('../../lib/constants'); const { BreakerState } = require('breakbeat').CircuitBreaker; class BackbeatConsumerMock extends BackbeatConsumer { @@ -397,4 +402,217 @@ describe('backbeatConsumer', () => { }); }); }); + + describe('shutdown', () => { + let consumer; + let mockConsumer; + + function makeMockConsumer(overrides) { + const mock = new EventEmitter(); + return Object.assign(mock, { + isConnected: () => true, + subscription: () => ['my-test-topic'], + assignments: () => [], + unsubscribe: sinon.stub(), + unassign: sinon.stub(), + assign: sinon.stub(), + commit: sinon.stub(), + pause: sinon.stub(), + resume: sinon.stub(), + consume: sinon.stub(), + disconnect: sinon.stub().callsFake(() => process.nextTick( + () => mock.emit('disconnected'))), + }, overrides); + } + + beforeEach(() => { + consumer = new BackbeatConsumerMock({ + kafka, + groupId: 'unittest-group', + topic: 'my-test-topic', + }); + mockConsumer = makeMockConsumer(); + consumer._consumer = mockConsumer; + }); + + afterEach(() => { + sinon.restore(); + }); + + describe('close', () => { + it('should not wait for an un-assign when nothing is assigned', + done => { + consumer.close(() => { + assert(mockConsumer.unsubscribe.calledOnce); + assert(mockConsumer.disconnect.calledOnce); + assert.strictEqual(consumer.listenerCount('unassign'), 0); + done(); + }); + }); + + it('should disconnect without waiting when reading the consumer ' + + 'state throws', done => { + const err = new Error('Local: Erroneous state'); + mockConsumer.subscription = sinon.stub().throws(err); + mockConsumer.assignments = sinon.stub().throws(err); + consumer.close(() => { + // unsubscribe is attempted best-effort regardless, but + // with no readable assignment we do not wait for an + // un-assign that may never come + assert.strictEqual(consumer.listenerCount('unassign'), 0); + assert(mockConsumer.disconnect.calledOnce); + done(); + }); + }); + + it('should not disconnect when not connected', done => { + mockConsumer.isConnected = () => false; + consumer.close(() => { + assert(mockConsumer.disconnect.notCalled); + done(); + }); + }); + + it('should wait for the revoke of a held assignment before ' + + 'disconnecting', done => { + mockConsumer.assignments = + () => [{ topic: 'my-test-topic', partition: 0 }]; + let closed = false; + consumer.close(() => { + closed = true; + assert(mockConsumer.disconnect.calledOnce); + done(); + }); + setImmediate(() => { + assert.strictEqual(closed, false); + assert(mockConsumer.disconnect.notCalled); + consumer.emit('unassign', unassignStatus.DRAINED); + }); + }); + + it('should wait for a parked revoke to finish draining before ' + + 'disconnecting', done => { + consumer._waitingDrain = true; + let closed = false; + consumer.close(() => { + closed = true; + assert(mockConsumer.disconnect.calledOnce); + done(); + }); + setImmediate(() => { + assert.strictEqual(closed, false); + assert(mockConsumer.disconnect.notCalled); + consumer.emit('unassign', unassignStatus.DRAINED); + }); + }); + + it('should give up waiting after the grace when no revoke ' + + 'starts, then disconnect', done => { + const clock = sinon.useFakeTimers({ toFake: ['setTimeout'] }); + mockConsumer.assignments = + () => [{ topic: 'my-test-topic', partition: 0 }]; + consumer.close(() => { + assert(mockConsumer.unassign.calledOnce); + assert(mockConsumer.disconnect.calledOnce); + done(); + }); + // no revoke arrives: the short grace elapses and we + // disconnect rather than waiting the full poll interval + clock.tick(SHUTDOWN_DRAIN_GRACE_MS + 1); + }); + + it('should keep waiting under the full budget while a revoke ' + + 'is draining', done => { + const clock = sinon.useFakeTimers({ toFake: ['setTimeout'] }); + consumer._waitingDrain = true; + let closed = false; + consumer.close(() => { + closed = true; + done(); + }); + // past the short grace: because a revoke is in progress we + // must not give up yet + clock.tick(SHUTDOWN_DRAIN_GRACE_MS + 1); + assert.strictEqual(closed, false); + assert(mockConsumer.disconnect.notCalled); + // the drain completes and emits 'unassign' + consumer._waitingDrain = false; + consumer.emit('unassign', unassignStatus.DRAINED); + }); + + it('should give up waiting for the disconnect event after ' + + 'DISCONNECT_TIMEOUT_MS', done => { + const clock = sinon.useFakeTimers({ toFake: ['setTimeout'] }); + mockConsumer.disconnect = sinon.stub(); + consumer.close(() => { + assert(mockConsumer.disconnect.calledOnce); + done(); + }); + process.nextTick(() => clock.tick(5001)); + }); + }); + + describe('_onRebalance during shutdown', () => { + it('should decline a partition grant with unassign()', () => { + consumer._closing = true; + let emitted = null; + consumer.once('unassign', status => { emitted = status; }); + consumer._onRebalance( + { code: CODES.ERRORS.ERR__ASSIGN_PARTITIONS }, + [{ topic: 'my-test-topic', partition: 0 }]); + assert(mockConsumer.assign.notCalled); + assert(mockConsumer.unassign.calledOnce); + assert.strictEqual(emitted, unassignStatus.SHUTDOWN); + }); + + it('should not signal un-assign for a declined grant while a ' + + 'revoke is still draining', () => { + consumer._closing = true; + consumer._waitingDrain = true; + let emitted = null; + consumer.once('unassign', status => { emitted = status; }); + consumer._onRebalance( + { code: CODES.ERRORS.ERR__ASSIGN_PARTITIONS }, + [{ topic: 'my-test-topic', partition: 0 }]); + assert(mockConsumer.unassign.calledOnce); + assert.strictEqual(emitted, null); + }); + + it('should answer a revoke immediately once disconnecting, ' + + 'without arming a drain', () => { + consumer._closing = true; + consumer._disconnecting = true; + let emitted = null; + consumer.once('unassign', status => { emitted = status; }); + consumer._onRebalance( + { code: CODES.ERRORS.ERR__REVOKE_PARTITIONS }, + [{ topic: 'my-test-topic', partition: 0 }]); + assert(mockConsumer.unassign.calledOnce); + assert.strictEqual(emitted, unassignStatus.SHUTDOWN); + assert.strictEqual(consumer._drainCallback, null); + assert.strictEqual(consumer._drainProcessQueueTimeout, null); + }); + + it('should keep the drain-first revoke handling when not ' + + 'shutting down', () => { + let emitted = null; + consumer.once('unassign', status => { emitted = status; }); + consumer._onRebalance( + { code: CODES.ERRORS.ERR__REVOKE_PARTITIONS }, + [{ topic: 'my-test-topic', partition: 0 }]); + assert(mockConsumer.commit.calledOnce); + assert(mockConsumer.unassign.calledOnce); + assert.strictEqual(emitted, unassignStatus.IDLE); + assert.strictEqual(consumer._waitingDrain, false); + }); + }); + + describe('_tryConsume during shutdown', () => { + it('should not consume once the shutdown has started', () => { + consumer._closing = true; + consumer._tryConsume(); + assert(mockConsumer.consume.notCalled); + }); + }); + }); });