telekinesis 2.0.0-java

data/lib/telekinesis/producer/async_producer.rb

@@ -0,0 +1,157 @@
+require "telekinesis/producer/async_producer_worker"
+
+module Telekinesis
+  module Producer
+    java_import java.util.concurrent.TimeUnit
+    java_import java.util.concurrent.Executors
+    java_import java.util.concurrent.ArrayBlockingQueue
+    java_import com.google.common.util.concurrent.ThreadFactoryBuilder
+
+    # An asynchronous producer that buffers events into a queue and uses a
+    # background thread to send them to Kinesis. Only available on JRuby.
+    #
+    # This class is thread-safe.
+    class AsyncProducer
+      # For convenience
+      MAX_PUT_RECORDS_SIZE = Telekinesis::Aws::KINESIS_MAX_PUT_RECORDS_SIZE
+
+      attr_reader :stream, :client, :failure_handler
+
+      # Create a new producer.
+      #
+      # AWS credentials may be specified by using the `:credentials` option and
+      # passing a hash containing your `:access_key_id` and `:secret_access_key`.
+      # If unspecified, credentials will be fetched from the environment, an
+      # ~/.aws/credentials file, or the current instance metadata.
+      #
+      # The producer's `:worker_count`, internal `:queue_size`, the `:send_size`
+      # of batches to Kinesis and how often workers send data to Kinesis, even
+      # if their batches aren't full (`:send_every_ms`) can be configured as
+      # well. They all have reasonable defaults.
+      #
+      # When requests to Kinesis fail, the configured `:failure_handler` will
+      # be called. If you don't specify a failure handler, a NoopFailureHandler
+      # is used.
+      def self.create(options = {})
+        stream = options[:stream]
+        client = Telekinesis::Aws::Client.build(options.fetch(:credentials, {}))
+        failure_handler = options.fetch(:failure_handler, NoopFailureHandler.new)
+        new(stream, client, failure_handler, options)
+      end
+
+      # Construct a new producer. Intended for internal use only - prefer
+      # #create unless it's strictly necessary.
+      def initialize(stream, client, failure_handler, options = {})
+        @stream = stream or raise ArgumentError, "stream may not be nil"
+        @client = client or raise ArgumentError, "client may not be nil"
+        @failure_handler = failure_handler or raise ArgumentError, "failure_handler may not be nil"
+        @shutdown = false
+
+        queue_size     = options.fetch(:queue_size, 1000)
+        send_every     = options.fetch(:send_every_ms, 1000)
+        worker_count   = options.fetch(:worker_count, 1)
+        raise ArgumentError(":worker_count must be > 0") unless worker_count > 0
+        send_size      = options.fetch(:send_size, MAX_PUT_RECORDS_SIZE)
+        raise ArgumentError(":send_size too large") if send_size > MAX_PUT_RECORDS_SIZE
+        retries        = options.fetch(:retries, 5)
+        raise ArgumentError(":retries must be >= 0") unless retries >= 0
+        retry_interval = options.fetch(:retry_interval, 1.0)
+        raise ArgumentError(":retry_interval must be > 0") unless retry_interval > 0
+
+        # NOTE: For testing.
+        @queue = options[:queue] || ArrayBlockingQueue.new(queue_size)
+
+        @lock = Telekinesis::JavaUtil::ReadWriteLock.new
+        @worker_pool = build_executor(worker_count)
+        @workers = worker_count.times.map do
+          AsyncProducerWorker.new(self, @queue, send_size, send_every, retries, retry_interval)
+        end
+
+        # NOTE: Start by default. For testing.
+        start unless options.fetch(:manual_start, false)
+      end
+
+      # Put a single key, value pair to Kinesis. Both key and value must be
+      # strings.
+      #
+      # This call returns immediately and returns true iff the producer is still
+      # accepting data. Data is put to Kinesis in the background.
+      def put(key, data)
+        put_all(key => data)
+      end
+
+      # Put all of the given key, value pairs to Kinesis. Both key and value
+      # must be Strings.
+      #
+      # This call returns immediately and returns true iff the producer is still
+      # accepting data. Data is put to Kinesis in the background.
+      def put_all(items)
+        # NOTE: The lock ensures that no new data can be added to the queue after
+        # the shutdown flag has been set. See the note in shutdown for details.
+        @lock.read_lock do
+          if @shutdown
+            false
+          else
+            items.each do |key, data|
+              @queue.put([key, data])
+            end
+            true
+          end
+        end
+      end
+
+      # Shut down this producer. After the call completes, the producer will not
+      # accept any more data, but will finish processing any data it has
+      # buffered internally.
+      #
+      # If block = true is passed, this call will block and wait for the producer
+      # to shut down before returning. This wait times out after duration has
+      # passed.
+      def shutdown(block = false, duration = 2, unit = TimeUnit::SECONDS)
+        # NOTE: Since a write_lock is exclusive, this prevents any data from being
+        # added to the queue while the SHUTDOWN tokens are being inserted. Without
+        # the lock, data can end up in the queue behind all of the shutdown tokens
+        # and be lost. This happens if the shutdown flag is be flipped by a thread
+        # calling shutdown after another thread has checked the "if @shutdown"
+        # condition in put but before it's called queue.put.
+        @lock.write_lock do
+          @shutdown = true
+          @workers.size.times do
+            @queue.put(AsyncProducerWorker::SHUTDOWN)
+          end
+        end
+
+        # Don't interrupt workers by calling shutdown_now.
+        @worker_pool.shutdown
+        await(duration, unit) if block
+      end
+
+      # Wait for this producer to shutdown.
+      def await(duration, unit = TimeUnit::SECONDS)
+        @worker_pool.await_termination(duration, unit)
+      end
+
+      # Return the number of events currently buffered by this producer. This
+      # doesn't include any events buffered in workers that are currently on
+      # their way to Kinesis.
+      def queue_size
+        @queue.size
+      end
+
+      protected
+
+      def start
+        @workers.each do |w|
+          @worker_pool.java_send(:submit, [java.lang.Runnable.java_class], w)
+        end
+      end
+
+      def build_executor(worker_count)
+        Executors.new_fixed_thread_pool(
+          worker_count,
+          ThreadFactoryBuilder.new.set_name_format("#{stream}-producer-worker-%d").build
+        )
+      end
+    end
+  end
+end

