fluent-plugin-kafka-enchanced 0.5.1

data/lib/fluent/plugin/in_kafka_group.rb

@@ -0,0 +1,218 @@
+require 'fluent/input'
+require 'fluent/time'
+require 'fluent/plugin/kafka_plugin_util'
+
+class Fluent::KafkaGroupInput < Fluent::Input
+  Fluent::Plugin.register_input('kafka_group', self)
+
+  config_param :brokers, :string, :default => 'localhost:9092',
+               :desc => "List of broker-host:port, separate with comma, must set."
+  config_param :consumer_group, :string,
+               :desc => "Consumer group name, must set."
+  config_param :topics, :string,
+               :desc => "Listening topics(separate with comma',')."
+  config_param :format, :string, :default => 'json',
+               :desc => "Supported format: (json|text|ltsv|msgpack)"
+  config_param :message_key, :string, :default => 'message',
+               :desc => "For 'text' format only."
+  config_param :add_prefix, :string, :default => nil,
+               :desc => "Tag prefix (Optional)"
+  config_param :add_suffix, :string, :default => nil,
+               :desc => "Tag suffix (Optional)"
+  config_param :retry_emit_limit, :integer, :default => nil,
+               :desc => "How long to stop event consuming when BufferQueueLimitError happens. Wait retry_emit_limit x 1s. The default is waiting until BufferQueueLimitError is resolved"
+  config_param :use_record_time, :bool, :default => false,
+               :desc => "Replace message timestamp with contents of 'time' field."
+  config_param :time_format, :string, :default => nil,
+               :desc => "Time format to be used to parse 'time' filed."
+
+  # Kafka consumer options
+  config_param :max_bytes, :integer, :default => 1048576,
+               :desc => "Maximum number of bytes to fetch."
+  config_param :max_wait_time, :integer, :default => nil,
+               :desc => "How long to block until the server sends us data."
+  config_param :min_bytes, :integer, :default => nil,
+               :desc => "Smallest amount of data the server should send us."
+  config_param :session_timeout, :integer, :default => nil,
+               :desc => "The number of seconds after which, if a client hasn't contacted the Kafka cluster"
+  config_param :offset_commit_interval, :integer, :default => nil,
+               :desc => "The interval between offset commits, in seconds"
+  config_param :offset_commit_threshold, :integer, :default => nil,
+               :desc => "The number of messages that can be processed before their offsets are committed"
+  config_param :start_from_beginning, :bool, :default => true,
+               :desc => "Whether to start from the beginning of the topic or just subscribe to new messages being produced"
+
+  include Fluent::KafkaPluginUtil::SSLSettings
+
+  class ForShutdown < StandardError
+  end
+
+  BufferError = if defined?(Fluent::Plugin::Buffer::BufferOverflowError)
+                  Fluent::Plugin::Buffer::BufferOverflowError
+                else
+                  Fluent::BufferQueueLimitError
+                end
+
+  unless method_defined?(:router)
+    define_method("router") { Fluent::Engine }
+  end
+
+  def initialize
+    super
+    require 'kafka'
+
+    @time_parser = nil
+  end
+
+  def _config_to_array(config)
+    config_array = config.split(',').map {|k| k.strip }
+    if config_array.empty?
+      raise Fluent::ConfigError, "kafka_group: '#{config}' is a required parameter"
+    end
+    config_array
+  end
+
+  private :_config_to_array
+
+  def configure(conf)
+    super
+
+    $log.info "Will watch for topics #{@topics} at brokers " \
+              "#{@brokers} and '#{@consumer_group}' group"
+
+    @topics = _config_to_array(@topics)
+
+    if conf['max_wait_ms']
+      log.warn "'max_wait_ms' parameter is deprecated. Use second unit 'max_wait_time' instead"
+      @max_wait_time = conf['max_wait_ms'].to_i / 1000
+    end
+
+    @parser_proc = setup_parser
+
+    @consumer_opts = {:group_id => @consumer_group}
+    @consumer_opts[:session_timeout] = @session_timeout if @session_timeout
+    @consumer_opts[:offset_commit_interval] = @offset_commit_interval if @offset_commit_interval
+    @consumer_opts[:offset_commit_threshold] = @offset_commit_threshold if @offset_commit_threshold
+
+    @fetch_opts = {}
+    @fetch_opts[:max_wait_time] = @max_wait_time if @max_wait_time
+    @fetch_opts[:min_bytes] = @min_bytes if @min_bytes
+
+    if @use_record_time and @time_format
+      @time_parser = Fluent::TextParser::TimeParser.new(@time_format)
+    end
+  end
+
+  def setup_parser
+    case @format
+    when 'json'
+      require 'yajl'
+      Proc.new { |msg| Yajl::Parser.parse(msg.value) }
+    when 'ltsv'
+      require 'ltsv'
+      Proc.new { |msg| LTSV.parse(msg.value).first }
+    when 'msgpack'
+      require 'msgpack'
+      Proc.new { |msg| MessagePack.unpack(msg.value) }
+    when 'text'
+      Proc.new { |msg| {@message_key => msg.value} }
+    end
+  end
+
+  def start
+    super
+
+    @kafka = Kafka.new(seed_brokers: @brokers,
+                       ssl_ca_cert: read_ssl_file(@ssl_ca_cert),
+                       ssl_client_cert: read_ssl_file(@ssl_client_cert),
+                       ssl_client_cert_key: read_ssl_file(@ssl_client_cert_key))
+    @consumer = setup_consumer
+    @thread = Thread.new(&method(:run))
+  end
+
+  def shutdown
+    # This nil assignment should be guarded by mutex in multithread programming manner.
+    # But the situation is very low contention, so we don't use mutex for now.
+    # If the problem happens, we will add a guard for consumer.
+    consumer = @consumer
+    @consumer = nil
+    consumer.stop
+
+    @thread.join
+    @kafka.close
+    super
+  end
+
+  def setup_consumer
+    consumer = @kafka.consumer(@consumer_opts)
+    @topics.each { |topic|
+      consumer.subscribe(topic, start_from_beginning: @start_from_beginning, max_bytes_per_partition: @max_bytes)
+    }
+    consumer
+  end
+
+  def run
+    while @consumer
+      begin
+        @consumer.each_batch(@fetch_opts) { |batch|
+          es = Fluent::MultiEventStream.new
+          tag = batch.topic
+          tag = @add_prefix + "." + tag if @add_prefix
+          tag = tag + "." + @add_suffix if @add_suffix
+
+          batch.messages.each { |msg|
+            begin
+              record = @parser_proc.call(msg)
+              if @use_record_time
+                if @time_format
+                  record_time = @time_parser.parse(record['time'])
+                else
+                  record_time = record['time']
+                end
+              else
+                record_time = Fluent::Engine.now
+              end
+              es.add(record_time, record)
+            rescue => e
+              log.warn "parser error in #{batch.topic}/#{batch.partition}", :error => e.to_s, :value => msg.value, :offset => msg.offset
+              log.debug_backtrace
+            end
+          }
+
+          unless es.empty?
+            emit_events(tag, es)
+          end
+        }
+      rescue ForShutdown
+      rescue => e
+        log.error "unexpected error during consuming events from kafka. Re-fetch events.", :error => e.to_s
+        log.error_backtrace
+      end
+    end
+  rescue => e
+    log.error "unexpected error during consumer object access", :error => e.to_s
+    log.error_backtrace
+  end
+
+  def emit_events(tag, es)
+    retries = 0
+    begin
+      router.emit_stream(tag, es)
+    rescue BufferError
+      raise ForShutdown if @consumer.nil?
+
+      if @retry_emit_limit.nil?
+        sleep 1
+        retry
+      end
+
+      if retries < @retry_emit_limit
+        retries += 1
+        sleep 1
+        retry
+      else
+        raise RuntimeError, "Exceeds retry_emit_limit"
+      end
+    end
+  end
+end

