fluent-plugin-kinesis 1.3.0 → 2.0.0

data/Rakefile

@@ -1,16 +1,16 @@
 #
-#  Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
-#  Licensed under the Amazon Software License (the "License").
-#  You may not use this file except in compliance with the License.
-#  A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
 #
-#  http://aws.amazon.com/asl/
+#     http://www.apache.org/licenses/LICENSE-2.0
 #
-#  or in the "license" file accompanying this file. This file is distributed
-#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-#  express or implied. See the License for the specific language governing
-#  permissions and limitations under the License.
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 
 require "bundler/gem_tasks"
 
@@ -23,8 +23,4 @@ Rake::TestTask.new do |test|
   test.options = '-v'
 end
 
-load 'kinesis_producer/tasks/binary.rake'
-Rake::Task[:build].enhance [:binaries]
-Rake::Task[:test].enhance [:binaries]
-
 load 'benchmark/task.rake'

data/benchmark/task.rake

@@ -1,72 +1,110 @@
 #
-#  Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
-#  Licensed under the Amazon Software License (the "License").
-#  You may not use this file except in compliance with the License.
-#  A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
 #
-#  http://aws.amazon.com/asl/
+#     http://www.apache.org/licenses/LICENSE-2.0
 #
-#  or in the "license" file accompanying this file. This file is distributed
-#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-#  express or implied. See the License for the specific language governing
-#  permissions and limitations under the License.
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 
-require_relative '../test/dummy_server'
+require_relative '../test/helper'
+require 'fluent/plugin/out_kinesis_streams'
+require 'fluent/plugin/out_kinesis_streams_aggregated'
+require 'fluent/plugin/out_kinesis_firehose'
+require 'benchmark'
+require 'net/empty_port'
 
-task :benchmark do
-  server = DummyServer.start
-  conf = profile_conf(ENV["TYPE"] || 'streams', ENV["RATE"] || 1000, server)
-  pid = spawn("bundle exec fluentd -i '#{conf}' -c benchmark/dummy.conf")
-  sleep 10
-  Process.kill("TERM", pid)
-  Process.wait
-  puts "Results: requets: #{server.requests.size}, raw_records: #{server.raw_records.size}, records: #{server.records.size}"
-  server.shutdown
+namespace :benchmark do
+  task :local do
+    KinesisBenchmark.new.run
+  end
+  task :remote do
+    KinesisBenchmark.new(false).run
+  end
 end
 
-def profile_conf(type, rate, server)
-  additional_conf = case type
-                  when 'streams', 'firehose'
-                    <<-EOS
-  endpoint https://localhost:#{server.port}
-  ssl_verify_peer false
-                    EOS
-                  when 'producer'
-                    <<-EOS
-  debug true
-  <kinesis_producer>
-    custom_endpoint localhost
-    port #{server.port}
-    verify_certificate false
-    record_max_buffered_time 1000
-    log_level error
-  </kinesis_producer>
-                    EOS
-                  end
+class KinesisBenchmark
+  def initialize(local = true)
+    @local = local
+    Fluent::Test.setup
+  end
 
-  conf = <<-EOS
-<source>
-  @type dummy
-  tag dummy
-  rate #{rate}
-</source>
+  def run
+    setup
+    benchmark(ENV['SIZE'] || 100, ENV['COUNT'] || 10000)
+    teardown
+  end
 
-<match dummy>
-  @type kinesis_#{type}
-  flush_interval 1
-  buffer_chunk_limit 1m
-  try_flush_interval 0.1
-  queued_chunk_flush_interval 0.01
-  @log_level debug
+  def setup
+    return if not @local
+    @port = Net::EmptyPort.empty_port
+    @server_pid = fork do
+      Process.setsid
+      server = DummyServer.start(port: @port)
+      Signal.trap("TERM") do
+        server.shutdown
+      end
+      server.thread.join
+    end
+    Net::EmptyPort.wait(@port, 3)
+  end
 
-  num_threads 100
+  def teardown
+    return if not @local
+    Process.kill "TERM", @server_pid
+    Process.waitpid @server_pid
+  end
 
-  region ap-northeast-1
-  stream_name fluent-plugin-test
+  def default_config
+    conf = %[
+      log_level error
+      region ap-northeast-1
+      data_key a
+    ]
+    if @local
+      conf += %[
+      endpoint https://localhost:#{@port}
+      ssl_verify_peer false
+      ]
+    end
+    conf
+  end
 