data/lib/telekinesis/producer/async_producer_worker.rb

@@ -0,0 +1,110 @@
+module Telekinesis
+  module Producer
+    java_import java.nio.ByteBuffer
+    java_import java.util.concurrent.TimeUnit
+    java_import com.amazonaws.services.kinesis.model.PutRecordsRequest
+    java_import com.amazonaws.services.kinesis.model.PutRecordsRequestEntry
+
+    class AsyncProducerWorker
+      SHUTDOWN = :shutdown
+
+      def initialize(producer, queue, send_size, send_every, retries, retry_interval)
+        @producer = producer
+        @queue = queue
+        @send_size = send_size
+        @send_every = send_every
+        @retries = retries
+        @retry_interval = retry_interval
+
+        @stream = producer.stream                   # for convenience
+        @client = producer.client                   # for convenience
+        @failure_handler = producer.failure_handler # for convenience
+
+        @buffer = []
+        @last_poll_at = current_time_millis
+        @shutdown = false
+      end
+
+      def run
+        loop do
+          next_wait = [0, (@last_poll_at + @send_every) - current_time_millis].max
+          next_item = @queue.poll(next_wait, TimeUnit::MILLISECONDS)
+
+          if next_item == SHUTDOWN
+            next_item, @shutdown = nil, true
+          end
+
+          unless next_item.nil?
+            buffer(next_item)
+          end
+
+          if buffer_full || (next_item.nil? && buffer_has_records)
+            put_records(get_and_reset_buffer, @retries, @retry_interval)
+          end
+
+          @last_poll_at = current_time_millis
+          break if @shutdown
+        end
+      rescue => e
+        # TODO: is there a way to encourage people to set up an uncaught exception
+        # hanlder and/or disable this?
+        bt = e.backtrace ? e.backtrace.map{|l| "!  #{l}"}.join("\n") : ""
+        warn "Producer background thread died!"
+        warn "#{e.class}: #{e.message}\n#{bt}"
+        raise e
+      end
+
+      protected
+
+      def current_time_millis
+        (Time.now.to_f * 1000).to_i
+      end
+
+      def buffer(item)
+        @buffer << item
+      end
+
+      def buffer_full
+        @buffer.size == @send_size
+      end
+
+      def buffer_has_records
+        !@buffer.empty?
+      end
+
+      def get_and_reset_buffer
+        ret, @buffer = @buffer, []
+        ret
+      end
+
+      def put_records(items, retries, retry_interval)
+        begin
+          failed = []
+          while retries > 0
+            retryable, unretryable = @client.put_records(@stream, items).partition do |_, _, code, _|
+              code == 'InternalFailure' || code == 'ProvisionedThroughputExceededException'
+            end
+            failed.concat(unretryable)
+
+            if retryable.empty?
+              break
+            else
+              items = retryable.map{|k, v, _, _| [k, v]}
+              retries -= 1
+            end
+          end
+          failed.concat(retryable) unless retryable.empty?
+          @failure_handler.on_record_failure(failed) unless failed.empty?
+        rescue Telekinesis::Aws::KinesisError => e
+          if e.cause && e.cause.is_retryable && (retries -= 1) > 0
+            sleep retry_interval
+            @failure_handler.on_kinesis_retry(e, items)
+            retry
+          else
+            @failure_handler.on_kinesis_failure(e, items)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/telekinesis/producer/noop_failure_handler.rb