data/lib/fluent/plugin/kafka_plugin_util.rb

@@ -0,0 +1,22 @@
+module Fluent
+  module KafkaPluginUtil
+    module SSLSettings
+      def self.included(klass)
+        klass.instance_eval {
+          # https://github.com/zendesk/ruby-kafka#encryption-and-authentication-using-ssl
+          config_param :ssl_ca_cert, :string, :default => nil,
+                       :desc => "a PEM encoded CA cert to use with and SSL connection."
+          config_param :ssl_client_cert, :string, :default => nil,
+                       :desc => "a PEM encoded client cert to use with and SSL connection. Must be used in combination with ssl_client_cert_key."
+          config_param :ssl_client_cert_key, :string, :default => nil,
+                       :desc => "a PEM encoded client cert key to use with and SSL connection. Must be used in combination with ssl_client_cert."
+        }
+      end
+
+      def read_ssl_file(path)
+        return nil if path.nil?
+        File.read(path)
+      end
+    end
+  end
+end

data/lib/fluent/plugin/kafka_producer_ext.rb

@@ -0,0 +1,225 @@
+require "set"
+require "kafka/partitioner"
+require "kafka/message_buffer"
+require "kafka/produce_operation"
+require "kafka/pending_message_queue"
+require "kafka/pending_message"
+require "kafka/compressor"
+require 'kafka/producer'
+
+# for out_kafka_buffered
+module Kafka
+  class Producer
+    def produce2(value, key: nil, topic:, partition: nil, partition_key: nil)
+      create_time = Time.now
+
+      message = PendingMessage.new(
+        value,
+        key,
+        topic,
+        partition,
+        partition_key,
+        create_time,
+        key.to_s.bytesize + value.to_s.bytesize
+      )
+
+      @target_topics.add(topic)
+      @pending_message_queue.write(message)
+
+      nil
+    end
+  end
+end
+
+# for out_kafka2
+module Kafka
+  class Client
+    def topic_producer(topic, compression_codec: nil, compression_threshold: 1, ack_timeout: 5, required_acks: :all, max_retries: 2, retry_backoff: 1, max_buffer_size: 1000, max_buffer_bytesize: 10_000_000)
+      compressor = Compressor.new(
+        codec_name: compression_codec,
+        threshold: compression_threshold,
+        instrumenter: @instrumenter,
+      )
+
+      TopicProducer.new(topic,
+        cluster: initialize_cluster,
+        logger: @logger,
+        instrumenter: @instrumenter,
+        compressor: compressor,
+        ack_timeout: ack_timeout,
+        required_acks: required_acks,
+        max_retries: max_retries,
+        retry_backoff: retry_backoff,
+        max_buffer_size: max_buffer_size,
+        max_buffer_bytesize: max_buffer_bytesize,
+      )
+    end
+  end
+
+  class TopicProducer
+    def initialize(topic, cluster:, logger:, instrumenter:, compressor:, ack_timeout:, required_acks:, max_retries:, retry_backoff:, max_buffer_size:, max_buffer_bytesize:)
+      @cluster = cluster
+      @logger = logger
+      @instrumenter = instrumenter
+      @required_acks = required_acks == :all ? -1 : required_acks
+      @ack_timeout = ack_timeout
+      @max_retries = max_retries
+      @retry_backoff = retry_backoff
+      @max_buffer_size = max_buffer_size
+      @max_buffer_bytesize = max_buffer_bytesize
+      @compressor = compressor
+
+      @topic = topic
+      @cluster.add_target_topics(Set.new([topic]))
+
+      # A buffer organized by topic/partition.
+      @buffer = MessageBuffer.new
+
+      # Messages added by `#produce` but not yet assigned a partition.
+      @pending_message_queue = PendingMessageQueue.new
+    end
+
+    def produce(value, key, partition, partition_key)
+      create_time = Time.now
+
+      message = PendingMessage.new(
+        value,
+        key,
+        @topic,
+        partition,
+        partition_key,
+        create_time,
+        key.to_s.bytesize + value.to_s.bytesize
+      )
+
+      @pending_message_queue.write(message)
+
+      nil
+    end
+
+    def deliver_messages
+      # There's no need to do anything if the buffer is empty.
+      return if buffer_size == 0
+
+      deliver_messages_with_retries
+    end
+
+    # Returns the number of messages currently held in the buffer.
+    #
+    # @return [Integer] buffer size.
+    def buffer_size
+      @pending_message_queue.size + @buffer.size
+    end
+
+    def buffer_bytesize
+      @pending_message_queue.bytesize + @buffer.bytesize
+    end
+
+    # Deletes all buffered messages.
+    #
+    # @return [nil]
+    def clear_buffer
+      @buffer.clear
+      @pending_message_queue.clear
+    end
+
+    # Closes all connections to the brokers.
+    #
+    # @return [nil]
+    def shutdown
+      @cluster.disconnect
+    end
+
+    private
+
+    def deliver_messages_with_retries
+      attempt = 0
+
+      #@cluster.add_target_topics(@target_topics)
+
+      operation = ProduceOperation.new(
+        cluster: @cluster,
+        buffer: @buffer,
+        required_acks: @required_acks,
+        ack_timeout: @ack_timeout,
+        compressor: @compressor,
+        logger: @logger,
+        instrumenter: @instrumenter,
+      )
+
+      loop do
+        attempt += 1
+
+        @cluster.refresh_metadata_if_necessary!
+
+        assign_partitions!
+        operation.execute
+
+        if @required_acks.zero?
+          # No response is returned by the brokers, so we can't know which messages
+          # have been successfully written. Our only option is to assume that they all
+          # have.
+          @buffer.clear
+        end
+
+        if buffer_size.zero?
+          break
+        elsif attempt <= @max_retries
+          @logger.warn "Failed to send all messages; attempting retry #{attempt} of #{@max_retries} after #{@retry_backoff}s"
+
+          sleep @retry_backoff
+        else
+          @logger.error "Failed to send all messages; keeping remaining messages in buffer"
+          break
+        end
+      end
+
+      unless @pending_message_queue.empty?
+        # Mark the cluster as stale in order to force a cluster metadata refresh.
+        @cluster.mark_as_stale!
+        raise DeliveryFailed, "Failed to assign partitions to #{@pending_message_queue.size} messages"
+      end
+
+      unless @buffer.empty?
+        partitions = @buffer.map {|topic, partition, _| "#{topic}/#{partition}" }.join(", ")
+
+        raise DeliveryFailed, "Failed to send messages to #{partitions}"
+      end
+    end
+
+    def assign_partitions!
+      failed_messages = []
+      partition_count = @cluster.partitions_for(@topic).count
+
+      @pending_message_queue.each do |message|
+        partition = message.partition
+
+        begin
+          if partition.nil?
+            partition = Partitioner.partition_for_key(partition_count, message)
+          end
+
+          @buffer.write(
+            value: message.value,
+            key: message.key,
+            topic: message.topic,
+            partition: partition,
+            create_time: message.create_time,
+          )
+        rescue Kafka::Error => e
+          failed_messages << message
+        end
+      end
+
+      if failed_messages.any?
+        failed_messages.group_by(&:topic).each do |topic, messages|
+          @logger.error "Failed to assign partitions to #{messages.count} messages in #{topic}"
+        end
+
+        @cluster.mark_as_stale!
+      end
+
+      @pending_message_queue.replace(failed_messages)
+    end
+  end
+end