fluent-plugin-kafka 0.19.7 → 0.19.9

data/lib/fluent/plugin/out_kafka2.rb

@@ -98,6 +98,7 @@ DESC
     include Fluent::KafkaPluginUtil::AwsIamSettings
     include Fluent::KafkaPluginUtil::SSLSettings
     include Fluent::KafkaPluginUtil::SaslSettings
+    include Fluent::KafkaPluginUtil::PartitionSettings
 
     def initialize
       super
@@ -197,6 +198,10 @@ DESC
         end
       end
 
+      if @share_producer && @buffer_config.flush_thread_count > 1
+        log.warn "share_producer sends one chunk at a time even with flush_thread_count #{@buffer_config.flush_thread_count}. Set share_producer false to send chunks in parallel"
+      end
+
       @producer_opts = {max_retries: @max_send_retries, required_acks: @required_acks, idempotent: @idempotent}
       @producer_opts[:ack_timeout] = @ack_timeout if @ack_timeout
       @producer_opts[:compression_codec] = @compression_codec.to_sym if @compression_codec
@@ -239,6 +244,7 @@ DESC
 
       if @share_producer
         @shared_producer = create_producer
+        @shared_producer_mutex = Mutex.new
       else
         @producers = {}
         @producers_mutex = Mutex.new
@@ -273,8 +279,10 @@ DESC
 
     def shutdown_producers
       if @share_producer