@@ -0,0 +1,12 @@
+module Telekinesis
+  module Producer
+    # A failure handler that does nothing.
+    #
+    # Nothing!
+    class NoopFailureHandler
+      def on_record_failure(item_error_tuples); end
+      def on_kinesis_retry(error, items); end
+      def on_kinesis_failure(error, items); end
+    end
+  end
+end

data/lib/telekinesis/producer/sync_producer.rb

@@ -0,0 +1,52 @@
+module Telekinesis
+  module Producer
+    # A synchronous Kinesis producer.
+    #
+    # This class is thread safe if and only if the underlying
+    # Telekines::Aws::Client is threadsafe. In practice, this means this client
+    # is threadsafe on JRuby and not thread safe elsewhere.
+    class SyncProducer
+      attr_reader :stream, :client
+
+      # Create a new Producer.
+      #
+      # AWS credentials may be specified by using the `:credentials` option and
+      # passing a hash containing your `:access_key_id` and `:secret_access_key`.
+      # If unspecified, credentials will be fetched from the environment, an
+      # ~/.aws/credentials file, or the current instance metadata.
+      #
+      # `:send_size` may also be used to configure the maximum batch size used
+      # in `put_all`. See `put_all` for more info.
+      def self.create(options = {})
+        stream = options[:stream]
+        client = Telekinesis::Aws::Client.build(options.fetch(:credentials, {}))
+        new(stream, client, failure_handler, options)
+      end
+
+      def initialize(stream, client, opts = {})
+        @stream = stream or raise ArgumentError, "stream may not be nil"
+        @client = client or raise ArgumentError, "client may not be nil"
+        @send_size = opts.fetch(:send_size, Telekinesis::Aws::KINESIS_MAX_PUT_RECORDS_SIZE)
+      end
+
+      # Put an individual k, v pair to Kinesis immediately. Both k and v must
+      # be strings.
+      #
+      # Returns once the call to Kinesis is complete.
+      def put(key, data)
+        @client.put_record(@stream, key, data)
+      end
+
+      # Put all of the [k, v] pairs to Kinesis in as few requests as possible.
+      # All of the ks and vs must be strings.
+      #
+      # Each request sends at most `:send_size` records. By default this is the
+      # Kinesis API limit of 500 records.
+      def put_all(items)
+        items.each_slice(@send_size).flat_map do |batch|
+          @client.put_records(@stream, batch)
+        end
+      end
+    end
+  end
+end

