fluent-plugin-kinesis 1.3.0 → 2.0.0

data/lib/fluent/plugin/kinesis_helper/aggregator.rb

@@ -0,0 +1,99 @@
+#
+# 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/configurable'
+require 'google/protobuf'
+
+Google::Protobuf::DescriptorPool.generated_pool.build do
+  add_message "AggregatedRecord" do
+    repeated :partition_key_table, :string, 1
+    repeated :explicit_hash_key_table, :string, 2
+    repeated :records, :message, 3, "Record"
+  end
+  add_message "Tag" do
+    optional :key, :string, 1
+    optional :value, :string, 2
+  end
+  add_message "Record" do
+    optional :partition_key_index, :uint64, 1
+    optional :explicit_hash_key_index, :uint64, 2
+    optional :data, :bytes, 3
+    repeated :tags, :message, 4, "Tag"
+  end
+end
+
+module Fluent
+  module KinesisHelper
+    class Aggregator
+      AggregatedRecord = Google::Protobuf::DescriptorPool.generated_pool.lookup("AggregatedRecord").msgclass
+      Tag = Google::Protobuf::DescriptorPool.generated_pool.lookup("Tag").msgclass
+      Record = Google::Protobuf::DescriptorPool.generated_pool.lookup("Record").msgclass
+
+      class InvalidEncodingError < ::StandardError; end
+
+      MagicNumber = ['F3899AC2'].pack('H*')
+
+      def aggregate(records, partition_key)
+        message = AggregatedRecord.encode(AggregatedRecord.new(
+          partition_key_table: ['a', partition_key],
+          records: records.map{|data|
+            Record.new(partition_key_index: 1, data: data)
+          },
+        ))
+        [MagicNumber, message, Digest::MD5.digest(message)].pack("A4A*A16")
+      end
+
+      def deaggregate(encoded)
+        unless aggregated?(encoded)
+          raise InvalidEncodingError, "Invalid MagicNumber #{encoded[0..3]}}"
+        end
+        message, digest = encoded[4..encoded.length-17], encoded[encoded.length-16..-1]
+        if Digest::MD5.digest(message) != digest
+          raise InvalidEncodingError, "Digest mismatch #{digest}"
+        end
+        decoded = AggregatedRecord.decode(message)
+        records = decoded.records.map(&:data)
+        partition_key = decoded.partition_key_table[1]
+        [records, partition_key]
+      end
+
+      def aggregated?(encoded)
+        encoded[0..3] == MagicNumber
+      end
+
+      def aggregated_size_offset(partition_key)
+        data = 'd'
+        encoded = aggregate([record(data)], partition_key)
+        finalize(encoded).size - data.size
+      end
+
+      module Mixin
+        AggregateOffset = 25
+        RecordOffset = 10
+
+        module Params
+          include Fluent::Configurable
+        end
+
+        def self.included(mod)
+          mod.include Params
+        end
+
+        def aggregator
+          @aggregator ||= Aggregator.new
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/kinesis_helper/api.rb

@@ -1,163 +1,194 @@
 #
-#  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 'fluent_plugin_kinesis/version'
+require 'fluent/configurable'
 
 module Fluent
   module KinesisHelper
     module API
+      MaxRecordSize = 1024 * 1024 # 1 MB
+
+      module APIParams
+        include Fluent::Configurable
+        config_param :max_record_size, :integer, default: MaxRecordSize
+      end
+
+      def self.included(mod)
+        mod.include APIParams
+      end
+
       def configure(conf)
         super
-        if @batch_request_max_count > self.class::BatchRequestLimitCount
-          raise ConfigError, "batch_request_max_count can't be grater than #{self.class::BatchRequestLimitCount}."
-        end
-        if @batch_request_max_size > self.class::BatchRequestLimitSize
-          raise ConfigError, "batch_request_max_size can't be grater than #{self.class::BatchRequestLimitSize}."
+        if @max_record_size > MaxRecordSize
+          raise ConfigError, "max_record_size can't be grater than #{MaxRecordSize/1024} KB."
         end
-        @region = client.config.region if @region.nil?
       end
 
-      def start
-        detach_multi_process do
-          super
+      module BatchRequest
+        module BatchRequestParams
+          include Fluent::Configurable
+          config_param :retries_on_batch_request, :integer, default: 8
+          config_param :reset_backoff_if_success, :bool,    default: true
+          config_param :batch_request_max_count,  :integer, default: nil
+          config_param :batch_request_max_size,   :integer, default: nil
         end