-#{additional_conf}
-</match>
-  EOS
-  conf
+  def create_driver(type, conf = default_config)
+    klass = case type
+            when :streams
+              Fluent::KinesisStreamsOutput
+            when :streams_aggregated
+              Fluent::KinesisStreamsAggregatedOutput
+            when :firehose
+              Fluent::KinesisFirehoseOutput
+            end
+    conf += case type
+            when :streams, :streams_aggregated
+              "stream_name fluent-plugin-test"
+            when :firehose
+              "delivery_stream_name fluent-plugin-test"
+            end
+    if fluentd_v0_12?
+      Fluent::Test::BufferedOutputTestDriver.new(klass) do
+      end.configure(conf)
+    else
+      Fluent::Test::Driver::Output.new(klass) do
+      end.configure(conf)
+    end
+  end
+
+  def benchmark(size, count)
+    record = {"a"=>"a"*size}
+    Benchmark.bmbm(20) do |x|
+      [:streams_aggregated, :streams, :firehose].each do |type|
+        x.report(type)  { driver_run(create_driver(type),  count.times.map{|i|record}) }
+      end
+    end
+  end
 end

data/fluent-plugin-kinesis.gemspec

@@ -1,22 +1,21 @@
 # coding: utf-8
 #
-#  Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
-#  Licensed under the Amazon Software License (the "License").
-#  You may not use this file except in compliance with the License.
-#  A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
 #
-#  http://aws.amazon.com/asl/
+#     http://www.apache.org/licenses/LICENSE-2.0
 #
-#  or in the "license" file accompanying this file. This file is distributed
-#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-#  express or implied. See the License for the specific language governing
-#  permissions and limitations under the License.
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 
 lib = File.expand_path('../lib', __FILE__)
 $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 require 'fluent_plugin_kinesis/version'
-require 'kinesis_producer/binary'
 
 Gem::Specification.new do |spec|
   spec.name          = "fluent-plugin-kinesis"
@@ -24,19 +23,17 @@ Gem::Specification.new do |spec|
   spec.author        = 'Amazon Web Services'
   spec.summary       = %q{Fluentd output plugin that sends events to Amazon Kinesis.}
   spec.homepage      = "https://github.com/awslabs/aws-fluent-plugin-kinesis"
