@point3/node-rdkafka 3.6.0-1

package/lib/kafka-consumer-stream.js

@@ -0,0 +1,397 @@
+/*
+ * node-rdkafka - Node.js wrapper for RdKafka C/C++ library
+ *
+ * Copyright (c) 2016 Blizzard Entertainment
+ *
+ * This software may be modified and distributed under the terms
+ * of the MIT license.  See the LICENSE.txt file for details.
+ */
+
+'use strict';
+
+module.exports = KafkaConsumerStream;
+
+var Readable = require('stream').Readable;
+var util = require('util');
+
+util.inherits(KafkaConsumerStream, Readable);
+
+/**
+ * ReadableStream integrating with the Kafka Consumer.
+ *
+ * This class is used to read data off of Kafka in a streaming way. It is
+ * useful if you'd like to have a way to pipe Kafka into other systems. You
+ * should generally not make this class yourself, as it is not even exposed
+ * as part of module.exports. Instead, you should KafkaConsumer.createReadStream.
+ *
+ * The stream implementation is slower than the continuous subscribe callback.
+ * If you don't care so much about backpressure and would rather squeeze
+ * out performance, use that method. Using the stream will ensure you read only
+ * as fast as you write.
+ *
+ * The stream detects if Kafka is already connected. If it is, it will begin
+ * reading. If it is not, it will connect and read when it is ready.
+ *
+ * This stream operates in objectMode. It streams {Consumer~Message}
+ *
+ * @param {Consumer} consumer - The Kafka Consumer object.
+ * @param {object} options - Options to configure the stream.
+ * @param {number} options.waitInterval - Number of ms to wait if Kafka reports
+ * that it has timed out or that we are out of messages (right now).
+ * @param {array} options.topics - Array of topics, or a function that parses
+ * metadata into an array of topics
+ * @constructor
+ * @extends stream.Readable
+ * @see Consumer~Message
+ */
+function KafkaConsumerStream(consumer, options) {
+  if (!(this instanceof KafkaConsumerStream)) {
+    return new KafkaConsumerStream(consumer, options);
+  }
+
+  if (options === undefined) {
+    options = { waitInterval: 1000 };
+  } else if (typeof options === 'number') {
+    options = { waitInterval: options };
+  } else if (options === null || typeof options !== 'object') {
+    throw new TypeError('"options" argument must be a number or an object');
+  }
+
+  var topics = options.topics;
+
+  if (typeof topics === 'function') {
+    // Just ignore the rest of the checks here
+  } else if (!Array.isArray(topics)) {
+    if (typeof topics !== 'string' && !(topics instanceof RegExp)) {
+      throw new TypeError('"topics" argument must be a string, regex, or an array');
+    } else {
+      topics = [topics];
+    }
+  }
+
+  options = Object.create(options);
+
+  var fetchSize = options.fetchSize || 1;
+
+  // Run in object mode by default.
+  if (options.objectMode === null || options.objectMode === undefined) {
+    options.objectMode = true;
+
+    // If they did not explicitly set high water mark, and we are running
+    // in object mode, set it to the fetch size + 2 to ensure there is room
+    // for a standard fetch
+    if (!options.highWaterMark) {
+      options.highWaterMark = fetchSize + 2;
+    }
+  }
+
+  if (options.objectMode !== true) {
+    this._read = this._read_buffer;
+  } else {
+    this._read = this._read_message;
+  }
+
+  Readable.call(this, options);
+
+  this.consumer = consumer;
+  this.topics = topics;
+  this.autoClose = options.autoClose === undefined ? true : !!options.autoClose;
+  this.waitInterval = options.waitInterval === undefined ? 1000 : options.waitInterval;
+  this.fetchSize = fetchSize;
+  this.connectOptions = options.connectOptions || {};
+  this.streamAsBatch = options.streamAsBatch || false;
+
+  // Hold the messages in here
+  this.messages = [];
+
+  var self = this;
+
+  this.consumer
+    .on('unsubscribed', function() {
+      // Invalidate the stream when we unsubscribe
+      self.push(null);
+    });
+
+  if (options.initOauthBearerToken) {
+    this.consumer.setOauthBearerToken(options.initOauthBearerToken);
+  }
+
+  // Call connect. Handles potentially being connected already
+  this.connect(this.connectOptions);
+
+  this.once('end', function() {
+    if (this.autoClose) {
+      this.destroy();
+    }
+  });
+
+}
+
+/**
+ * Refresh OAuthBearer token, initially provided in factory method.
+ * Expiry is always set to maximum value, as the callback of librdkafka
+ * for token refresh is not used.
+ *
+ * @param {string} tokenStr - OAuthBearer token string
+ * @see connection.cc
+ */
+KafkaConsumerStream.prototype.refreshOauthBearerToken = function (tokenStr) {
+  this.consumer.setOauthBearerToken(tokenStr);
+};
+
+/**
+ * Internal stream read method. This method reads message objects.
+ * @param {number} size - This parameter is ignored for our cases.
+ * @private
+ */
+KafkaConsumerStream.prototype._read_message = function(size) {
+  if (this.messages.length > 0) {
+    return this.push(this.messages.shift());
+  }
+
+  if (!this.consumer) {
+    // This consumer is set to `null` in the close function
+    return;
+  }
+
+  if (!this.consumer.isConnected()) {
+    this.consumer.once('ready', function() {
+      // This is the way Node.js does it
+      // https://github.com/nodejs/node/blob/master/lib/fs.js#L1733
+      this._read(size);
+    }.bind(this));
+    return;
+  }
+
+  if (this.destroyed) {
+    return;
+  }
+
+  var self = this;
+
+  // If the size (number of messages) we are being advised to fetch is
+  // greater than or equal to the fetch size, use the fetch size.
+  // Only opt to use the size in case it is LESS than the fetch size.
+  // Essentially, we want to use the smaller value here
+  var fetchSize = size >= this.fetchSize ? this.fetchSize : size;
+
+  this.consumer.consume(fetchSize, onread);
+
+  // Retry function. Will wait up to the wait interval, with some
+  // random noise if one is provided. Otherwise, will go immediately.
+  function retry() {
+    if (!self.waitInterval) {
+      setImmediate(function() {
+        self._read(size);
+      });
+    } else {
+      setTimeout(function() {
+        self._read(size);
+      }, self.waitInterval * Math.random()).unref();
+    }
+  }
+
+  function onread(err, messages) {
+
+    // If there was an error we still want to emit it.
+    // Essentially, if the user does not register an error
+    // handler, it will still cause the stream to blow up.
+    //
+    // But... if one is provided, consumption will move on
+    // as normal
+    if (err) {
+      self.emit('error', err);
+    }
+
+    // If there are no messages it means we reached EOF or a timeout.
+    // Do what we used to do
+
+    if (err || messages.length < 1) {
+      // If we got an error or if there were no messages, initiate a retry
+      retry();
+      return;
+    } else {
+      if (self.streamAsBatch) {
+        self.push(messages);
+      } else {
+        for (var i = 0; i < messages.length; i++) {
+          self.messages.push(messages[i]);
+        }
+
+        // Now that we have added them all the inner messages buffer,
+        // we can just push the most recent one
+        self.push(self.messages.shift());
+      }
+    }
+
+  }
+};
+
+/**
+ * Internal stream read method. This method reads message buffers.
+ * @param {number} size - This parameter is ignored for our cases.
+ * @private
+ */
+KafkaConsumerStream.prototype._read_buffer = function(size) {
+  if (this.messages.length > 0) {
+    return this.push(this.messages.shift());
+  }
+
+  if (!this.consumer) {
+    // This consumer is set to `null` in the close function
+    return;
+  }
+
+  if (!this.consumer.isConnected()) {
+    this.consumer.once('ready', function() {
+      // This is the way Node.js does it
+      // https://github.com/nodejs/node/blob/master/lib/fs.js#L1733
+      this._read(size);
+    }.bind(this));
+    return;
+  }
+
+  if (this.destroyed) {
+    return;
+  }
+
+  var self = this;
+
+  // If the size (number of messages) we are being advised to fetch is
+  // greater than or equal to the fetch size, use the fetch size.
+  // Only opt to use the size in case it is LESS than the fetch size.
+  // Essentially, we want to use the smaller value here
+  var fetchSize = size >= this.fetchSize ? this.fetchSize : size;
+
+  this.consumer.consume(fetchSize, onread);
+
+  // Retry function. Will wait up to the wait interval, with some
+  // random noise if one is provided. Otherwise, will go immediately.
+  function retry() {
+    if (!self.waitInterval) {
+      setImmediate(function() {
+        self._read(size);
+      });
+    } else {
+      setTimeout(function() {
+        self._read(size);
+      }, self.waitInterval * Math.random()).unref();
+    }
+  }
+
+  function onread(err, messages) {
+    // If there was an error we still want to emit it.
+    // Essentially, if the user does not register an error
+    // handler, it will still cause the stream to blow up.
+    //
+    // But... if one is provided, consumption will move on
+    // as normal
+    if (err) {
+      self.emit('error', err);
+    }
+
+    // If there are no messages it means we reached EOF or a timeout.
+    // Do what we used to do
+
+    if (err || messages.length < 1) {
+      // If we got an error or if there were no messages, initiate a retry
+      retry();
+      return;
+    } else {
+      if (self.streamAsBatch) {
+        self.push(messages);
+      } else {
+        for (var i = 0; i < messages.length; i++) {
+          self.messages.push(messages[i].value);
+        }
+        // Now that we have added them all the inner messages buffer,
+        // we can just push the most recent one
+        self.push(self.messages.shift());
+      }
+    }
+
+  }
+};
+
+KafkaConsumerStream.prototype.connect = function(options) {
+  var self = this;
+
+  function connectCallback(err, metadata) {
+    if (err) {
+      self.emit('error', err);
+      self.destroy();
+      return;
+    }
+
+    try {
+      // Subscribe to the topics as well so we will be ready
+      // If this throws the stream is invalid
+
+      // This is the magic part. If topics is a function, before we subscribe,
+      // pass the metadata in
+      if (typeof self.topics === 'function') {
+        var topics = self.topics(metadata);
+        self.consumer.subscribe(topics);
+      } else {
+        self.consumer.subscribe(self.topics);
+      }
+    } catch (e) {
+      self.emit('error', e);
+      self.destroy();
+      return;
+    }
+
+    // start the flow of data
+    self.read();
+  }
+
+  if (!this.consumer.isConnected()) {
+    self.consumer.connect(options, connectCallback);
+  } else {
+    // Immediately call the connect callback
+    setImmediate(function() {
+      connectCallback(null, self.consumer._metadata);
+    });
+  }
+
+};
+
+KafkaConsumerStream.prototype.destroy = function() {
+  if (this.destroyed) {
+    return;
+  }
+  this.destroyed = true;
+  this.close();
+};
+
+KafkaConsumerStream.prototype.close = function(cb) {
+  var self = this;
+  if (cb) {
+    this.once('close', cb);
+  }
+
+  if (!self.consumer._isConnecting && !self.consumer._isConnected) {
+    // If we aren't even connected just exit. We are done.
+    close();
+    return;
+  }
+
+  if (self.consumer._isConnecting) {
+    self.consumer.once('ready', function() {
+      // Don't pass the CB because it has already been passed.
+      self.close();
+    });
+    return;
+  }
+
+  if (self.consumer._isConnected) {
+    self.consumer.unsubscribe();
+    self.consumer.disconnect(function() {
+      close();
+    });
+  }
+
+  function close() {
+    self.emit('close');
+  }
+};