data/lib/telekinesis/producer/warn_failure_handler.rb

@@ -0,0 +1,25 @@
+module Telekinesis
+  module Producer
+    # A simple FailureHandler that logs errors with `warn`. Available as an
+    # example and an easy default.
+    class WarnFailureHandler
+      def on_record_failure(item_err_pairs)
+        warn "Puts for #{item_err_pairs.size} records failed!"
+      end
+
+      # Do nothing on retry. Let it figure itself out.
+      def on_kinesis_retry(err, items); end
+
+      def on_kinesis_failure(err, items)
+        warn "PutRecords request with #{items.size} items failed!"
+        warn format_bt(err)
+      end
+
+      protected
+
+      def format_bt(e)
+        e.backtrace ? e.backtrace.map{|l| "!  #{l}"}.join("\n") : ""
+      end
+    end
+  end
+end

data/lib/telekinesis/producer.rb

@@ -0,0 +1,4 @@
+require "telekinesis/producer/sync_producer"
+require "telekinesis/producer/noop_failure_handler"
+require "telekinesis/producer/warn_failure_handler"
+require "telekinesis/producer/async_producer"

data/lib/telekinesis/version.rb

@@ -0,0 +1,3 @@
+module Telekinesis
+  VERSION = '2.0.0'
+end

data/lib/telekinesis.rb

@@ -0,0 +1,14 @@
+module Telekinesis; end
+
+unless RUBY_PLATFORM.match(/java/)
+  raise "Sorry! Telekinesis is only supported on JRuby"
+end
+
+require "telekinesis/version"
+require "telekinesis/telekinesis-#{Telekinesis::VERSION}.jar"
+require "telekinesis/java_util"
+require "telekinesis/logging/java_logging"
+require "telekinesis/aws"
+
+require "telekinesis/producer"
+require "telekinesis/consumer"

data/telekinesis.gemspec

@@ -0,0 +1,21 @@
+$:.push File.expand_path("../lib", __FILE__)
+require "telekinesis/version"
+
+Gem::Specification.new do |spec|
+  spec.name         = "telekinesis"
+  spec.version      = Telekinesis::VERSION
+  spec.author       = "Ben Linsay"
+  spec.email        = "ben@kickstarter.com"
+  spec.summary      = "High level clients for Amazon Kinesis"
+  spec.homepage     = "https://github.com/kickstarter/telekinesis"
+
+  spec.platform     = "java"
+  spec.files        = `git ls-files`.split($/) + Dir.glob("lib/telekinesis/*.jar")
+  spec.require_paths = ["lib"]
+  spec.add_dependency "aws-sdk"
+
+  spec.add_development_dependency "rake"
+  spec.add_development_dependency "nokogiri"
+  spec.add_development_dependency "minitest"
+  spec.add_development_dependency "shoulda-context"
+end

data/test/aws/test_client_adapter.rb

@@ -0,0 +1,29 @@
+require_relative '../test_helper'
+
+class ClientAdapterTest < Minitest::Test
+  StubResponse = Struct.new(:error_code, :error_message)
+
+  class EvenRecordsAreErrors < Telekinesis::Aws::ClientAdapter
+    def do_put_records(stream, items)
+      items.each_with_index.map do |_, idx|
+        err, message = idx.even? ? ["error-#{idx}", "message-#{idx}"] : [nil, nil]
+        StubResponse.new(err, message)
+      end
+    end
+  end
+
+  context "ClientAdapter" do
+    context "put_records" do
+      setup do
+        @client = EvenRecordsAreErrors.new(nil)
+        @items = 10.times.map{|i| ["key-#{i}", "value-#{i}"]}
+        @expected = 10.times.select{|i| i.even?}
+                            .map{|i| ["key-#{i}", "value-#{i}", "error-#{i}", "message-#{i}"]}
+      end
+
+      should "zip error responses with records" do
+        assert(@expected, @client.put_records('stream', @items))
+      end
+    end
+  end
+end

