fluent-plugin-kinesis-modified 3.1.3

data/Rakefile

@@ -0,0 +1,26 @@
+#
+# 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 "bundler/gem_tasks"
+
+require 'rake/testtask'
+
+task default: [:test]
+Rake::TestTask.new do |test|
+  test.libs << 'lib' << 'test'
+  test.test_files = FileList['test/**/test_*.rb']
+  test.options = '-v'
+end
+
+load 'benchmark/task.rake'

data/benchmark/task.rake

@@ -0,0 +1,106 @@
+#
+# 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_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'
+
+namespace :benchmark do
+  task :local do
+    KinesisBenchmark.new.run
+  end
+  task :remote do
+    KinesisBenchmark.new(false).run
+  end
+end
+
+class KinesisBenchmark
+  def initialize(local = true)
+    @local = local
+    Fluent::Test.setup
+  end
+
+  def run
+    setup
+    benchmark(ENV['SIZE'] || 100, ENV['COUNT'] || 10000)
+    teardown
+  end
+
+  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
+
+  def teardown
+    return if not @local
+    Process.kill "TERM", @server_pid
+    Process.waitpid @server_pid
+  end
+
+  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
+
+  def create_driver(type, conf = default_config)
+    klass = case type
+            when :streams
+              Fluent::Plugin::KinesisStreamsOutput
+            when :streams_aggregated
+              Fluent::Plugin::KinesisStreamsAggregatedOutput
+            when :firehose
+              Fluent::Plugin::KinesisFirehoseOutput
+            end
+    conf += case type
+            when :streams, :streams_aggregated
+              "stream_name fluent-plugin-test"
+            when :firehose
+              "delivery_stream_name fluent-plugin-test"
+            end
+
+    Fluent::Test::Driver::Output.new(klass) do
+    end.configure(conf)
+  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

@@ -0,0 +1,63 @@
+# coding: utf-8
+#
+# 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.
+
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'fluent_plugin_kinesis/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = "fluent-plugin-kinesis-modified"
+  spec.version       = FluentPluginKinesis::VERSION
+  spec.author        = 'iHandy'
+  spec.summary       = %q{Fluentd output plugin that sends events to Amazon Kinesis. Forked from fluent-plugin-kinesis version 3.1.0}
+  spec.homepage      = "https://github.com/BlackNight95/aws-fluent-plugin-kinesis/tree/feature_modified_aggregation"
+  spec.license       = "Apache-2.0"
+
+  spec.files         = `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) }
+  spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
+  spec.require_paths = ["lib"]
+  spec.required_ruby_version = '>= 2.1'
+
+  spec.add_dependency "fluentd", ">= 0.14.10", "< 2"
+
+  # This plugin is sometimes used with s3 plugin, so watch out for conflicts
+  #   https://rubygems.org/gems/fluent-plugin-s3
+  # Exclude v1.5 to avoid aws-sdk dependency problem due to this issue
+  #   https://github.com/aws/aws-sdk-ruby/issues/1872
+  # Exclude aws-sdk-kinesis v1.4 to avoid aws-sdk-core dependency problem with td-agent v3.1.1
+  #   NoMethodError: undefined method `event=' for #<Seahorse::Model::Shapes::ShapeRef:*>
+  #   https://github.com/aws/aws-sdk-ruby/commit/03d60f9d3d821e645bd2a3efca066f37350ef906#diff-c69f15af8ea3eb9ab152659476e04608R401
+  #   https://github.com/aws/aws-sdk-ruby/commit/571c2d0e5ff9c24ff72893a08a74790db591fb57#diff-a55155f04aa6559460a0814e264eb0cdR43
+  spec.add_dependency "aws-sdk-kinesis", "~> 1", "!= 1.4", "!= 1.5"
+  # Exclude aws-sdk-firehose v1.9 to avoid aws-sdk-core dependency problem with td-agent v3.2.1
+  #   LoadError: cannot load such file -- aws-sdk-core/plugins/endpoint_discovery.rb
+  #   https://github.com/aws/aws-sdk-ruby/commit/85d8538a62255e58d9e176ee524a9f94354b51a0#diff-d51486091a10ada65b308b7f45966af1R18
+  #   https://github.com/aws/aws-sdk-ruby/commit/7c9584bc6473100df9aec9333ab491ad4faeeca8#diff-be94f87e58e00329a6c0e03e43d5c292
+  spec.add_dependency "aws-sdk-firehose", "~> 1", "!= 1.5", "!= 1.9"
+
+  spec.add_dependency "google-protobuf", "~> 3"
+
+  spec.add_development_dependency "bundler", ">= 1.10"
+  spec.add_development_dependency "rake", ">= 10.0"
+  spec.add_development_dependency "test-unit", ">= 3.0.8"
+  spec.add_development_dependency "test-unit-rr", ">= 1.0.3"
+  spec.add_development_dependency "pry", ">= 0.10.1"
+  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 "mocha", ">= 1.1.0"
+  spec.add_development_dependency "webmock", ">= 1.24.2"
+  spec.add_development_dependency "fakefs", ">= 0.8.1"
+end