-      end
 
-      private
-
-      def client_options
-        options = {
-          user_agent_suffix: "fluent-plugin-kinesis/#{request_type}/#{FluentPluginKinesis::VERSION}",
-          credentials: credentials,
-        }
-        options.update(region:          @region)          unless @region.nil?
-        options.update(http_proxy:      @http_proxy)      unless @http_proxy.nil?
-        options.update(endpoint:        @endpoint)        unless @endpoint.nil?
-        options.update(ssl_verify_peer: @ssl_verify_peer) unless @ssl_verify_peer.nil?
-        if @debug
-          options.update(logger: Logger.new(log.out))
-          options.update(log_level: :debug)
+        def self.included(mod)
+          mod.include BatchRequestParams
         end
-        options
-      end
-
-      def split_to_batches(records)
-        batch_by_limit(records, @batch_request_max_count, @batch_request_max_size)
-      end
 
-      def batch_by_limit(records, max_count, max_size)
-        result, buf, size = records.inject([[],[],0]){|(result, buf, size), record|
-          record_size = size_of_values(record)
-          if buf.size+1 > max_count or size+record_size > max_size
-            result << buf
-            buf = []
-            size = 0
+        def configure(conf)
+          super
+          if @batch_request_max_count.nil?
+            @batch_request_max_count = self.class::BatchRequestLimitCount
+          elsif @batch_request_max_count > self.class::BatchRequestLimitCount
+            raise ConfigError, "batch_request_max_count can't be grater than #{self.class::BatchRequestLimitCount}."
           end
-          buf << record
-          size += record_size
-          [result, buf, size]
-        }
-        result << buf
-      end
+          if @batch_request_max_size.nil?
+            @batch_request_max_size = self.class::BatchRequestLimitSize
+          elsif @batch_request_max_size > self.class::BatchRequestLimitSize
+            raise ConfigError, "batch_request_max_size can't be grater than #{self.class::BatchRequestLimitSize}."
+          end
+        end
 
-      def size_of_values(record)
-        record.values_at(:data, :partition_key).compact.map(&:size).inject(:+) || 0
-      end
+        def size_of_values(record)
+          record.compact.map(&:size).inject(:+) || 0
+        end
 
-      def batch_request_with_retry(batch, retry_count=0, backoff: nil)
-        backoff ||= Backoff.new
-        res = batch_request(batch)
-        if failed_count(res) > 0
-          failed_records = collect_failed_records(batch, res)
-          if retry_count < @retries_on_batch_request
-            backoff.reset if @reset_backoff_if_success and any_records_shipped?(res)
-            sleep(backoff.next)
-            log.warn(truncate 'Retrying to request batch. Retry count: %d, Retry records: %d' % [retry_count, failed_records.size])
-            retry_batch = failed_records.map{|r| r[:original] }
-            batch_request_with_retry(retry_batch, retry_count + 1, backoff: backoff)
-          else
-            give_up_retries(failed_records)
+        private
+
+        def split_to_batches(records, &block)
+          batch = []
+          size = 0
+          records.each do |record|
+            record_size = size_of_values(record)
+            if batch.size+1 > @batch_request_max_count or size+record_size > @batch_request_max_size
+              yield(batch, size)
+              batch = []
+              size = 0
+	    end
+            batch << record
+            size += record_size
           end
+          yield(batch, size) if batch.size > 0
         end
-      end
-
-      def any_records_shipped?(res)
-        results(res).size > failed_count(res)
-      end
 
-      def collect_failed_records(records, res)
-        failed_records = []
-        results(res).each_with_index do |record, index|
-          next unless record[:error_code]
-          failed_records.push(
-            original:      records[index],
-            error_code:    record[:error_code],
-            error_message: record[:error_message]
-          )
+        def batch_request_with_retry(batch, retry_count=0, backoff: nil, &block)
+          backoff ||= Backoff.new
+          res = yield(batch)
+          if failed_count(res) > 0
+            failed_records = collect_failed_records(batch, res)
+            if retry_count < @retries_on_batch_request
+              backoff.reset if @reset_backoff_if_success and any_records_shipped?(res)
+              wait_second = backoff.next
+              msg = 'Retrying to request batch. Retry count: %3d, Retry records: %3d, Wait seconds %3.2f' % [retry_count+1, failed_records.size, wait_second]
+              log.warn(truncate msg)
+              # TODO: sleep() doesn't wait the given seconds sometime.
+              # The root cause is unknown so far, so I'd like to add debug print only. It should be fixed in the future.
+              log.debug("#{Thread.current.object_id} sleep start")
+              sleep(wait_second)
+              log.debug("#{Thread.current.object_id} sleep finish")
+              batch_request_with_retry(retry_records(failed_records), retry_count+1, backoff: backoff, &block)
+            else
+              give_up_retries(failed_records)
+            end
+          end
         end