data/test/aws/test_java_client_adapter.rb

@@ -0,0 +1,72 @@
+require_relative '../test_helper'
+
+class JavaClientAdapterTest < Minitest::Test
+  java_import com.amazonaws.services.kinesis.model.PutRecordRequest
+  java_import com.amazonaws.services.kinesis.model.PutRecordsRequest
+
+  SomeStruct = Struct.new(:field)
+  StubResponse = Struct.new(:records)
+
+  class EchoClient
+    def put_record(*args)
+      args
+    end
+
+    def put_records(*args)
+      StubResponse.new(args)
+    end
+  end
+
+  context "JavaClientAdapter" do
+    setup do
+      @client = Telekinesis::Aws::JavaClientAdapter.new(EchoClient.new)
+    end
+
+    context "#put_record" do
+      setup do
+        # No exceptions, coerced to string. [args, expected]
+        @data = [
+          [['stream', 'key', 'value'], ['stream', 'key', 'value']],
+          [['stream', 123, 123], ['stream', '123', '123']],
+          [['stream', SomeStruct.new('key'), SomeStruct.new('value')], ['stream', '#<struct JavaClientAdapterTest::SomeStruct field="key">', '#<struct JavaClientAdapterTest::SomeStruct field="value">']],
+        ]
+      end
+
+      should "generate aws.PutRecordsRequest" do
+        @data.each do |args, expected|
+          request, = @client.put_record(*args)
+          expected_stream, expected_key, expected_value = expected
+
+          assert_equal(expected_stream, request.stream_name)
+          assert_equal(expected_key, request.partition_key)
+          assert_equal(expected_value, String.from_java_bytes(request.data.array))
+        end
+      end
+    end
+
+    context "#do_put_records" do
+      setup do
+        # No exceptions, coerced to string. [args, expected]
+        @data = [
+          [
+            ['stream', [['key', 'value'], [123, 123], [SomeStruct.new('key'), SomeStruct.new('value')]]],
+            ['stream', [['key', 'value'], ['123', '123'], ['#<struct JavaClientAdapterTest::SomeStruct field="key">', '#<struct JavaClientAdapterTest::SomeStruct field="value">']]]
+          ],
+        ]
+      end
+
+      should "generate aws.PutRecordsRequest" do
+        @data.each do |args, expected|
+          request, = @client.send(:do_put_records, *args)
+          expected_stream, expected_items = expected
+
+          assert_equal(expected_stream, request.stream_name)
+          expected_items.zip(request.records) do |(expected_key, expected_value), record|
+            assert_equal(expected_key, record.partition_key)
+            assert_equal(expected_value, String.from_java_bytes(record.data.array))
+          end
+        end
+      end
+    end
+  end
+end

data/test/producer/test_async_producer.rb