-  spec.license       = "Amazon Software License"
+  spec.license       = "Apache-2.0"
 
   spec.files         = `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) }
-  spec.files         += KinesisProducer::Binary::Files.values
   spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]
-  spec.required_ruby_version = '>= 2.0.0'
+  spec.required_ruby_version = '>= 2.1'
 
-  spec.add_dependency "fluentd", ">= 0.10.58", "< 2"
-  spec.add_dependency "protobuf", ">= 3.5.5"
-  spec.add_dependency "aws-sdk", "~> 2"
-  spec.add_dependency "concurrent-ruby", "~> 1"
-  spec.add_dependency "os", ">= 0.9.6"
+  spec.add_dependency "fluentd", ">= 0.12.35", "< 2"
+  spec.add_dependency "aws-sdk-kinesis", "~> 1"
+  spec.add_dependency "aws-sdk-firehose", "~> 1"
+  spec.add_dependency "google-protobuf", "~> 3"
 
   spec.add_development_dependency "bundler", "~> 1.10"
   spec.add_development_dependency "rake", "~> 10.0"
@@ -46,7 +43,6 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "pry-byebug", ">= 3.3.0"
   spec.add_development_dependency "pry-stack_explorer", ">= 0.4.9.2"
   spec.add_development_dependency "net-empty_port", ">= 0.0.2"
-  spec.add_development_dependency "rubyzip", ">= 1.0.0"
   spec.add_development_dependency "mocha", ">= 1.1.0"
   spec.add_development_dependency "webmock", ">= 1.24.2"
   spec.add_development_dependency "fakefs", ">= 0.8.1"

data/gemfiles/Gemfile.fluentd-0.12

@@ -1,20 +1,20 @@
 #
-#  Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
-#  Licensed under the Amazon Software License (the "License").
-#  You may not use this file except in compliance with the License.
-#  A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
 #
-#  http://aws.amazon.com/asl/
+#     http://www.apache.org/licenses/LICENSE-2.0
 #
-#  or in the "license" file accompanying this file. This file is distributed
-#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-#  express or implied. See the License for the specific language governing
-#  permissions and limitations under the License.
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 
 source 'https://rubygems.org'
 
 # Specify your gem's dependencies in fluent-plugin-kinesis.gemspec
 gemspec path: ".."
 
-gem "fluentd", "~> 0.12.26"
+gem "fluentd", ">= 0.12.35", "< 0.14"

data/lib/fluent/plugin/kinesis.rb

@@ -0,0 +1,166 @@
+#
+# Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+require 'fluent/plugin/kinesis_helper/client'
+require 'fluent/plugin/kinesis_helper/api'
+require 'zlib'
+
+module Fluent
+  class KinesisOutput < BufferedOutput
+    def self.fluentd_v0_12?
+      @fluentd_v0_12 ||= Gem.loaded_specs['fluentd'].version < Gem::Version.create('0.14')
+    end
+
+    include Fluent::SetTimeKeyMixin
+    include Fluent::SetTagKeyMixin
+
+    include KinesisHelper::Client
+    include KinesisHelper::API
+
+    class SkipRecordError < ::StandardError
+      def initialize(message, record)
+        super message
+        @record_message = if record.is_a? Array
+          record.reverse.map(&:to_s).join(', ')
+        else
+          record.to_s
+        end
+      end
+
+      def to_s
+        super + ": " + @record_message
+      end
+    end
+    class KeyNotFoundError < SkipRecordError
+      def initialize(key, record)
+        super "Key '#{key}' doesn't exist", record
+      end
+    end
+    class ExceedMaxRecordSizeError < SkipRecordError
+      def initialize(size, record)
+        super "Record size limit exceeded in #{size/1024} KB", record
+      end
+    end
+    class InvalidRecordError < SkipRecordError
+      def initialize(record)
+        super "Invalid type of record", record
+      end
+    end
+
+    config_param :data_key,              :string,  default: nil
+    config_param :log_truncate_max_size, :integer, default: 1024
+    config_param :compression,           :string,  default: nil
+    config_section :format do
+      config_set_default :@type, 'json'
+    end
+    config_section :inject do
+      config_set_default :time_type, 'string'
+      config_set_default :time_format, '%Y-%m-%dT%H:%M:%S.%N%z'
+    end
+
+    config_param :debug, :bool, default: false
+
+    if fluentd_v0_12?
+      config_param :format, :string,  default: 'json'
+    else
+      helpers :formatter, :inject
+    end
+
+    def configure(conf)
+      super
+      @data_formatter = data_formatter_create(conf)
+    end
+
+    def multi_workers_ready?
+      true
+    end
+
+    private
+
+    def fluentd_v0_12?
+      self.class.fluentd_v0_12?
+    end
+
+    def data_formatter_create(conf)
+      if fluentd_v0_12?
+        formatter = Fluent::Plugin.new_formatter(@format)
+        formatter.configure(conf)
+      else
+        formatter = formatter_create
+      end
+      compressor = compressor_create
+      if @data_key.nil?
+        ->(tag, time, record) {
+          unless fluentd_v0_12?
+            record = inject_values_to_record(tag, time, record)
+          end
+          compressor.call(formatter.format(tag, time, record).chomp.b)
+        }
+      else
+        ->(tag, time, record) {
+          raise InvalidRecordError, record unless record.is_a? Hash
+          raise KeyNotFoundError.new(@data_key, record) if record[@data_key].nil?
+          compressor.call(record[@data_key].to_s.b)
+        }
+      end
+    end
+
+    def compressor_create
+      case @compression
+      when "zlib"
+        ->(data) { Zlib::Deflate.deflate(data) }
+      else
+        ->(data) { data }
+      end
+    end
+
+    def format_for_api(&block)
+      converted = block.call
+      size = size_of_values(converted)
+      if size > @max_record_size
+        raise ExceedMaxRecordSizeError.new(size, converted)
+      end
+      converted.to_msgpack
+    rescue SkipRecordError => e
+      log.error(truncate e)
+      ''
+    end
+
+    def write_records_batch(chunk, &block)
+      if fluentd_v0_12?
+        unique_id = chunk.unique_id.unpack('H*').first
+      else
+        unique_id = chunk.dump_unique_id_hex(chunk.unique_id)
+      end
+      records = chunk.to_enum(:msgpack_each)
+      split_to_batches(records) do |batch, size|
+        log.debug(sprintf "Write chunk %s / %3d records / %4d KB", unique_id, batch.size, size/1024)
+        batch_request_with_retry(batch, &block)
+        log.debug("Finish writing chunk")
+      end
+    end
+
+    def request_type
+      self.class::RequestType
+    end
+
+    def truncate(msg)
+      if @log_truncate_max_size == 0 or (msg.to_s.size <= @log_truncate_max_size)
+        msg.to_s
+      else
+        msg.to_s[0...@log_truncate_max_size]
+      end
+    end
+  end
+end