Add timeout waiting for broker to be ready fixes #750

SOHU-Co · Aug 30, 2017 · 40e7ba1 · 40e7ba1
1 parent a580e80
commit 40e7ba1
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 6 deletions.
diff --git a/lib/kafkaClient.js b/lib/kafkaClient.js
@@ -476,7 +476,8 @@ KafkaClient.prototype.createBroker = function (host, port, longpolling) {
         if (error) {
           logger.error('error initialize broker after reconnect', error);
         } else {
-          brokerWrapper.emit('ready');
+          const readyEventName = brokerWrapper.getReadyEventName();
+          self.emit(readyEventName);
         }
         self.emit('reconnect');
       });
@@ -485,7 +486,8 @@ KafkaClient.prototype.createBroker = function (host, port, longpolling) {
         if (error) {
           logger.error('error initialize broker after connect', error);
         } else {
-          brokerWrapper.emit('ready');
+          const readyEventName = brokerWrapper.getReadyEventName();
+          self.emit(readyEventName);
         }
         self.emit('connect');
       });
@@ -609,6 +611,27 @@ function getSupportedForRequestType (broker, requestType) {
   };
 }
 
+KafkaClient.prototype.waitUntilReady = function (broker, callback) {
+  let timeoutId = null;
+
+  function onReady () {
+    logger.debug('broker is now ready');
+    clearTimeout(timeoutId);
+    timeoutId = null;
+    callback(null);
+  }
+
+  const timeout = this.options.requestTimeout;
+  const readyEventName = broker.getReadyEventName();
+
+  timeoutId = setTimeout(() => {
+    this.removeListener(readyEventName, onReady);
+    callback(new TimeoutError(`Request timed out after ${timeout}ms`));
+  }, timeout);
+
+  this.once(readyEventName, onReady);
+};
+
 KafkaClient.prototype.sendRequest = function (request, callback) {
   logger.debug('sending request');
   const payloads = this.payloadsByLeader(request.data.payloads);
@@ -644,10 +667,7 @@ KafkaClient.prototype.sendRequest = function (request, callback) {
     const broker = this.brokerForLeader(leader);
     if (broker.apiSupport == null) {
       logger.debug('missing apiSupport waiting until broker is ready...');
-      broker.once('ready', function () {
-        logger.debug('broker is now ready');
-        callback(null);
-      });
+      this.waitUntilReady(broker, callback);
     } else {
       logger.debug('has apiSupport broker is ready');
       callback(null);

diff --git a/lib/wrapper/BrokerWrapper.js b/lib/wrapper/BrokerWrapper.js
@@ -31,6 +31,11 @@ var BrokerWrapper = function (socket, noAckBatchOptions, idleConnectionMs) {
 
 util.inherits(BrokerWrapper, EventEmitter);
 
+BrokerWrapper.prototype.getReadyEventName = function () {
+  const lp = this.socket.longpolling ? '-longpolling' : '';
+  return `${this.socket.addr}${lp}-ready`;
+};
+
 BrokerWrapper.prototype.isConnected = function () {
   return !this.socket.destroyed && !this.socket.closing && !this.socket.error;
 };

diff --git a/test/test.kafkaClient.js b/test/test.kafkaClient.js
@@ -155,6 +155,59 @@ describe('Kafka Client', function () {
     });
   });
 
+  describe('#waitUntilReady', function () {
+    let sandbox, client, clock;
+
+    before(function () {
+      sandbox = sinon.sandbox.create();
+      clock = sandbox.useFakeTimers();
+      client = new Client({ kafkaHost: '127.0.0.1:9092', autoConnect: false, requestTimeout: 4000 });
+    });
+
+    afterEach(function () {
+      sandbox.restore();
+    });
+
+    it('should yield error timeout if broker is not ready by requestTimeout', function (done) {
+      const fakeBroker = new BrokerWrapper(new FakeSocket());
+      const readyKey = 'broker.host-ready';
+
+      sandbox.stub(fakeBroker, 'getReadyEventName').returns(readyKey);
+      sandbox.spy(client, 'removeListener');
+      sandbox.spy(client, 'once');
+
+      client.waitUntilReady(fakeBroker, function (error) {
+        error.should.not.be.empty;
+        error.should.be.an.instanceOf(TimeoutError);
+        sinon.assert.calledOnce(fakeBroker.getReadyEventName);
+        sinon.assert.calledWith(client.removeListener, readyKey, sinon.match.func);
+        sinon.assert.calledWith(client.once, readyKey, sinon.match.func);
+        done();
+      });
+
+      clock.tick(client.options.requestTimeout + 1);
+    });
+
+    it('should yield if broker is ready before requestTimeout', function (done) {
+      const fakeBroker = new BrokerWrapper(new FakeSocket());
+      const readyKey = 'broker.host-ready';
+
+      sandbox.stub(fakeBroker, 'getReadyEventName').returns(readyKey);
+      sandbox.spy(client, 'removeListener');
+      sandbox.spy(client, 'once');
+
+      client.waitUntilReady(fakeBroker, function (error) {
+        should(error).be.empty;
+        sinon.assert.calledOnce(fakeBroker.getReadyEventName);
+        sinon.assert.calledWith(client.once, readyKey, sinon.match.func);
+        done();
+      });
+
+      clock.tick(client.options.requestTimeout - 1);
+      client.emit(readyKey);
+    });
+  });
+
   describe('#wrapTimeoutIfNeeded', function () {
     let sandbox, wrapTimeoutIfNeeded, client, clock;