-        @shared_producer.shutdown
-        @shared_producer = nil
+        @shared_producer_mutex.synchronize {
+          @shared_producer.shutdown if @shared_producer
+          @shared_producer = nil
+        }
       else
         @producers_mutex.synchronize {
           shutdown_threads = @producers.map { |key, producer|
@@ -292,7 +300,7 @@ DESC
 
     def get_producer
       if @share_producer
-        @shared_producer
+        @shared_producer ||= create_producer
       else
         @producers_mutex.synchronize {
           producer = @producers[Thread.current.object_id]
@@ -305,6 +313,17 @@ DESC
       end
     end
 
+    def discard_producer
+      if @share_producer
+        producer, @shared_producer = @shared_producer, nil
+      else
+        producer = @producers_mutex.synchronize { @producers.delete(Thread.current.object_id) }
+      end
+      producer.shutdown if producer
+    rescue => e
+      log.warn "failed to shut down the discarded producer:", :error => e.to_s, :error_class => e.class.to_s
+    end
+
     def setup_formatter(conf)
       type = conf['@type']
       case type
@@ -314,8 +333,8 @@ DESC
           Oj.default_options = Fluent::DEFAULT_OJ_OPTIONS
           Proc.new { |tag, time, record| Oj.dump(record) }
         rescue LoadError
-          require 'yajl'
-          Proc.new { |tag, time, record| Yajl::Encoder.encode(record) }
+          require 'json'
+          Proc.new { |tag, time, record| JSON.generate(record) }
         end
       when 'ltsv'
         require 'ltsv'
@@ -326,10 +345,20 @@ DESC
       end
     end
 
-    # TODO: optimize write performance
     def write(chunk)
       @writing_threads_mutex.synchronize { @writing_threads.add(Thread.current) }
 
+      if @share_producer
+        @shared_producer_mutex.synchronize { write_chunk(chunk) }
+      else
+        write_chunk(chunk)
+      end
+    ensure
+      @writing_threads_mutex.synchronize { @writing_threads.delete(Thread.current) }
+    end
+
+    # TODO: optimize write performance
+    def write_chunk(chunk)
       tag = chunk.metadata.tag
       topic = if @topic
                 extract_placeholders(@topic, chunk)
@@ -351,6 +380,7 @@ DESC
             partition_key = (@exclude_partition_key ? record.delete(@partition_key_key) : record[@partition_key_key]) || @default_partition_key
             partition = (@exclude_partition ? record.delete(@partition_key) : record[@partition_key]) || @default_partition
             message_key = (@exclude_message_key ? record.delete(@message_key_key) : record[@message_key_key]) || @default_message_key
+            partition = coerce_partition(partition) unless partition.nil?
 
             if mutate_headers
               headers = base_headers.clone
@@ -416,12 +446,11 @@ DESC
       log.warn "Exception ignored in tag : #{tag}" if ignore
       # For safety, refresh client and its producers
       refresh_client(false)
+      discard_producer
       # raise UnrecoverableError for backup ignored exception chunk
       raise Fluent::UnrecoverableError if ignore && exception_backup
       # Raise exception to retry sendind messages
       raise e unless ignore
-    ensure
-      @writing_threads_mutex.synchronize { @writing_threads.delete(Thread.current) }
     end
   end
 end

data/lib/fluent/plugin/out_kafka_buffered.rb

@@ -119,7 +119,7 @@ DESC
       @seed_brokers = []
       z = Zookeeper.new(@zookeeper)
       z.get_children(:path => @zookeeper_path)[:children].each do |id|
-        broker = Yajl.load(z.get(:path => @zookeeper_path + "/#{id}")[:data])
+        broker = JSON.parse(z.get(:path => @zookeeper_path + "/#{id}")[:data], allow_duplicate_key: true)
         if @ssl_client_cert
           @seed_brokers.push(pickup_ssl_endpoint(broker))
         else
@@ -259,8 +259,8 @@ DESC
         Oj.default_options = Fluent::DEFAULT_OJ_OPTIONS
         Proc.new { |tag, time, record| Oj.dump(record) }
       rescue LoadError
-        require 'yajl'
-        Proc.new { |tag, time, record| Yajl::Encoder.encode(record) }
+        require 'json'
+        Proc.new { |tag, time, record| JSON.generate(record) }
       end
     elsif @output_data_type == 'ltsv'
       require 'ltsv'

data/lib/fluent/plugin/out_rdkafka.rb

@@ -130,12 +130,14 @@ DESC
       :"bootstrap.servers" => @brokers,
     }
 
-    if @ssl_ca_cert && @ssl_ca_cert[0]
+    if (@ssl_ca_cert && @ssl_ca_cert[0]) || @ssl_ca_certs_from_system || @ssl_client_cert || @ssl_client_cert_key
       ssl = true
-      config[:"ssl.ca.location"] = @ssl_ca_cert[0]
+      config[:"ssl.ca.location"] = @ssl_ca_cert[0] if @ssl_ca_cert && @ssl_ca_cert[0]
       config[:"ssl.certificate.location"] = @ssl_client_cert if @ssl_client_cert
       config[:"ssl.key.location"] = @ssl_client_cert_key if @ssl_client_cert_key
       config[:"ssl.key.password"] = @ssl_client_cert_key_password if @ssl_client_cert_key_password
+      config[:"ssl.endpoint.identification.algorithm"] = @ssl_verify_hostname ? "https" : "none"
+      config[:"enable.ssl.certificate.verification"] = true
     end
 
     if @principal
@@ -228,8 +230,8 @@ DESC
         Oj.default_options = Fluent::DEFAULT_OJ_OPTIONS
         Proc.new { |tag, time, record| Oj.dump(record) }
       rescue LoadError
-        require 'yajl'
-        Proc.new { |tag, time, record| Yajl::Encoder.encode(record) }
+        require 'json'
+        Proc.new { |tag, time, record| JSON.generate(record) }
       end
     elsif @output_data_type == 'ltsv'
       require 'ltsv'

data/lib/fluent/plugin/out_rdkafka2.rb

@@ -7,14 +7,16 @@ require 'rdkafka'
 
 begin
   rdkafka_version = Gem::Version::create(Rdkafka::VERSION)
-  if rdkafka_version < Gem::Version.create('0.12.0')
-    require_relative 'rdkafka_patch/0_11_0'
-  elsif rdkafka_version == Gem::Version.create('0.12.0')
-    require_relative 'rdkafka_patch/0_12_0'
+  if rdkafka_version >= Gem::Version.create('0.16.0')
+    require_relative 'rdkafka_patch/0_16_0'
   elsif rdkafka_version >= Gem::Version.create('0.14.0')
     require_relative 'rdkafka_patch/0_14_0'
-  elsif rdkafka_version >= Gem::Version.create('0.16.0')
-    require_relative 'rdkafka_patch/0_16_0'
+  elsif rdkafka_version >= Gem::Version.create('0.13.0')
+    require_relative 'rdkafka_patch/0_13_0'
+  elsif rdkafka_version >= Gem::Version.create('0.12.0')
+    require_relative 'rdkafka_patch/0_12_0'
+  else
+    require_relative 'rdkafka_patch/0_11_0'
   end
 rescue LoadError, NameError
   raise "unable to patch rdkafka."
@@ -97,7 +99,6 @@ DESC
     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 :unrecoverable_error_codes, :array, :default => ["topic_authorization_failed", "msg_size_too_large"],
                  :desc => 'Handle some of the error codes should be unrecoverable if specified'
 
@@ -111,6 +112,8 @@ DESC
 
     include Fluent::KafkaPluginUtil::SSLSettings
     include Fluent::KafkaPluginUtil::SaslSettings
+    include Fluent::KafkaPluginUtil::PartitionSettings
+    include Fluent::KafkaPluginUtil::RdkafkaSecuritySettings
 
     class EnqueueRate
       class LimitExceeded < StandardError
@@ -250,39 +253,7 @@ DESC
     end
 
     def build_config
-      config = {:"bootstrap.servers" => @brokers}
-
-      if @ssl_ca_cert && @ssl_ca_cert[0]
-        ssl = true
-        config[:"ssl.ca.location"] = @ssl_ca_cert[0]
-        config[:"ssl.certificate.location"] = @ssl_client_cert if @ssl_client_cert
-        config[:"ssl.key.location"] = @ssl_client_cert_key if @ssl_client_cert_key
-        config[:"ssl.key.password"] = @ssl_client_cert_key_password if @ssl_client_cert_key_password
-      end
-
-      if @principal
-        sasl = true
-        config[:"sasl.mechanisms"] = "GSSAPI"
-        config[:"sasl.kerberos.principal"] = @principal
-        config[:"sasl.kerberos.service.name"] = @service_name if @service_name
-        config[:"sasl.kerberos.keytab"] = @keytab if @keytab
-      end
-
-      if @username && @password
-        sasl = true
-        config[:"sasl.mechanisms"] = 'PLAIN'
-      end
-
-      if ssl && sasl
-        security_protocol = "SASL_SSL"
-      elsif ssl && !sasl
-        security_protocol = "SSL"
-      elsif !ssl && sasl
-        security_protocol = "SASL_PLAINTEXT"
-      else
-        security_protocol = "PLAINTEXT"
-      end
-      config[:"security.protocol"] = security_protocol
+      config = {:"bootstrap.servers" => @brokers}.merge(rdkafka_security_config)
 
       config[:"compression.codec"] = @compression_codec if @compression_codec
       config[:"message.send.max.retries"] = @max_send_retries if @max_send_retries
@@ -292,14 +263,14 @@ DESC
       config[:"queue.buffering.max.messages"] = @rdkafka_buffering_max_messages if @rdkafka_buffering_max_messages
       config[:"message.max.bytes"] = @rdkafka_message_max_bytes if @rdkafka_message_max_bytes
       config[:"batch.num.messages"] = @rdkafka_message_max_num if @rdkafka_message_max_num
-      config[:"sasl.username"] = @username if @username
-      config[:"sasl.password"] = @password if @password
       config[:"enable.idempotence"] = @idempotent if @idempotent
 
       @rdkafka_options.each { |k, v|
         config[k.to_sym] = v
       }
 
+      validate_sasl_over_ssl(config)
+
       config
     end
 
@@ -406,6 +377,8 @@ DESC
             record.delete(@topic_key) if @exclude_topic_key
             partition = (@exclude_partition ? record.delete(@partition_key) : record[@partition_key]) || @default_partition
             message_key = (@exclude_message_key ? record.delete(@message_key_key) : record[@message_key_key]) || @default_message_key
+            partition = coerce_partition(partition) unless partition.nil?
+            message_key = message_key.to_s unless message_key.nil?
 
             @headers_from_record_accessors.each do |key, header_accessor|
               headers[key] = header_accessor.call(record)
@@ -451,6 +424,7 @@ DESC
           log.warn "Delivery failed and matched regexp pattern #{@discard_kafka_delivery_failed_regex}. Discard events:", :error => e.to_s, :error_class => e.class.to_s, :tag => tag
 	else
 	  log.warn "Send exception occurred: #{e} at #{e.backtrace.first}"
+          raise Fluent::UnrecoverableError, "Rejected due to #{e}" if unrecoverable_error?(e)
           # Raise exception to retry sendind messages
           raise e
 	end
@@ -507,15 +481,14 @@ DESC
 
             raise e
           else
-            if unrecoverable_error_codes.include?(e.code.to_s)
-              # some of the errors should be handled as an unrecoverable error
-              raise Fluent::UnrecoverableError, "Rejected due to #{e}"
-            else
-              raise e
-            end
+            raise e
           end
         end
       end
     end
+
+    def unrecoverable_error?(e)
+      e.is_a?(Rdkafka::RdkafkaError) && unrecoverable_error_codes.include?(e.code.to_s)
+    end
   end
 end

data/lib/fluent/plugin/rdkafka_patch/0_13_0.rb

@@ -0,0 +1,43 @@
+class Rdkafka::NativeKafka
+  # return false if producer is forcefully closed, otherwise return true
+  def close(timeout=nil, object_id=nil)
+    return true if closed?
+
+    synchronize do
+      # Indicate to the outside world that we are closing
+      @closing = true
+
+      thread_status = :unknown
+      if @polling_thread
+        # Indicate to polling thread that we're closing
+        @polling_thread[:closing] = true
+
+        # Wait for the polling thread to finish up,
+        # this can be aborted in practice if this
+        # code runs from a finalizer.
+        thread_status = @polling_thread.join(timeout)
+      end
+
+      # Destroy the client after locking both mutexes
+      @poll_mutex.lock
+
+      # This check prevents a race condition, where we would enter the close in two threads
+      # and after unlocking the primary one that hold the lock but finished, ours would be unlocked
+      # and would continue to run, trying to destroy inner twice
+      if @inner
+        Rdkafka::Bindings.rd_kafka_destroy(@inner)
+        @inner = nil
+      end
+
+      !thread_status.nil?
+    end
+  end
+end
+
+class Rdkafka::Producer
+  def close(timeout = nil)
+    return true if closed?
+    ObjectSpace.undefine_finalizer(self)
+    @native_kafka.close(timeout)
+  end
+end

data/lib/fluent/plugin/rdkafka_patch/0_16_0.rb

@@ -24,7 +24,7 @@ class Rdkafka::NativeKafka
       # This check prevents a race condition, where we would enter the close in two threads
       # and after unlocking the primary one that hold the lock but finished, ours would be unlocked
       # and would continue to run, trying to destroy inner twice
-      retun unless @inner
+      return unless @inner
 
       Rdkafka::Bindings.rd_kafka_destroy(@inner)
       @inner = nil
@@ -40,7 +40,7 @@ class Rdkafka::Producer
     return true if closed?
     ObjectSpace.undefine_finalizer(self)
 
-    @native_kafka.close(timeout) do
+    closed = @native_kafka.close(timeout) do
       # We need to remove the topics references objects before we destroy the producer,
       # otherwise they would leak out
       @topics_refs_map.each_value do |refs|
@@ -51,5 +51,6 @@ class Rdkafka::Producer
     end
 
     @topics_refs_map.clear
+    closed
   end
 end

data/test/plugin/test_in_kafka.rb

@@ -1,6 +1,7 @@
 require 'helper'
 require 'fluent/test/driver/input'
 require 'securerandom'
+require 'json'
 
 class KafkaInputTest < Test::Unit::TestCase
   def setup
@@ -34,6 +35,158 @@ class KafkaInputTest < Test::Unit::TestCase
     assert_false d.instance.multi_workers_ready?
   end
 
+  class TopicWatcherTest < self
+    FakeMessage = Struct.new(:value, :key, :offset, :create_time)
+
+    class FakeRouter
+      attr_reader :emitted
+
+      def initialize
+        @emitted = []
+      end
+
+      def emit_stream(tag, es)
+        records = []
+        es.each { |time, record| records << record }
+        @emitted << [tag, records]
+      end
+    end
+
+    class OverflowRouter < FakeRouter
+      def initialize(overflow_tag)
+        super()
+        @overflow_tag = overflow_tag
+      end
+
+      def emit_stream(tag, es)
+        raise Fluent::KafkaInput::BufferChunkOverflowError, "too large" if tag == @overflow_tag
+        super
+      end
+    end
+
+    def create_topic_watcher(messages, router, tag_source: :record, add_prefix: nil, add_suffix: nil)
+      kafka = Object.new
+      kafka.define_singleton_method(:fetch_messages) { |**args| messages }
+      parser = Proc.new { |msg, te| JSON.parse(msg.value) }
+
+      Fluent::KafkaInput::TopicWatcher.new(
+        Fluent::KafkaInput::TopicEntry.new(TOPIC_NAME, 0, 0),
+        kafka,
+        1,
+        parser,
+        add_prefix,
+        add_suffix,
+        nil,
+        router,
+        nil,
+        :now,
+        'time',
+        tag_source,
+        'tag')
+    end
+
+    def tagged_message(tag, message, offset)
+      FakeMessage.new({'tag' => tag, 'message' => message}.to_json, nil, offset, Time.now)
+    end
+
+    def test_consume_with_tag_source_record_emits_stream_per_tag
+      messages = [
+        tagged_message('app.trusted', 'record 1', 0),
+        tagged_message('app.trusted', 'record 2', 1),
+        tagged_message('attacker.controlled', 'record 3', 2),
+      ]
+      router = FakeRouter.new
+      create_topic_watcher(messages, router).consume
+
+      assert_equal([['app.trusted', ['record 1', 'record 2']],
+                    ['attacker.controlled', ['record 3']]],
+                   router.emitted.map { |tag, records| [tag, records.map { |r| r['message'] }] })
+    end
+
+    def test_consume_with_tag_source_record_applies_prefix_and_suffix_per_tag
+      messages = [
+        tagged_message('app.trusted', 'record 1', 0),
+        tagged_message('attacker.controlled', 'record 2', 1),
+      ]
+      router = FakeRouter.new
+      create_topic_watcher(messages, router, add_prefix: 'prefix', add_suffix: 'suffix').consume
+
+      assert_equal(['prefix.app.trusted.suffix', 'prefix.attacker.controlled.suffix'],
+                   router.emitted.map(&:first))
+    end
+
+    def test_consume_with_tag_source_record_skips_unparsable_message
+      messages = [
+        tagged_message('app.trusted', 'record 1', 0),
+        FakeMessage.new('this is not json', nil, 1, Time.now),
+        tagged_message('attacker.controlled', 'record 2', 2),
+      ]
+      router = FakeRouter.new
+      create_topic_watcher(messages, router).consume
+
+      assert_equal([['app.trusted', ['record 1']],
+                    ['attacker.controlled', ['record 2']]],
+                   router.emitted.map { |tag, records| [tag, records.map { |r| r['message'] }] })
+    end
+
+    def test_consume_with_tag_source_record_skips_invalid_tag
+      messages = [
+        tagged_message('app.trusted', 'record 1', 0),
+        tagged_message(7, 'record 2', 1),
+        FakeMessage.new({'message' => 'record 3'}.to_json, nil, 2, Time.now),
+        tagged_message('', 'record 4', 3),
+        tagged_message('attacker.controlled', 'record 5', 4),
+      ]
+      router = FakeRouter.new
+      watcher = create_topic_watcher(messages, router)
+      watcher.consume
+
+      assert_equal([['app.trusted', ['record 1']],
+                    ['attacker.controlled', ['record 5']]],
+                   router.emitted.map { |tag, records| [tag, records.map { |r| r['message'] }] })
+      assert_equal(5, watcher.instance_variable_get(:@next_offset))
+    end
+
+    def test_consume_advances_offset_when_every_message_is_skipped
+      messages = [
+        FakeMessage.new('this is not json', nil, 0, Time.now),
+        tagged_message(nil, 'record 1', 1),
+      ]
+      router = FakeRouter.new
+      watcher = create_topic_watcher(messages, router)
+      watcher.consume
+
+      assert_equal([], router.emitted)
+      assert_equal(2, watcher.instance_variable_get(:@next_offset))
+    end
+
+    def test_consume_with_tag_source_topic_emits_single_stream
+      messages = [
+        tagged_message('app.trusted', 'record 1', 0),
+        tagged_message('attacker.controlled', 'record 2', 1),
+      ]
+      router = FakeRouter.new
+      create_topic_watcher(messages, router, tag_source: :topic).consume
+
+      assert_equal([[TOPIC_NAME, ['record 1', 'record 2']]],
+                   router.emitted.map { |tag, records| [tag, records.map { |r| r['message'] }] })
+    end
+
+    def test_consume_skips_buffer_chunk_overflow
+      messages = [
+        tagged_message('app.large', 'record 1', 0),
+        tagged_message('app.small', 'record 2', 1),
+      ]
+      router = OverflowRouter.new('app.large')
+      watcher = create_topic_watcher(messages, router)
+      watcher.consume
+
+      assert_equal([['app.small', ['record 2']]],
+                   router.emitted.map { |tag, records| [tag, records.map { |r| r['message'] }] })
+      assert_equal(2, watcher.instance_variable_get(:@next_offset))
+    end
+  end
+
   class ConsumeTest < self
     def setup
       @kafka = Kafka.new(["localhost:9092"], client_id: 'kafka')

data/test/plugin/test_in_kafka_group.rb

@@ -36,6 +36,19 @@ class KafkaGroupInputTest < Test::Unit::TestCase
     assert_true d.instance.multi_workers_ready?
   end
 
+  def test_emit_events_with_buffer_chunk_overflow
+    d = create_driver
+    es = Fluent::MultiEventStream.new
+    router = Object.new
+    mock(router).emit_stream("test", es).once { raise Fluent::Plugin::Buffer::BufferChunkOverflowError, "too large" }
+    stub(d.instance).router { router }
+
+    assert_nothing_raised do
+      d.instance.emit_events("test", es)
+    end
+    assert_true d.logs.any? { |log| log.include?("Skipped records larger than the buffer chunk limit size") && log.include?("too large") }
+  end
+
   class ConsumeTest < self
     def setup
       @kafka = Kafka.new(["localhost:9092"], client_id: 'kafka')