kcl-rb 1.0.0

data/lib/kcl/checkpoints/sentinel.rb

@@ -0,0 +1,17 @@
+# Enumeration of the sentinel values of checkpoints.
+# Used during initialization of ShardConsumers to determine the starting point
+# in the shard and to flag that a shard has been completely processed.
+module Kcl::Checkpoints
+  module Sentinel
+    # Start from the first available record in the shard.
+    TRIM_HORIZON = 'TRIM_HORIZON'.freeze
+    # Start from the latest record in the shard.
+    LATEST       = 'LATEST'.freeze
+    # We've completely processed all records in this shard.
+    SHARD_END    = 'SHARD_END'.freeze
+    # Start from the record at or after the specified server-side timestamp.
+    AT_TIMESTAMP = 'AT_TIMESTAMP'.freeze
+    # Continue from the sequence number in the shard.
+    AFTER_SEQUENCE_NUMBER = 'AFTER_SEQUENCE_NUMBER'.freeze
+  end
+end

data/lib/kcl/config.rb

@@ -0,0 +1,35 @@
+class Kcl::Config
+  attr_accessor :aws_region,
+    :aws_access_key_id,
+    :aws_secret_access_key,
+    :dynamodb_endpoint,
+    :dynamodb_table_name,
+    :dynamodb_read_capacity,
+    :dynamodb_write_capacity,
+    :dynamodb_failover_seconds,
+    :kinesis_endpoint,
+    :kinesis_stream_name,
+    :logger,
+    :log_level,
+    :max_lease_count,
+    :use_ssl,
+    :worker_count
+
+  # Set default values
+  def initialize
+    @aws_region              = nil
+    @aws_access_key_id       = nil
+    @aws_secret_access_key   = nil
+    @dynamodb_endpoint       = 'https://localhost:4566'
+    @dynamodb_table_name     = nil
+    @dynamodb_read_capacity  = 10
+    @dynamodb_write_capacity = 10
+    @dynamodb_failover_seconds = 10
+    @kinesis_endpoint        = 'https://localhost:4566'
+    @kinesis_stream_name     = nil
+    @logger                  = nil
+    @max_lease_count         = 1
+    @use_ssl                 = false
+    @worker_count            = 1
+  end
+end

data/lib/kcl/errors.rb

@@ -0,0 +1,6 @@
+module Kcl::Errors
+  class IllegalArgumentError < StandardError; end
+  class CheckpointNotFoundError < StandardError; end
+  class SequenceNumberNotFoundError < StandardError; end
+  class LeaseNotAquiredError < StandardError; end
+end

data/lib/kcl/logger.rb

@@ -0,0 +1,3 @@
+require 'logger'
+
+class Kcl::Logger < ::Logger; end

data/lib/kcl/proxies/dynamo_db_proxy.rb

@@ -0,0 +1,132 @@
+require 'aws-sdk-dynamodb'
+
+module Kcl::Proxies
+  class DynamoDbProxy
+    attr_reader :client
+
+    def initialize(config)
+      @client = Aws::DynamoDB::Client.new(
+        {
+          access_key_id: config.aws_access_key_id,
+          secret_access_key: config.aws_secret_access_key,
+          region: config.aws_region,
+          endpoint: config.dynamodb_endpoint,
+          ssl_verify_peer: config.use_ssl
+        }
+      )
+    end
+
+    # @params [String] table_name
+    def exists?(table_name)
+      @client.describe_table({ table_name: table_name })
+      true
+    rescue Aws::DynamoDB::Errors::NotFound,
+           Aws::DynamoDB::Errors::ResourceNotFoundException
+      false
+    end
+
+    # @params [String] table_name
+    # @params [Array]  attributes
+    # @params [Array]  schema
+    # @params [Hash]   throughputs
+    def create_table(table_name, attributes = [], schema = [], throughputs = {})
+      @client.create_table(
+        {
+          table_name: table_name,
+          attribute_definitions: attributes,
+          key_schema: schema,
+          provisioned_throughput: throughputs
+        }
+      )
+    end
+
+    # @params [String] table_name
+    def delete_table(table_name)
+      @client.delete_table({ table_name: table_name })
+      true
+    rescue Aws::DynamoDB::Errors::ResourceNotFoundException
+      false
+    end
+
+    # @params [String] table_name
+    # @params [Hash]   conditions
+    # @return [Hash]
+    def get_item(table_name, conditions)
+      response = @client.get_item(
+        {
+          table_name: table_name,
+          key: conditions
+        }
+      )
+      response.item
+    rescue Aws::DynamoDB::Errors::ResourceNotFoundException
+      nil
+    end
+
+    # @params [String]  table_name
+    # @params [Hash]    item
+    # @return [Boolean]
+    def put_item(table_name, item)
+      @client.put_item(
+        {
+          table_name: table_name,
+          item: item
+        }
+      )
+      true
+    rescue Aws::DynamoDB::Errors::ResourceNotFoundException
+      false
+    end
+
+    # @params [String]  table_name
+    # @params [Hash]    conditions
+    # @params [String]  update_expression
+    # @return [Boolean]
+    def update_item(table_name, conditions, update_expression)
+      @client.update_item(
+        {
+          table_name: table_name,
+          key: conditions,
+          update_expression: update_expression
+        }
+      )
+      true
+    rescue Aws::DynamoDB::Errors::ResourceNotFoundException
+      false
+    end
+
+    # @params [String]  table_name
+    # @params [Hash]    item
+    # @params [String]  condition_expression
+    # @params [Hash]    expression_attributes
+    # @return [Boolean]
+    def conditional_update_item(table_name, item, condition_expression, expression_attributes)
+      @client.put_item(
+        {
+          table_name: table_name,
+          item: item,
+          condition_expression: condition_expression,
+          expression_attribute_values: expression_attributes
+        }
+      )
+      true
+    rescue Aws::DynamoDB::Errors::ResourceNotFoundException
+      false
+    end
+
+    # @params [String]  table_name
+    # @params [Hash]    conditions
+    # @return [Boolean]
+    def remove_item(table_name, conditions)
+      @client.delete_item(
+        {
+          table_name: table_name,
+          key: conditions
+        }
+      )
+      true
+    rescue Aws::DynamoDB::Errors::ResourceNotFoundException
+      false
+    end
+  end
+end

