fluent-plugin-kafka 0.16.2 → 0.17.2

data/lib/fluent/plugin/out_rdkafka.rb

@@ -48,7 +48,7 @@ Set true to remove partition from data
 DESC
    config_param :exclude_message_key, :bool, :default => false,
                :desc => <<-DESC
-Set true to remove partition key from data
+Set true to remove message key from data
 DESC
    config_param :exclude_topic_key, :bool, :default => false,
                 :desc => <<-DESC
@@ -65,6 +65,7 @@ DESC
 The codec the producer uses to compress messages.
 Supported codecs: (gzip|snappy)
 DESC
+  config_param :use_event_time, :bool, :default => false, :desc => 'Use fluentd event time for rdkafka timestamp'
   config_param :max_send_limit_bytes, :size, :default => nil
   config_param :rdkafka_buffering_max_ms, :integer, :default => nil
   config_param :rdkafka_buffering_max_messages, :integer, :default => nil
@@ -91,23 +92,29 @@ DESC
   def configure(conf)
     super
     log.instance_eval {
-      def add(level, &block)
-        return unless block
+      def add(level, message = nil)
+        if message.nil?
+          if block_given?
+            message = yield
+          else
+            return
+          end
+        end
 
         # Follow rdkakfa's log level. See also rdkafka-ruby's bindings.rb: https://github.com/appsignal/rdkafka-ruby/blob/e5c7261e3f2637554a5c12b924be297d7dca1328/lib/rdkafka/bindings.rb#L117
         case level
         when Logger::FATAL
-          self.fatal(block.call)
+          self.fatal(message)
         when Logger::ERROR
-          self.error(block.call)
+          self.error(message)
         when Logger::WARN
-          self.warn(block.call)
+          self.warn(message)
         when Logger::INFO
-          self.info(block.call)
+          self.info(message)
         when Logger::DEBUG
-          self.debug(block.call)
+          self.debug(message)
         else
-          self.trace(block.call)
+          self.trace(message)
         end
       end
     }
@@ -280,7 +287,7 @@ DESC
         end
 
         producer = get_producer