data/gemfiles/Gemfile.fluentd-0.14.10

@@ -0,0 +1,20 @@
+#
+# 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.
+
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in fluent-plugin-kinesis.gemspec
+gemspec path: ".."
+
+gem "fluentd", "0.14.10"

data/gemfiles/Gemfile.td-agent-3.1.1

@@ -0,0 +1,31 @@
+#
+# 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.
+
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in fluent-plugin-kinesis.gemspec
+gemspec path: ".."
+
+# Specify related gems for td-agent v3.1.1
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.1.1/config/projects/td-agent3.rb#L22
+gem "fluentd", "1.0.2"
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.1.0/plugin_gems.rb#L16-L23
+gem "jmespath", "1.3.1"
+gem "aws-partitions", "1.42.0"
+gem "aws-sigv4", "1.0.2"
+gem "aws-sdk-core", "3.11.0"
+gem "aws-sdk-kms", "1.3.0"
+gem "aws-sdk-sqs", "1.3.0"
+gem "aws-sdk-s3", "1.8.0"
+gem "fluent-plugin-s3", "1.1.0"

data/gemfiles/Gemfile.td-agent-3.2.0

@@ -0,0 +1,31 @@
+#
+# 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.
+
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in fluent-plugin-kinesis.gemspec
+gemspec path: ".."
+
+# Specify related gems for td-agent v3.2.0
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.2.0/config/projects/td-agent3.rb#L27
+gem "fluentd", "1.2.2"
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.2.0/plugin_gems.rb#L16-L23
+gem "jmespath", "1.4.0"
+gem "aws-partitions", "1.87.0"
+gem "aws-sigv4", "1.0.2"
+gem "aws-sdk-core", "3.21.2"
+gem "aws-sdk-kms", "1.5.0"
+gem "aws-sdk-sqs", "1.3.0"
+gem "aws-sdk-s3", "1.13.0"
+gem "fluent-plugin-s3", "1.1.3"

data/gemfiles/Gemfile.td-agent-3.2.1

@@ -0,0 +1,31 @@
+#
+# 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.
+
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in fluent-plugin-kinesis.gemspec
+gemspec path: ".."
+
+# Specify related gems for td-agent v3.2.1
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.2.1/config/projects/td-agent3.rb#L27
+gem "fluentd", "1.2.6"
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.2.1/plugin_gems.rb#L16-L23
+gem "jmespath", "1.4.0"
+gem "aws-partitions", "1.105.0"
+gem "aws-sigv4", "1.0.3"
+gem "aws-sdk-core", "3.30.0"
+gem "aws-sdk-kms", "1.9.0"
+gem "aws-sdk-sqs", "1.7.0"
+gem "aws-sdk-s3", "1.21.0"
+gem "fluent-plugin-s3", "1.1.6"