data/lib/kcl/proxies/kinesis_proxy.rb

@@ -0,0 +1,56 @@
+require 'aws-sdk-kinesis'
+
+module Kcl::Proxies
+  class KinesisProxy
+    attr_reader :client
+
+    def initialize(config)
+      @client = Aws::Kinesis::Client.new(
+        {
+          access_key_id: config.aws_access_key_id,
+          secret_access_key: config.aws_secret_access_key,
+          region: config.aws_region,
+          endpoint: config.kinesis_endpoint,
+          ssl_verify_peer: config.use_ssl
+        }
+      )
+      @stream_name = config.kinesis_stream_name
+    end
+
+    # @return [Array]
+    def shards
+      res = @client.describe_stream({ stream_name: @stream_name })
+      res.stream_description.shards
+    end
+
+    # @param [String] shard_id
+    # @param [String] shard_iterator_type
+    # @return [String]
+    def get_shard_iterator(shard_id, shard_iterator_type = nil, sequence_number = nil)
+      params = {
+        stream_name: @stream_name,
+        shard_id: shard_id,
+        shard_iterator_type: shard_iterator_type || Kcl::Checkpoints::Sentinel::LATEST
+      }
+      if shard_iterator_type == Kcl::Checkpoints::Sentinel::AFTER_SEQUENCE_NUMBER
+        params[:starting_sequence_number] = sequence_number
+      end
+      res = @client.get_shard_iterator(params)
+      res.shard_iterator
+    end
+
+    # @param [String] shard_iterator
+    # @return [Hash]
+    def get_records(shard_iterator)
+      res = @client.get_records({ shard_iterator: shard_iterator })
+      { records: res.records, next_shard_iterator: res.next_shard_iterator }
+    end
+
+    # @param [Hash] data
+    # @return [Hash]
+    def put_record(data)
+      res = @client.put_record(data)
+      { shard_id: res.shard_id, sequence_number: res.sequence_number }
+    end
+  end
+end

data/lib/kcl/record_processor.rb

@@ -0,0 +1,13 @@
+class Kcl::RecordProcessor
+  def after_initialize(_initialization_input)
+    raise NotImplementedError.new("You must implement #{self.class}##{__method__}")
+  end
+
+  def process_records(_records_input)
+    raise NotImplementedError.new("You must implement #{self.class}##{__method__}")
+  end
+
+  def shutdown(_shutdown_input)
+    raise NotImplementedError.new("You must implement #{self.class}##{__method__}")
+  end
+end

data/lib/kcl/record_processor_factory.rb

@@ -0,0 +1,5 @@
+class Kcl::RecordProcessorFactory
+  def create_processor
+    raise NotImplementedError.new("You must implement #{self.class}##{__method__}")
+  end
+end

data/lib/kcl/types/extended_sequence_number.rb