@@ -0,0 +1,158 @@
+require_relative "test_helper"
+
+class AsyncProducerTest < Minitest::Test
+  java_import java.util.concurrent.TimeUnit
+  java_import java.util.concurrent.CountDownLatch
+  java_import java.util.concurrent.ArrayBlockingQueue
+
+  StubClient = Struct.new(:welp)
+
+  class LatchQueue
+    def initialize
+      @under = ArrayBlockingQueue.new(100)
+      @latch = CountDownLatch.new(1)
+      @putting = CountDownLatch.new(1)
+    end
+
+    def count_down
+      @latch.count_down
+    end
+
+    def wait_for_put
+      @putting.await
+    end
+
+    def put(item)
+      @putting.count_down
+      @latch.await
+      @under.put(item)
+    end
+  end
+
+  def build_producer
+    opts = {
+      queue: @queue,
+      manual_start: true,
+      worker_count: @worker_count,
+    }
+    Telekinesis::Producer::AsyncProducer.new(
+      @stream,
+      StubClient.new,
+      Telekinesis::Producer::NoopFailureHandler.new,
+      opts
+    )
+  end
+
+  context "AsyncProducer" do
+    setup do
+      @stream = 'test'  # ignored
+      @worker_count = 3 # arbitrary
+    end
+
+    context "put" do
+      setup do
+        @queue = ArrayBlockingQueue.new(100)
+        build_producer.put("hi", "there")
+      end
+
+      should "add the k,v pair to the queue" do
+        assert_equal([["hi", "there"]], @queue.to_a)
+      end
+    end
+
+    context "put_all" do
+      setup do
+        @items = 10.times.map{|i| ["key-#{i}", "value-#{i}"]}
+        @queue = ArrayBlockingQueue.new(100)
+        build_producer.put_all(@items)
+      end
+
+      should "add all items to the queue" do
+        assert_equal(@items, @queue.to_a)
+      end
+    end
+
+    context "after shutdown" do
+      setup do
+        @queue = ArrayBlockingQueue.new(100)
+        @producer = build_producer
+        @producer.shutdown
+      end
+
+      should "shutdown all workers" do
+        assert_equal([Telekinesis::Producer::AsyncProducerWorker::SHUTDOWN] * @worker_count, @queue.to_a)
+      end
+
+      should "not accept events while shut down" do
+        refute(@producer.put("key", "value"))
+      end
+    end
+
+    context "with a put in progress" do
+      setup do
+        @queue = LatchQueue.new
+        @producer = build_producer
+
+        # Thread blocks waiting for the latch in LatchQueue. Don't do any other
+        # set up until this thread is in the critical section.
+        Thread.new do
+          @producer.put("k", "v")
+        end
+        @queue.wait_for_put
+
+        # Thread blocks waiting for the write_lock in AsyncProducer. Once it's
+        # unblocked it signals by counting down shutdown_latch.
+        @shutdown_latch = CountDownLatch.new(1)
+        Thread.new do
+          @producer.shutdown
+          @shutdown_latch.count_down
+        end
+      end
+
+      should "block on shutdown until the put is done" do
+        # Check that the latch hasn't been triggered yet. Return immediately
+        # from the check - don't bother waiting.
+        refute(@shutdown_latch.await(0, TimeUnit::MILLISECONDS))
+        @queue.count_down
+        # NOTE: The assert is here to fail the test if it times out. This could
+        # effectively just be an await with no duration.
+        assert(@shutdown_latch.await(2, TimeUnit::SECONDS))
+      end
+    end
+
+    context "with a shutdown in progress" do
+      setup do
+        @queue = LatchQueue.new
+        @producer = build_producer
+
+        # Thread blocks waiting to insert :shutdown into the queue because of
+        # the latch in LatchQueue. Don't do any other test set up until this
+        # thread is in the critical section.
+        Thread.new do
+          @producer.shutdown
+        end
+        @queue.wait_for_put
+
+        # This thread blocks waiting for the lock in AsyncProducer. Once it's
+        # done the put continues and then it signals completion by counting
+        # down finished_put_latch.
+        @finished_put_latch = CountDownLatch.new(1)
+        Thread.new do
+          @put_result = @producer.put("k", "v")
+          @finished_put_latch.count_down
+        end
+      end
+
+      should "block on a put" do
+        # Thread is already waiting in the critical section. Just check that
+        # the call hasn't exited yet and return immediately.
+        refute(@finished_put_latch.await(0, TimeUnit::MILLISECONDS))
+        @queue.count_down
+        # NOTE: The assert is here to fail the test if it times out. This could
+        # effectively just be an await with no duration.
+        assert(@finished_put_latch.await(2, TimeUnit::SECONDS))
+        refute(@put_result, "Producer should reject a put after shutdown")
+      end
+    end
+  end
+end