data/gemfiles/Gemfile.td-agent-3.3.0

@@ -0,0 +1,31 @@
+#
+# 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.
+
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in fluent-plugin-kinesis.gemspec
+gemspec path: ".."
+
+# Specify related gems for td-agent v3.3.0
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.3.0/config/projects/td-agent3.rb#L27
+gem "fluentd", "1.3.3"
+# https://github.com/treasure-data/omnibus-td-agent/blob/v3.3.0/plugin_gems.rb#L16-L23
+gem "jmespath", "1.4.0"
+gem "aws-partitions", "1.127.0"
+gem "aws-sigv4", "1.0.3"
+gem "aws-sdk-core", "3.44.2"
+gem "aws-sdk-kms", "1.13.0"
+gem "aws-sdk-sqs", "1.10.0"
+gem "aws-sdk-s3", "1.30.0"
+gem "fluent-plugin-s3", "1.1.7"

data/lib/fluent/plugin/kinesis.rb

@@ -0,0 +1,185 @@
+#
+# 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/output'
+require 'fluent/plugin/kinesis_helper/client'
+require 'fluent/plugin/kinesis_helper/api'
+require 'zlib'
+
+module Fluent
+  module Plugin
+    class KinesisOutput < Fluent::Plugin::Output
+      include Fluent::MessagePackFactory::Mixin
+      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
+
+      desc "Formatter calls chomp and removes separator from the end of each record. This option is for compatible format with plugin v2. (default: false)"
+      # https://github.com/awslabs/aws-fluent-plugin-kinesis/issues/142
+      config_param :chomp_record,          :bool,    default: false
+
+      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
+      config_param :max_records_per_call, :integer, default: 128
+
+      helpers :formatter, :inject
+
+      def configure(conf)
+        super
+        @data_formatter = data_formatter_create(conf)
+      end
+
+      def multi_workers_ready?
+        true
+      end
+
+      private
+
+      def data_formatter_create(conf)
+        formatter = formatter_create
+        compressor = compressor_create
+        if @data_key.nil?
+          if @chomp_record
+            ->(tag, time, record) {
+              record = inject_values_to_record(tag, time, record)
+              # Formatter calls chomp and removes separator from the end of each record.
+              # This option is for compatible format with plugin v2.
+              # https://github.com/awslabs/aws-fluent-plugin-kinesis/issues/142
+              compressor.call(formatter.format(tag, time, record).chomp.b)
+            }
+          else
+            ->(tag, time, record) {
+              record = inject_values_to_record(tag, time, record)
+              compressor.call(formatter.format(tag, time, record).b)
+            }
+          end
+        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)
+        unique_id = chunk.dump_unique_id_hex(chunk.unique_id)
+        chunk.open do |io|
+          records = msgpack_unpacker(io).to_enum
+          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
+      end
+
+      def write_records_batch2(chunk, &block)
+        unique_id = chunk.dump_unique_id_hex(chunk.unique_id)
+        chunk.open do |io|
+          records = msgpack_unpacker(io).to_enum
+          batches = []
+          sizes = []
+          split_to_batches(records) do |batch, size|
+            batches << batch
+            sizes << size            
+          end
+          (0..batches.size).step(@max_records_per_call) do |i|
+            batches_slice = batches[i..i+@max_records_per_call-1]
+            size = sizes[i..i+@max_records_per_call-1].inject(:+)
+            batches_size = batches_slice.map(&:size).inject(:+)
+            if not size.nil?
+              log.debug(sprintf "Write chunk %s / %3d batches / %3d records / %4d KB", unique_id, batches_slice.size, batches_size, size/1024)
+              batch_request_with_retry(batches_slice, &block)
+              log.debug("Finish writing chunk")
+            end
+          end
+        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
+end