@@ -0,0 +1,89 @@
+require 'bigdecimal'
+
+module Kcl::Types
+  class ExtendedSequenceNumber
+    attr_reader :sequence_number, :sub_sequence_number
+
+    TRIM_HORIZON_VALUE = BigDecimal(-2)
+    LATEST_VALUE       = BigDecimal(-1)
+    AT_TIMESTAMP_VALUE = BigDecimal(-3)
+
+    # @return [Kcl::Types::ExtendedSequenceNumber]
+    def self.latest
+      @_latest ||= self.new(Kcl::Checkpoints::Sentinel::LATEST)
+    end
+
+    # @return [Kcl::Types::ExtendedSequenceNumber]
+    def self.shard_end
+      @_shard_end ||= self.new(Kcl::Checkpoints::Sentinel::SHARD_END)
+    end
+
+    # @return [Kcl::Types::ExtendedSequenceNumber]
+    def self.trim_horizon
+      @_trim_horizon ||= self.new(Kcl::Checkpoints::Sentinel::TRIM_HORIZON)
+    end
+
+    # @param [String] str
+    # @return [Boolean]
+    def self.digits_or_sentinel?(str)
+      digits?(str) || sentinel?(str)
+    end
+
+    # @param [String] str
+    # @return [Boolean]
+    def self.sentinel?(str)
+      case str
+      when Kcl::Checkpoints::Sentinel::TRIM_HORIZON,
+        Kcl::Checkpoints::Sentinel::LATEST,
+        Kcl::Checkpoints::Sentinel::SHARD_END,
+        Kcl::Checkpoints::Sentinel::AT_TIMESTAMP
+        true
+      else
+        false
+      end
+    end
+
+    # @param [String] str
+    # @return [Boolean]
+    def self.digits?(str)
+      return false if str.nil? || str.empty?
+      (str =~ /\A[0-9]+\z/) != nil
+    end
+
+    # @param [String] sequence_number
+    # @param [Number] sub_sequence_number
+    def initialize(sequence_number, sub_sequence_number = 0)
+      @sequence_number     = sequence_number
+      @sub_sequence_number = sub_sequence_number
+    end
+
+    # @return [BigDecimal]
+    def value
+      if self.class.digits?(@sequence_number)
+        return BigDecimal(@sequence_number)
+      end
+
+      case @sequence_number
+      when Kcl::Checkpoints::Sentinel::LATEST
+        LATEST_VALUE
+      when Kcl::Checkpoints::Sentinel::TRIM_HORIZON
+        TRIM_HORIZON_VALUE
+      when Kcl::Checkpoints::Sentinel::AT_TIMESTAMP
+        AT_TIMESTAMP_VALUE
+      else
+        raise Kcl::Errors::IllegalArgumentError.new(
+          'Expected a string of digits, TRIM_HORIZON, LATEST or AT_TIMESTAMP but received ' + @sequence_number
+        )
+      end
+    end
+
+    # @param [Kcl::Types::ExtendedSequenceNumber] extended_sequence_number
+    # @return [Boolean]
+    def equals(extended_sequence_number)
+      if @sequence_number != extended_sequence_number.sequence_number
+        return false
+      end
+      @sub_sequence_number == extended_sequence_number.sub_sequence_number
+    end
+  end
+end

data/lib/kcl/types/initialization_input.rb

@@ -0,0 +1,13 @@
+module Kcl::Types
+  # Container for the parameters to the RecordProcessor's method.
+  class InitializationInput
+    attr_reader :shard_id, :extended_sequence_number
+
+    # @param [String] shard_id
+    # @param [Kcl::Types::ExtendedSequenceNumber] extended_sequence_number
+    def initialize(shard_id, extended_sequence_number)
+      @shard_id                 = shard_id
+      @extended_sequence_number = extended_sequence_number
+    end
+  end
+end

data/lib/kcl/types/records_input.rb

@@ -0,0 +1,15 @@
+module Kcl::Types
+  # Container for the parameters to the IRecordProcessor's method.
+  class RecordsInput
+    attr_reader :records, :millis_behind_latest, :record_checkpointer
+
+    # @param [Array] records
+    # @param [Number] millis_behind_latest
+    # @param [Kcl::Workers::RecordCheckpointer] record_checkpointer
+    def initialize(records, millis_behind_latest, record_checkpointer)
+      @records              = records
+      @millis_behind_latest = millis_behind_latest
+      @record_checkpointer  = record_checkpointer
+    end
+  end
+end

data/lib/kcl/types/shutdown_input.rb

@@ -0,0 +1,13 @@
+module Kcl::Types
+  # Container for the parameters to the IRecordProcessor's method.
+  class ShutdownInput
+    attr_reader :shutdown_reason, :record_checkpointer
+
+    # @param [Kcl::Worker::ShutdownReason] shutdown_reason
+    # @param [Kcl::Workers::RecordCheckpointer] record_checkpointer
+    def initialize(shutdown_reason, record_checkpointer)
+      @shutdown_reason     = shutdown_reason
+      @record_checkpointer = record_checkpointer
+    end
+  end
+end

data/lib/kcl/version.rb

@@ -0,0 +1,3 @@
+module Kcl
+  VERSION = '1.0.0'.freeze
+end

data/lib/kcl/worker.rb