-        failed_records
-      end
 
-      def failed_count(res)
-        failed_field = case request_type
-                       when :streams;  :failed_record_count
-                       when :firehose; :failed_put_count
-                       end
-        res[failed_field]
-      end
+        def any_records_shipped?(res)
+          results(res).size > failed_count(res)
+        end
 
-      def results(res)
-        result_field = case request_type
-                       when :streams;  :records
-                       when :firehose; :request_responses
+        def collect_failed_records(records, res)
+          failed_records = []
+          results(res).each_with_index do |record, index|
+            next unless record[:error_code]
+            original = case request_type
+                       when :streams, :firehose; records[index]
+                       when :streams_aggregated; records
                        end
-        res[result_field]
-      end
+            failed_records.push(
+              original:      original,
+              error_code:    record[:error_code],
+              error_message: record[:error_message]
+            )
+          end
+          failed_records
+        end
 
-      def give_up_retries(failed_records)
-        failed_records.each {|record|
-          log.error(truncate 'Could not put record, Error: %s/%s, Record: %s' % [
-            record[:error_code],
-            record[:error_message],
-            record[:original]
-          ])
-        }
-      end
+        def retry_records(failed_records)
+          case request_type
+          when :streams, :firehose
+            failed_records.map{|r| r[:original] }
+          when :streams_aggregated
+            failed_records.first[:original]
+          end
+        end
 
-      class Backoff
-        def initialize
-          @count = 0
+        def failed_count(res)
+          failed_field = case request_type
+                         when :streams;            :failed_record_count
+                         when :streams_aggregated; :failed_record_count
+                         when :firehose;           :failed_put_count
+                         end
+          res[failed_field]
         end
 
-        def next
-          value = calc(@count)
-          @count += 1
-          value
+        def results(res)
+          result_field = case request_type
+                         when :streams;            :records
+                         when :streams_aggregated; :records
+                         when :firehose;           :request_responses
+                         end
+          res[result_field]
         end
 
-        def reset
-          @count = 0
+        def give_up_retries(failed_records)
+          failed_records.each {|record|
+            log.error(truncate 'Could not put record, Error: %s/%s, Record: %s' % [
+              record[:error_code],
+              record[:error_message],
+              record[:original]
+            ])
+          }
         end
 
-        private
+        class Backoff
+          def initialize
+            @count = 0
+          end
 
-        def calc(count)
-          (2 ** count) * scaling_factor
-        end
+          def next
+            value = calc(@count)
+            @count += 1
+            value
+          end
+
+          def reset
+            @count = 0
+          end
+
+          private
 
-        def scaling_factor
-          0.3 + (0.5-rand) * 0.1
+          def calc(count)
+            (2 ** count) * scaling_factor
+          end
+
+          def scaling_factor
+            0.3 + (0.5-rand) * 0.1
+          end
         end
       end
     end

data/lib/fluent/plugin/kinesis_helper/client.rb

@@ -1,20 +1,75 @@
 #
-#  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 'fluent/configurable'
 
 module Fluent
   module KinesisHelper
     module Client