-        handler = enqueue_with_retry(producer, topic, record_buf, message_key, partition)
+        handler = enqueue_with_retry(producer, topic, record_buf, message_key, partition, time)
         handler
       }.each { |handler|
         handler.wait(max_wait_timeout: @rdkafka_delivery_handle_poll_timeout) if @rdkafka_delivery_handle_poll_timeout != 0
@@ -292,11 +299,11 @@ DESC
     raise e
   end
 
-  def enqueue_with_retry(producer, topic, record_buf, message_key, partition)
+  def enqueue_with_retry(producer, topic, record_buf, message_key, partition, time)
     attempt = 0
     loop do
       begin
-        handler = producer.produce(topic: topic, payload: record_buf, key: message_key, partition: partition)
+        handler = producer.produce(topic: topic, payload: record_buf, key: message_key, partition: partition, timestamp: @use_event_time ? Time.at(time) : nil)
         return handler
       rescue Exception => e
         if e.respond_to?(:code) && e.code == :queue_full

data/lib/fluent/plugin/out_rdkafka2.rb

@@ -56,6 +56,8 @@ DESC
                  :desc => <<-DESC
 Set true to remove topic key from data
 DESC
+    config_param :exclude_fields, :array, :default => [], value_type: :string,
+                 :desc => 'Fields to remove from data where the value is a jsonpath to a record value'
     config_param :headers, :hash, default: {}, symbolize_keys: true, value_type: :string,
                  :desc => 'Kafka message headers'
     config_param :headers_from_record, :hash, default: {}, symbolize_keys: true, value_type: :string,
@@ -72,6 +74,7 @@ DESC
 The codec the producer uses to compress messages. Used for compression.codec
 Supported codecs: (gzip|snappy)
 DESC
+    config_param :use_event_time, :bool, :default => false, :desc => 'Use fluentd event time for rdkafka timestamp'
     config_param :max_send_limit_bytes, :size, :default => nil
     config_param :discard_kafka_delivery_failed, :bool, :default => false
     config_param :rdkafka_buffering_max_ms, :integer, :default => nil, :desc => 'Used for queue.buffering.max.ms'
@@ -84,6 +87,7 @@ DESC
 
     config_param :max_enqueue_retries, :integer, :default => 3
     config_param :enqueue_retry_backoff, :integer, :default => 3
+    config_param :max_enqueue_bytes_per_second, :size, :default => nil, :desc => 'The maximum number of enqueueing bytes per second'
 
     config_param :service_name, :string, :default => nil, :desc => 'Used for sasl.kerberos.service.name'
     config_param :ssl_client_cert_key_password, :string, :default => nil, :desc => 'Used for ssl.key.password'
@@ -99,34 +103,96 @@ DESC
     include Fluent::KafkaPluginUtil::SSLSettings
     include Fluent::KafkaPluginUtil::SaslSettings
 
+    class EnqueueRate
+      class LimitExceeded < StandardError
+        attr_reader :next_retry_clock
+        def initialize(next_retry_clock)
+          @next_retry_clock = next_retry_clock
+        end
+      end
+
+      def initialize(limit_bytes_per_second)
+        @mutex = Mutex.new
+        @start_clock = Fluent::Clock.now
+        @bytes_per_second = 0
+        @limit_bytes_per_second = limit_bytes_per_second
+        @commits = {}
+      end
+
+      def raise_if_limit_exceeded(bytes_to_enqueue)
+        return if @limit_bytes_per_second.nil?
+
+        @mutex.synchronize do
+          @commits[Thread.current] = {
+            clock: Fluent::Clock.now,
+            bytesize: bytes_to_enqueue,
+          }
+
+          @bytes_per_second += @commits[Thread.current][:bytesize]
+          duration = @commits[Thread.current][:clock] - @start_clock
+
+          if duration < 1.0
+            if @bytes_per_second > @limit_bytes_per_second
+              raise LimitExceeded.new(@start_clock + 1.0)
+            end
+          else
+            @start_clock = @commits[Thread.current][:clock]
+            @bytes_per_second = @commits[Thread.current][:bytesize]
+          end
+        end
+      end
+
+      def revert
+        return if @limit_bytes_per_second.nil?
+
+        @mutex.synchronize do
+          return unless @commits[Thread.current]
+          return unless @commits[Thread.current][:clock]
+          if @commits[Thread.current][:clock] >= @start_clock
+            @bytes_per_second -= @commits[Thread.current][:bytesize]
+          end
+          @commits[Thread.current] = nil
+        end
+      end
+    end
+
     def initialize
       super
 
       @producers = nil
       @producers_mutex = nil
       @shared_producer = nil
+      @enqueue_rate = nil
+      @writing_threads_mutex = Mutex.new
+      @writing_threads = Set.new
     end
 
     def configure(conf)
       super
       log.instance_eval {
-        def add(level, &block)
-          return unless block
+        def add(level, message = nil)
+          if message.nil?
+            if block_given?
+              message = yield
+            else
+              return
+            end
+          end
 
           # Follow rdkakfa's log level. See also rdkafka-ruby's bindings.rb: https://github.com/appsignal/rdkafka-ruby/blob/e5c7261e3f2637554a5c12b924be297d7dca1328/lib/rdkafka/bindings.rb#L117
           case level
           when Logger::FATAL
-            self.fatal(block.call)
+            self.fatal(message)
           when Logger::ERROR
-            self.error(block.call)
+            self.error(message)
           when Logger::WARN
-            self.warn(block.call)
+            self.warn(message)
           when Logger::INFO
-            self.info(block.call)
+            self.info(message)
           when Logger::DEBUG
-            self.debug(block.call)
+            self.debug(message)
           else
-            self.trace(block.call)
+            self.trace(message)
           end
         end
       }
@@ -158,6 +224,12 @@ DESC
       @headers_from_record.each do |key, value|
         @headers_from_record_accessors[key] = record_accessor_create(value)
       end
+
+      @exclude_field_accessors = @exclude_fields.map do |field|
+        record_accessor_create(field)
+      end
+
+      @enqueue_rate = EnqueueRate.new(@max_enqueue_bytes_per_second) unless @max_enqueue_bytes_per_second.nil?
     end
 
     def build_config
@@ -221,8 +293,19 @@ DESC
       true
     end
 
+    def wait_writing_threads
+      done = false
+      until done do
+        @writing_threads_mutex.synchronize do
+          done = true if @writing_threads.empty?
+        end
+        sleep(1) unless done
+      end
+    end
+
     def shutdown
       super
+      wait_writing_threads
       shutdown_producers
     end
 
@@ -279,6 +362,7 @@ DESC
     end
 
     def write(chunk)
+      @writing_threads_mutex.synchronize { @writing_threads.add(Thread.current) }
       tag = chunk.metadata.tag
       topic = if @topic
                 extract_placeholders(@topic, chunk)
@@ -305,6 +389,12 @@ DESC
               headers[key] = header_accessor.call(record)
             end
 
+            unless @exclude_fields.empty?
+              @exclude_field_accessors.each do |exclude_field_acessor|
+                exclude_field_acessor.delete(record)
+              end
+            end
+
             record_buf = @formatter_proc.call(tag, time, record)
             record_buf_bytes = record_buf.bytesize
             if @max_send_limit_bytes && record_buf_bytes > @max_send_limit_bytes
@@ -316,7 +406,7 @@ DESC
             next
           end
 
-          handler = enqueue_with_retry(producer, topic, record_buf, message_key, partition, headers)
+          handler = enqueue_with_retry(producer, topic, record_buf, message_key, partition, headers, time)
           if @rdkafka_delivery_handle_poll_timeout != 0
             handlers << handler
           end
@@ -333,14 +423,22 @@ DESC
         # Raise exception to retry sendind messages
         raise e
       end
+    ensure
+      @writing_threads_mutex.synchronize { @writing_threads.delete(Thread.current) }
     end
 
-    def enqueue_with_retry(producer, topic, record_buf, message_key, partition, headers)
+    def enqueue_with_retry(producer, topic, record_buf, message_key, partition, headers, time)
       attempt = 0
       loop do
         begin
-          return producer.produce(topic: topic, payload: record_buf, key: message_key, partition: partition, headers: headers)
+          @enqueue_rate.raise_if_limit_exceeded(record_buf.bytesize) if @enqueue_rate
+          return producer.produce(topic: topic, payload: record_buf, key: message_key, partition: partition, headers: headers, timestamp: @use_event_time ? Time.at(time) : nil)
+        rescue EnqueueRate::LimitExceeded => e
+          @enqueue_rate.revert if @enqueue_rate
+          duration = e.next_retry_clock - Fluent::Clock.now
+          sleep(duration) if duration > 0.0
         rescue Exception => e
+          @enqueue_rate.revert if @enqueue_rate
           if e.respond_to?(:code) && e.code == :queue_full
             if attempt <= @max_enqueue_retries
               log.warn "Failed to enqueue message; attempting retry #{attempt} of #{@max_enqueue_retries} after #{@enqueue_retry_backoff}s"

data/test/helper.rb

@@ -8,6 +8,7 @@ rescue Bundler::BundlerError => e
   exit e.status_code
 end
 require 'test/unit'
+require 'test/unit/rr'
 
 $LOAD_PATH.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 $LOAD_PATH.unshift(File.dirname(__FILE__))
@@ -22,6 +23,12 @@ unless ENV.has_key?('VERBOSE')
 end
 
 require 'fluent/plugin/out_kafka'
+require 'fluent/plugin/out_kafka_buffered'
+require 'fluent/plugin/out_kafka2'
+require 'fluent/plugin/in_kafka'
+require 'fluent/plugin/in_kafka_group'
+
+require "fluent/test/driver/output"
 
 class Test::Unit::TestCase
 end

data/test/plugin/test_in_kafka.rb

@@ -0,0 +1,66 @@
+require 'helper'
+require 'fluent/test/driver/input'
+require 'securerandom'
+
+class KafkaInputTest < Test::Unit::TestCase
+  def setup
+    Fluent::Test.setup
+  end
+
+  TOPIC_NAME = "kafka-input-#{SecureRandom.uuid}"
+
+  CONFIG = %[
+    @type kafka
+    brokers localhost:9092
+    format text
+    @label @kafka
+    topics #{TOPIC_NAME}
+  ]
+
+  def create_driver(conf = CONFIG)
+    Fluent::Test::Driver::Input.new(Fluent::KafkaInput).configure(conf)
+  end
+
+
+  def test_configure
+    d = create_driver
+    assert_equal TOPIC_NAME, d.instance.topics
+    assert_equal 'text', d.instance.format
+    assert_equal 'localhost:9092', d.instance.brokers
+  end
+
+  def test_multi_worker_support
+    d = create_driver
+    assert_false d.instance.multi_workers_ready?
+  end
+
+  class ConsumeTest < self
+    def setup
+      @kafka = Kafka.new(["localhost:9092"], client_id: 'kafka')
+      @producer = @kafka.producer
+    end
+
+    def teardown
+      @kafka.delete_topic(TOPIC_NAME)
+      @kafka.close
+    end
+
+    def test_consume
+      conf = %[
+        @type kafka
+        brokers localhost:9092
+        format text
+        @label @kafka
+        topics #{TOPIC_NAME}
+      ]
+      d = create_driver
+
+      d.run(expect_records: 1, timeout: 10) do
+        @producer.produce("Hello, fluent-plugin-kafka!", topic: TOPIC_NAME)
+        @producer.deliver_messages
+      end
+      expected = {'message'  => 'Hello, fluent-plugin-kafka!'}
+      assert_equal expected, d.events[0][2]
+    end
+  end
+end

data/test/plugin/test_in_kafka_group.rb

@@ -0,0 +1,67 @@
+require 'helper'
+require 'fluent/test/driver/input'
+require 'securerandom'
+
+class KafkaGroupInputTest < Test::Unit::TestCase
+  def setup
+    Fluent::Test.setup
+  end
+
+  TOPIC_NAME = "kafka-input-#{SecureRandom.uuid}"
+
+  CONFIG = %[
+    @type kafka
+    brokers localhost:9092
+    consumer_group fluentd
+    format text
+    @label @kafka
+    topics #{TOPIC_NAME}
+  ]
+
+  def create_driver(conf = CONFIG)
+    Fluent::Test::Driver::Input.new(Fluent::KafkaGroupInput).configure(conf)
+  end
+
+
+  def test_configure
+    d = create_driver
+    assert_equal [TOPIC_NAME], d.instance.topics
+    assert_equal 'text', d.instance.format
+    assert_equal 'localhost:9092', d.instance.brokers
+  end
+
+  def test_multi_worker_support
+    d = create_driver
+    assert_true d.instance.multi_workers_ready?
+  end
+
+  class ConsumeTest < self
+    def setup
+      @kafka = Kafka.new(["localhost:9092"], client_id: 'kafka')
+      @producer = @kafka.producer
+    end
+
+    def teardown
+      @kafka.delete_topic(TOPIC_NAME)
+      @kafka.close
+    end
+
+    def test_consume
+      conf = %[
+        @type kafka
+        brokers localhost:9092
+        format text
+        @label @kafka
+        topics #{TOPIC_NAME}
+      ]
+      d = create_driver
+
+      d.run(expect_records: 1, timeout: 10) do
+        @producer.produce("Hello, fluent-plugin-kafka!", topic: TOPIC_NAME)
+        @producer.deliver_messages
+      end
+      expected = {'message'  => 'Hello, fluent-plugin-kafka!'}
+      assert_equal expected, d.events[0][2]
+    end
+  end
+end

data/test/plugin/test_kafka_plugin_util.rb

@@ -1,14 +1,8 @@
 require 'helper'
 require 'fluent/plugin/kafka_plugin_util'
 
-class File
-    def File::read(path)
-        path
-    end
-end
-
 class KafkaPluginUtilTest < Test::Unit::TestCase
-    
+
     def self.config_param(name, type, options)
     end
     include Fluent::KafkaPluginUtil::SSLSettings
@@ -20,19 +14,31 @@ class KafkaPluginUtilTest < Test::Unit::TestCase
     end
 
     def test_read_ssl_file_when_nil
-        assert_equal(nil, read_ssl_file(nil))
+      stub(File).read(anything) do |path|
+        path
+      end
+      assert_equal(nil, read_ssl_file(nil))
     end
 
     def test_read_ssl_file_when_empty_string
-        assert_equal(nil, read_ssl_file(""))
+      stub(File).read(anything) do |path|
+        path
+      end
+      assert_equal(nil, read_ssl_file(""))
     end
 
     def test_read_ssl_file_when_non_empty_path
-        assert_equal("path", read_ssl_file("path"))
+      stub(File).read(anything) do |path|
+        path
+      end
+      assert_equal("path", read_ssl_file("path"))
     end
 
     def test_read_ssl_file_when_non_empty_array
-        assert_equal(["a","b"], read_ssl_file(["a","b"]))
+      stub(File).read(anything) do |path|
+        path
+      end
+      assert_equal(["a","b"], read_ssl_file(["a","b"]))
     end
 
-end
+end

data/test/plugin/test_out_kafka.rb

@@ -43,6 +43,16 @@ class KafkaOutputTest < Test::Unit::TestCase
     d = create_driver
   end
 
+  data("crc32" => "crc32",
+      "murmur2" => "murmur2")
+  def test_partitioner_hash_function(data)
+    hash_type = data
+    d = create_driver(CONFIG + %[partitioner_hash_function #{hash_type}])
+    assert_nothing_raised do
+      d.instance.refresh_client
+    end
+  end
+
   def test_mutli_worker_support
     d = create_driver
     assert_equal true, d.instance.multi_workers_ready?

data/test/plugin/test_out_kafka2.rb

@@ -0,0 +1,116 @@
+require 'helper'
+require 'fluent/test/helpers'
+require 'fluent/test/driver/input'
+require 'fluent/test/driver/output'
+require 'securerandom'
+
+class Kafka2OutputTest < Test::Unit::TestCase
+  include Fluent::Test::Helpers
+
+  def setup
+    Fluent::Test.setup
+  end
+
+  def base_config
+    config_element('ROOT', '', {"@type" => "kafka2"}, [
+                     config_element('format', "", {"@type" => "json"})
+                   ])
+  end
+
+  def config(default_topic: "kitagawakeiko")
+    base_config + config_element('ROOT', '', {"default_topic" => default_topic,
+                                              "brokers" => "localhost:9092"}, [
+                                 ])
+  end
+
+  def create_driver(conf = config, tag='test')
+    Fluent::Test::Driver::Output.new(Fluent::Kafka2Output).configure(conf)
+  end
+
+  def test_configure
+    assert_nothing_raised(Fluent::ConfigError) {
+      create_driver(base_config)
+    }
+
+    assert_nothing_raised(Fluent::ConfigError) {
+      create_driver(config)
+    }
+
+    assert_nothing_raised(Fluent::ConfigError) {
+      create_driver(config + config_element('buffer', "", {"@type" => "memory"}))
+    }
+
+    d = create_driver
+    assert_equal 'kitagawakeiko', d.instance.default_topic
+    assert_equal ['localhost:9092'], d.instance.brokers
+  end
+
+  data("crc32" => "crc32",
+       "murmur2" => "murmur2")
+  def test_partitioner_hash_function(data)
+    hash_type = data
+    d = create_driver(config + config_element('ROOT', '', {"partitioner_hash_function" => hash_type}))
+    assert_nothing_raised do
+      d.instance.refresh_client
+    end
+  end
+
+  def test_mutli_worker_support
+    d = create_driver
+    assert_equal true, d.instance.multi_workers_ready?
+  end
+
+  class WriteTest < self
+    TOPIC_NAME = "kafka-output-#{SecureRandom.uuid}"
+
+    INPUT_CONFIG = %[
+      @type kafka
+      brokers localhost:9092
+      format json
+      @label @kafka
+      topics #{TOPIC_NAME}
+    ]
+
+    def create_target_driver(conf = INPUT_CONFIG)
+      Fluent::Test::Driver::Input.new(Fluent::KafkaInput).configure(conf)
+    end
+
+    def setup
+      @kafka = Kafka.new(["localhost:9092"], client_id: 'kafka')
+    end
+
+    def teardown
+      @kafka.delete_topic(TOPIC_NAME)
+      @kafka.close
+    end
+
+    def test_write
+      target_driver = create_target_driver
+      expected_message = {"a" => 2}
+      target_driver.run(expect_records: 1, timeout: 5) do
+        sleep 2
+        d = create_driver(config(default_topic: TOPIC_NAME))
+        d.run do
+          d.feed("test", event_time, expected_message)
+        end
+      end
+      actual_messages = target_driver.events.collect { |event| event[2] }
+      assert_equal([expected_message], actual_messages)
+    end
+
+    def test_exclude_fields
+      conf = config(default_topic: TOPIC_NAME) +
+             config_element('ROOT', '', {"exclude_fields" => "$.foo"}, [])
+      target_driver = create_target_driver
+      target_driver.run(expect_records: 1, timeout: 5) do
+        sleep 2
+        d = create_driver(conf)
+        d.run do
+          d.feed('test', event_time, {'a' => 'b', 'foo' => 'bar', 'message' => 'test'})
+        end
+      end
+      actual_messages = target_driver.events.collect { |event| event[2] }
+      assert_equal([{'a' => 'b', 'message' => 'test'}], actual_messages)
+    end
+  end
+end