@@ -0,0 +1,159 @@
+require 'eventmachine'
+
+class Kcl::Worker
+  PROCESS_INTERVAL = 1 # by sec
+
+  def self.run(id, record_processor_factory)
+    worker = self.new(id, record_processor_factory)
+    worker.start
+  end
+
+  def initialize(id, record_processor_factory)
+    @id = id
+    @record_processor_factory = record_processor_factory
+    @live_shards  = {} # Map<String, Boolean>
+    @shards       = {} # Map<String, Kcl::Workers::ShardInfo>
+    @kinesis      = nil # Kcl::Proxies::KinesisProxy
+    @checkpointer = nil # Kcl::Checkpointer
+    @timer        = nil
+  end
+
+  # Start consuming data from the stream,
+  # and pass it to the application record processors.
+  def start
+    Kcl.logger.info("Start worker at #{object_id}")
+
+    EM.run do
+      trap_signals
+
+      @timer = EM::PeriodicTimer.new(PROCESS_INTERVAL) do
+        sync_shards!
+        consume_shards! if available_lease_shard?
+      end
+    end
+
+    cleanup
+    Kcl.logger.info("Finish worker at #{object_id}")
+  rescue => e
+    Kcl.logger.error("#{e.class}: #{e.message}")
+    raise e
+  end
+
+  # Shutdown gracefully
+  def shutdown(signal = :NONE)
+    unless @timer.nil?
+      @timer.cancel
+      @timer = nil
+    end
+    EM.stop
+
+    Kcl.logger.info("Shutdown worker with signal #{signal} at #{object_id}")
+  rescue => e
+    Kcl.logger.error("#{e.class}: #{e.message}")
+    raise e
+  end
+
+  # Cleanup resources
+  def cleanup
+    @live_shards  = {}
+    @shards       = {}
+    @kinesis      = nil
+    @checkpointer = nil
+  end
+
+  # Add new shards and delete unused shards
+  def sync_shards!
+    @live_shards.transform_values! { |_| false }
+
+    kinesis.shards.each do |shard|
+      @live_shards[shard.shard_id] = true
+      next if @shards[shard.shard_id]
+      @shards[shard.shard_id] = Kcl::Workers::ShardInfo.new(
+        shard.shard_id,
+        shard.parent_shard_id,
+        shard.sequence_number_range
+      )
+      Kcl.logger.info("Found new shard at shard_id: #{shard.shard_id}")
+    end
+
+    @live_shards.each do |shard_id, alive|
+      next if alive
+      checkpointer.remove_lease(@shards[shard_id])
+      @shards.delete(shard_id)
+      Kcl.logger.info("Remove shard at shard_id: #{shard_id}")
+    end
+
+    @shards
+  end
+
+  # Count the number of leases hold by worker excluding the processed shard
+  # @return [Boolean]
+  def available_lease_shard?
+    leased_count = @shards.values.inject(0) do |num, shard|
+      shard.lease_owner == @id && !shard.completed? ? num + 1 : num
+    end
+    Kcl.config.max_lease_count > leased_count
+  end
+
+  # Process records by shard
+  def consume_shards!
+    threads = []
+    @shards.each do |shard_id, shard|
+      # already owner of the shard
+      next if shard.lease_owner == @id
+
+      begin
+        shard = checkpointer.fetch_checkpoint(shard)
+      rescue Kcl::Errors::CheckpointNotFoundError
+        Kcl.logger.info("Not found checkpoint of shard at #{shard.to_h}")
+        next
+      end
+      # shard is closed and processed all records
+      next if shard.completed?
+
+      shard = checkpointer.lease(shard, @id)
+
+      threads << Thread.new do
+        begin
+          consumer = Kcl::Workers::Consumer.new(
+            shard,
+            @record_processor_factory.create_processor,
+            kinesis,
+            checkpointer
+          )
+          consumer.consume!
+        ensure
+          shard = checkpointer.remove_lease_owner(shard)
+          Kcl.logger.info("Finish to consume shard at shard_id: #{shard_id}")
+        end
+      end
+    end
+    threads.each(&:join)
+  end
+
+  private
+
+  def kinesis
+    if @kinesis.nil?
+      @kinesis = Kcl::Proxies::KinesisProxy.new(Kcl.config)
+      Kcl.logger.info('Created Kinesis session in worker')
+    end
+    @kinesis
+  end
+
+  def checkpointer
+    if @checkpointer.nil?
+      @checkpointer = Kcl::Checkpointer.new(Kcl.config)
+      Kcl.logger.info('Created Checkpoint in worker')
+    end
+    @checkpointer
+  end
+
+  def trap_signals
+    [:HUP, :INT, :TERM].each do |signal|
+      trap signal do
+        EM.add_timer(0) { shutdown(signal) }
+      end
+    end
+  end
+end