+      module ClientParams
+        include Fluent::Configurable
+        config_param :region, :string,  default: nil
+
+        config_param :http_proxy,      :string, default: nil, secret: true
+        config_param :endpoint,        :string, default: nil
+        config_param :ssl_verify_peer, :bool,   default: true
+
+        config_param :aws_key_id,  :string, default: nil, secret: true
+        config_param :aws_sec_key, :string, default: nil, secret: true
+        config_section :assume_role_credentials, multi: false do
+          desc "The Amazon Resource Name (ARN) of the role to assume"
+          config_param :role_arn, :string, secret: true
+          desc "An identifier for the assumed role session"
+          config_param :role_session_name, :string
+          desc "An IAM policy in JSON format"
+          config_param :policy, :string, default: nil
+          desc "The duration, in seconds, of the role session (900-3600)"
+          config_param :duration_seconds, :integer, default: nil
+          desc "A unique identifier that is used by third parties when assuming roles in their customers' accounts."
+          config_param :external_id, :string, default: nil, secret: true
+        end
+        config_section :instance_profile_credentials, multi: false do
+          desc "Number of times to retry when retrieving credentials"
+          config_param :retries, :integer, default: nil
+          desc "IP address (default:169.254.169.254)"
+          config_param :ip_address, :string, default: nil
+          desc "Port number (default:80)"
+          config_param :port, :integer, default: nil
+          desc "Number of seconds to wait for the connection to open"
+          config_param :http_open_timeout, :float, default: nil
+          desc "Number of seconds to wait for one block to be read"
+          config_param :http_read_timeout, :float, default: nil
+          # config_param :delay, :integer or :proc, :default => nil
+          # config_param :http_degub_output, :io, :default => nil
+        end
+        config_section :shared_credentials, multi: false do
+          desc "Path to the shared file. (default: $HOME/.aws/credentials)"
+          config_param :path, :string, default: nil
+          desc "Profile name. Default to 'default' or ENV['AWS_PROFILE']"
+          config_param :profile_name, :string, default: nil
+        end
+      end
+
+      def self.included(mod)
+        mod.include ClientParams
+      end
+
+      def configure(conf)
+        super
+        @region = client.config.region if @region.nil?
+      end
+
       def client
         @client ||= client_class.new(client_options)
       end
@@ -23,14 +78,72 @@ module Fluent
 
       def client_class
         case request_type
-        when :streams
+        when :streams, :streams_aggregated
+          require 'aws-sdk-kinesis'
           Aws::Kinesis::Client
         when :firehose
+          require 'aws-sdk-firehose'
           Aws::Firehose::Client
-        when :producer
-          KinesisProducer::Library
         end
       end
+
+      def client_options
+        options = setup_credentials
+        options.update(
+          user_agent_suffix: "fluent-plugin-kinesis/#{request_type}/#{FluentPluginKinesis::VERSION}"
+        )
+        options.update(region:          @region)          unless @region.nil?
+        options.update(http_proxy:      @http_proxy)      unless @http_proxy.nil?
+        options.update(endpoint:        @endpoint)        unless @endpoint.nil?
+        options.update(ssl_verify_peer: @ssl_verify_peer) unless @ssl_verify_peer.nil?
+        if @debug
+          options.update(logger: Logger.new(log.out))
+          options.update(log_level: :debug)
+        end
+        options
+      end
+
+      def setup_credentials
+        options = {}
+        credentials_options = {}
+        case
+        when @aws_key_id && @aws_sec_key
+          options[:access_key_id] = @aws_key_id
+          options[:secret_access_key] = @aws_sec_key
+        when @assume_role_credentials
+          c = @assume_role_credentials
+          credentials_options[:role_arn] = c.role_arn
+          credentials_options[:role_session_name] = c.role_session_name
+          credentials_options[:policy] = c.policy if c.policy
+          credentials_options[:duration_seconds] = c.duration_seconds if c.duration_seconds
+          credentials_options[:external_id] = c.external_id if c.external_id
+          if @region
+            credentials_options[:client] = Aws::STS::Client.new(region: @region)
+          end
+          options[:credentials] = Aws::AssumeRoleCredentials.new(credentials_options)
+        when @instance_profile_credentials
+          c = @instance_profile_credentials
+          credentials_options[:retries] = c.retries if c.retries
+          credentials_options[:ip_address] = c.ip_address if c.ip_address
+          credentials_options[:port] = c.port if c.port
+          credentials_options[:http_open_timeout] = c.http_open_timeout if c.http_open_timeout
+          credentials_options[:http_read_timeout] = c.http_read_timeout if c.http_read_timeout
+          if ENV["AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"]
+            options[:credentials] = Aws::ECSCredentials.new(credentials_options)
+          else
+            options[:credentials] = Aws::InstanceProfileCredentials.new(credentials_options)
+          end
+        when @shared_credentials
+          c = @shared_credentials
+          credentials_options[:path] = c.path if c.path
+          credentials_options[:profile_name] = c.profile_name if c.profile_name
+          options[:credentials] = Aws::SharedCredentials.new(credentials_options)
+        else
+          # Use default credentials
+          # See http://docs.aws.amazon.com/sdkforruby/api/Aws/S3/Client.html
+        end
+        options
+      end
     end
   end
 end