Allow token refresh cb to be an async function

Author: milindl

e2e/oauthbearer_cb.spec.js

@@ -14,69 +14,88 @@ var eventListener = require('./listener');
 
 var kafkaBrokerList = process.env.KAFKA_HOST || 'localhost:9092';
 
-describe('Client with oauthbearer_cb', function () {
-  const oauthbearer_config = 'key=value';
-  let oauthbearer_cb_called = 0;
-  let oauthbearer_cb = function (config, cb) {
-    console.log("Called oauthbearer_cb with given config: " + config);
-    t.equal(config, oauthbearer_config);
-    oauthbearer_cb_called++;
-
-    // The broker is not expected to be configured for oauthbearer authentication.
-    // We just want to make sure that token refresh callback is triggered.
-    cb(new Error('oauthbearer_cb error'), null);
-  };
-
-  const commonConfig = {
-    'metadata.broker.list': kafkaBrokerList,
-    'debug': 'all',
-    'security.protocol': 'SASL_PLAINTEXT',
-    'sasl.mechanisms': 'OAUTHBEARER',
-    'oauthbearer_token_refresh_cb': oauthbearer_cb,
-    'sasl.oauthbearer.config': oauthbearer_config,
-  }
-
-  const checkClient = function (client, done) {
-    eventListener(client);
-
-    client.on('error', function (e) {
-      t.match(e.message, /oauthbearer_cb error/);
-    });
+const oauthbearer_config = 'key=value';
+let oauthbearer_cb_called = 0;
+
+let oauthbearer_cb_callback = function (config, cb) {
+  console.log("Called oauthbearer_cb with given config: " + config);
+  t.equal(config, oauthbearer_config);
+  oauthbearer_cb_called++;
+
+  // The broker is not expected to be configured for oauthbearer authentication.
+  // We just want to make sure that token refresh callback is triggered.
+  cb(new Error('oauthbearer_cb error'), null);
+};
+
+let oauthbearer_cb_async = async function (config) {
+  console.log("Called oauthbearer_cb with given config: " + config);
+  t.equal(config, oauthbearer_config);
+  oauthbearer_cb_called++;
+
+  // The broker is not expected to be configured for oauthbearer authentication.
+  // We just want to make sure that token refresh callback is triggered.
+  throw new Error('oauthbearer_cb error');
+};
+
+for (const oauthbearer_cb of [oauthbearer_cb_async, oauthbearer_cb_callback]) {
+  describe('Client with ' + (oauthbearer_cb.name), function () {
+
+    const commonConfig = {
+      'metadata.broker.list': kafkaBrokerList,
+      'debug': 'all',
+      'security.protocol': 'SASL_PLAINTEXT',
+      'sasl.mechanisms': 'OAUTHBEARER',
+      'oauthbearer_token_refresh_cb': oauthbearer_cb,
+      'sasl.oauthbearer.config': oauthbearer_config,
+    }
+
+    const checkClient = function (client, done, useCb) {
+      eventListener(client);
 
-    client.connect();
+      client.on('error', function (e) {
+        t.match(e.message, /oauthbearer_cb error/);
+      });
 
-    // We don't actually expect the connection to succeed, but we want to
-    // make sure that the oauthbearer_cb is called so give it a couple seconds.
-    setTimeout(() => {
-      client.disconnect();
-      client = null;
-      t.equal(oauthbearer_cb_called >= 1, true);
+      // The default timeout for the connect is 30s, so even if we
+      // call disconnect() midway, the test ends up being at least 30s.
+      client.connect({timeout: 2000});
+
+      // We don't actually expect the connection to succeed, but we want to
+      // make sure that the oauthbearer_cb is called so give it a couple seconds.
+      setTimeout(() => {
+        t.equal(oauthbearer_cb_called >= 1, true);
+        client.disconnect(() => {
+          done();
+        });
+        client = null;
+        if (!useCb) // for admin client, where disconnect is sync.
+          done();
+      }, 2000);
+    }
+
+    beforeEach(function (done) {
+      oauthbearer_cb_called = 0;
       done();
-    }, 2000);
-  }
+    });
 
-  beforeEach(function (done) {
-    oauthbearer_cb_called = 0;
-    done();
-  });
+    it('as producer', function (done) {
+      let producer = new Kafka.Producer(commonConfig);
+      checkClient(producer, done, true);
+      producer = null;
+    }).timeout(5000);
 
-  it('as producer', function (done) {
-    let producer = new Kafka.Producer(commonConfig);
-    checkClient(producer, done);
-    producer = null;
-  }).timeout(2500);
-
-  it('as consumer', function (done) {
-    const config = Object.assign({ 'group.id': 'gid' }, commonConfig);
-    let consumer = new Kafka.KafkaConsumer(config);
-    checkClient(consumer, done);
-    consumer = null;
-  }).timeout(2500);
-
-  it('as admin', function (done) {
-    let admin = new Kafka.AdminClient.create(commonConfig);
-    checkClient(admin, done);
-    admin = null;
-  }).timeout(2500);
-
-});
+    it('as consumer', function (done) {
+      const config = Object.assign({ 'group.id': 'gid' }, commonConfig);
+      let consumer = new Kafka.KafkaConsumer(config);
+      checkClient(consumer, done, true);
+      consumer = null;
+    }).timeout(5000);
+
+    it('as admin', function (done) {
+      let admin = new Kafka.AdminClient.create(commonConfig);
+      checkClient(admin, done, false);
+      admin = null;
+    }).timeout(5000);
+
+  });
+}

lib/client.js

@@ -109,7 +109,10 @@ function Client(globalConf, SubClientType, topicConf) {
         // we might want to disconnect without ever completing connection.
         return;
       }
-      savedCallback(oauthbearer_config, (err, token) => {
+
+      // This sets the token or error within librdkafka, and emits any
+      // errors on the emitter.
+      const postProcessTokenRefresh = (err, token) => {
         try {
           if (err) {
             throw err;
@@ -130,7 +133,19 @@ function Client(globalConf, SubClientType, topicConf) {
           this._client.setOAuthBearerTokenFailure(e.message);
           this.emit('error', e);
         }
-      });
+      };
+      const returnPromise = savedCallback(oauthbearer_config, postProcessTokenRefresh);
+
+      // If it looks like a promise, and quacks like a promise, it is a promise
+      // (or an async function). We expect the callback NOT to have been called
+      // in such a case.
+      if (returnPromise && (typeof returnPromise.then === 'function')) {
+        returnPromise.then((token) => {
+          postProcessTokenRefresh(null, token);
+        }).catch(err => {
+          postProcessTokenRefresh(err);
+        });
+      }
     }
   }