telekinesis 2.0.0-java

data/ext/src/main/java/com/kickstarter/jruby/Telekinesis.java

@@ -0,0 +1,103 @@
+package com.kickstarter.jruby;
+
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer;
+import com.amazonaws.services.kinesis.clientlibrary.lib.worker.KinesisClientLibConfiguration;
+import com.amazonaws.services.kinesis.clientlibrary.lib.worker.Worker;
+import com.amazonaws.services.kinesis.clientlibrary.types.ShutdownReason;
+import com.amazonaws.services.kinesis.model.Record;
+
+import java.util.List;
+
+/**
+ * A shim that makes it possible to use the Kinesis Client Library from JRuby.
+ * Without the shim, {@code initialize} method in
+ * {@link com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor}
+ * conflicts with the special {@code initialize} method in Ruby. The shim
+ * interface renames {@code initialize} to {@code init}.
+ * <p />
+ *
+ * For convenience a {@link #newWorker(KinesisClientLibConfiguration, IRecordProcessorFactory)}
+ * method is provided, so you can use closure conversion in JRuby to specify an
+ * {@link IRecordProcessorFactory}. For example
+ *
+ * <p />
+ *
+ * <pre>
+ *     some_thing = ...
+ *
+ *     com.kickstarter.jruby.Telekinesis.new_worker(my_config) do
+ *       MyRecordProcessor.new(some_thing, some_other_thing)
+ *     end
+ * </pre>
+ */
+public class Telekinesis {
+    /**
+     * Create a new KCL {@link Worker} that processes records using the given
+     * {@link IRecordProcessorFactory}.
+     */
+    public static Worker newWorker(final KinesisClientLibConfiguration config, final IRecordProcessorFactory factory) {
+        return new Worker(new com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorFactory() {
+            @Override
+            public com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor createProcessor() {
+                return new RecordProcessorShim(factory.createProcessor());
+            }
+        }, config);
+    }
+
+    // ========================================================================
+    /**
+     * A shim that wraps a {@link IRecordProcessor} so it can get used by the KCL.
+     */
+    private static class RecordProcessorShim implements com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor {
+        private final IRecordProcessor underlying;
+
+        public RecordProcessorShim(final IRecordProcessor underlying) { this.underlying = underlying; }
+
+        @Override
+        public void initialize(final String shardId) {
+            underlying.init(shardId);
+        }
+
+        @Override
+        public void processRecords(final List<Record> records, final IRecordProcessorCheckpointer checkpointer) {
+            underlying.processRecords(records, checkpointer);
+        }
+
+        @Override
+        public void shutdown(final IRecordProcessorCheckpointer checkpointer, final ShutdownReason reason) {
+            underlying.shutdown(checkpointer, reason);
+        }
+    }
+
+    /**
+     * A parallel {@link com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor}
+     * that avoids naming conflicts with reserved words in Ruby.
+     */
+    public static interface IRecordProcessor {
+        /**
+         * @see com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor#initialize(String)
+         */
+        void init(final String shardId);
+
+        /**
+         * @see com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor#processRecords(List, IRecordProcessorCheckpointer)
+         */
+        void processRecords(List<Record> records, IRecordProcessorCheckpointer checkpointer);
+
+        /**
+         * @see com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor#shutdown(IRecordProcessorCheckpointer, ShutdownReason)
+         */
+        void shutdown(IRecordProcessorCheckpointer checkpointer, ShutdownReason reason);
+    }
+
+    /**
+     * A parallel {@link com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorFactory}
+     * for {@link IRecordProcessor}.
+     */
+    public static interface IRecordProcessorFactory {
+        /**
+         * @see com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorFactory#createProcessor()
+         */
+        IRecordProcessor createProcessor();
+    }
+}

data/lib/telekinesis/aws/client_adapter.rb

@@ -0,0 +1,61 @@
+module Telekinesis
+  module Aws
+    # NOTE: wrapping the cause is necessary since JRuby isn't 2.1 compatible (yet)
+    class KinesisError < RuntimeError
+      attr_reader :cause
+
+      def initialize(cause)
+        @cause = cause
+      end
+    end
+
+    # Base class for other ClientAdapters. Client adapters exist to make
+    # switching between platforms easy and painless.
+    #
+    # The base adapter defines the interface and provides convience methods.
+    class ClientAdapter
+      # Build a new client given AWS credentials.
+      #
+      # Credentials must be supplied as a hash that contains symbolized
+      # :access_key_id and :secret_access_key keys.
+      def self.build(credentials)
+        raise NotImplementedError
+      end
+
+      def initialize(client)
+        @client = client
+      end
+
+      # Make a put_record call to the underlying client. Must return an object
+      # that responds to `shard_id` and `sequence_number`.
+      def put_record(stream, key, value)
+        raise NotImplementedError
+      end
+
+      # Make a put_records call to the underlying client. If the request
+      # succeeds but returns errors for some records, the original [key, value]
+      # pair is zipped with the [error_code, error_message] pair and the
+      # offending records are returned.
+      def put_records(stream, items)
+        response = do_put_records(stream, items)
+        failures = items.zip(response).reject{|_, r| r.error_code.nil?}
+
+        failures.map do |(k, v), r|
+          [k, v, r.error_code, r.error_message]
+        end
+      end
+
+      protected
+
+      # Put an enumerable of [key, value] pairs to the given stream. Returns an
+      # enumerable of response objects the same size as the given list of items.
+      #
+      # Response objects must respond to `error_code` and `error_message`. Any
+      # response with a nil error_code is considered successful.
+      def do_put_records(stream, items)
+        raise NotImplementedError
+      end
+    end
+  end
+end
+

data/lib/telekinesis/aws/java_client_adapter.rb

@@ -0,0 +1,72 @@
+module Telekinesis
+  module Aws
+    java_import java.nio.ByteBuffer
+    java_import com.amazonaws.AmazonClientException
+    java_import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
+    java_import com.amazonaws.services.kinesis.AmazonKinesisClient
+    java_import com.amazonaws.services.kinesis.model.PutRecordRequest
+    java_import com.amazonaws.services.kinesis.model.PutRecordsRequest
+    java_import com.amazonaws.services.kinesis.model.PutRecordsRequestEntry
+
+    # A ClientAdapter that wraps the AWS Java SDK.
+    #
+    # Since the underlying Java client is thread safe, this adapter is thread
+    # safe.
+    class JavaClientAdapter < ClientAdapter
+      # Build a new client adapter. `credentials` is a hash keyed with
+      # `:access_key_id` and `:secret_access_key`. If this hash is left blank
+      # (the default) the client uses the DefaultAWSCredentialsProviderChain to
+      # look for credentials.
+      def self.build(credentials = {})
+        client = AmazonKinesisClient.new(build_credentials_provider(credentials))
+        new(client)
+      end
+
+      def self.build_credentials_provider(credentials)
+        if credentials.empty?
+          DefaultAWSCredentialsProviderChain.new
+        else
+          StaticCredentialsProvider.new(
+            BasicAWSCredentials.new(
+              credentials[:access_key_id],
+              credentials[:secret_access_key]
+            )
+          )
+        end
+      end
+
+      def put_record(stream, key, value)
+        r = PutRecordRequest.new.tap do |request|
+          request.stream_name = stream
+          request.partition_key = key.to_s
+          request.data = ByteBuffer.wrap(value.to_s.to_java_bytes)
+        end
+        @client.put_record(r)
+      rescue AmazonClientException => e
+        raise KinesisError.new(e)
+      end
+
+      protected
+
+      def do_put_records(stream, items)
+        result = @client.put_records(build_put_records_request(stream, items))
+        result.records
+      rescue AmazonClientException => e
+        raise KinesisError.new(e)
+      end
+
+      def build_put_records_request(stream, items)
+        entries = items.map do |key, value|
+          PutRecordsRequestEntry.new.tap do |entry|
+            entry.partition_key = key.to_s
+            entry.data = ByteBuffer.wrap(value.to_s.to_java_bytes)
+          end
+        end
+        PutRecordsRequest.new.tap do |request|
+          request.stream_name = stream
+          request.records = entries
+        end
+      end
+    end
+  end
+end

data/lib/telekinesis/aws/ruby_client_adapter.rb

@@ -0,0 +1,40 @@
+module Telekinesis
+  module Aws
+    # A ClientAdapter that wraps the ruby aws-sdk gem (version 2).
+    #
+    # Since the aws-sdk gem does not appear to be thread-safe, this adapter
+    # should not be considered thread safe.
+    class RubyClientAdapter < ClientAdapter
+      # Build a new client adapter. Credentials are passed directly to the
+      # constructor for Aws::Kinesis::Client.
+      #
+      # See: http://docs.aws.amazon.com/sdkforruby/api/Aws/Kinesis/Client.html
+      def self.build(credentials)
+        new(Aws::Kinesis::Client.new(credentials))
+      end
+
+      def put_record(stream, key, value)
+        @client.put_record(stream: stream, partition_key: key, data: value)
+      rescue Aws::Errors::ServiceError => e
+        raise KinesisError.new(e)
+      end
+
+      protected
+
+      def do_put_records(stream, items)
+        @client.put_records(build_put_records_request(stream, items)).flat_map do |page|
+          page.records
+        end
+      rescue Aws::Errors::ServiceError => e
+        raise KinesisError.new(e)
+      end
+
+      def build_put_records_request(stream, items)
+        {
+          stream: stream,
+          records: items.map{|k, v| {partition_key: k, data: v}}
+        }
+      end
+    end
+  end
+end

data/lib/telekinesis/aws.rb

@@ -0,0 +1,9 @@
+require "telekinesis/aws/client_adapter.rb"
+require "telekinesis/aws/java_client_adapter"
+
+module Telekinesis
+  module Aws
+    KINESIS_MAX_PUT_RECORDS_SIZE = 500
+    Client = JavaClientAdapter
+  end
+end

data/lib/telekinesis/consumer/base_processor.rb

@@ -0,0 +1,12 @@
+module Telekinesis
+  module Consumer
+    # A RecordProcessor with no-op implementations of all of the required
+    # IRecordProcessor methods. Override it to implement simple IRecordProcessors
+    # that don't need to do anything special on init or shutdown.
+    class BaseProcessor
+      def init(shard_id); end
+      def process_records(records, checkpointer); end
+      def shutdown(checkpointer, reason); end
+    end
+  end
+end

data/lib/telekinesis/consumer/block.rb

@@ -0,0 +1,22 @@
+module Telekinesis
+  module Consumer
+    # A RecordProcessor that uses the given block to process records. Useful to
+    # quickly define a consumer.
+    #
+    # Telekinesis::Consumer::Worker.new(stream: 'my-stream', app: 'tail') do
+    #   Telekinesis::Consumer::Block.new do |records, checkpointer|
+    #     records.each {|r| puts r}
+    #   end
+    # end
+    class Block < BaseProcessor
+      def initialize(&block)
+        raise ArgumentError, "No block given" unless block_given?
+        @block = block
+      end
+
+      def process_records(records, checkpointer)
+        @block.call(records, checkpointer)
+      end
+    end
+  end
+end

data/lib/telekinesis/consumer/distributed_consumer.rb

@@ -0,0 +1,114 @@
+module Telekinesis
+  module Consumer
+    java_import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStream
+    java_import com.amazonaws.services.kinesis.clientlibrary.lib.worker.KinesisClientLibConfiguration
+
+    class DistributedConsumer
+      # Create a new consumer that consumes data from a Kinesis stream.
+      # DistributedConsumers use DynamoDB to register as part of the same
+      # application and evenly distribute work between them. See the
+      # AWS Docs for more information:
+      #
+      # http://docs.aws.amazon.com/kinesis/latest/dev/developing-consumer-apps-with-kcl.html
+      #
+      # DistributedConsumers are configured with a hash. The Kinesis `:stream`
+      # to consume from is required.
+      #
+      # DistribtuedConsumers operate in groups. All consumers with the same
+      # `:app` id use dynamo to attempt to distribute work evenly among
+      # themselves. The `:worker_id` is used to distinguish individual clients
+      # (`:worker_id` defaults to the current hostname. If you plan to run more
+      # than one DistributedConsumer in the same `:app` per host, make sure you
+      # set this to something unique!).
+      #
+      # Any other valid KCL Worker `:options` may be passed as a hash.
+      #
+      # For example, to configure a `tail` app on `some-stream` and use the
+      # default `:worker_id`, you might pass the following configuration to your
+      # DistributedConsumer.
+      #
+      #     config = {
+      #       app: 'tail',
+      #       stream: 'some-stream',
+      #       options: {initial_position_in_stream: 'TRIM_HORIZON'}
+      #     }
+      #
+      # To actually process the stream, a DistribtuedConsumer creates
+      # record processors. These are objects that correspond to the KCL's
+      # RecordProcessor interface - processors must implement `init`,
+      # `process_records`, and `shutdown` methods.
+      #
+      # http://docs.aws.amazon.com/kinesis/latest/dev/kinesis-record-processor-implementation-app-java.html#kinesis-record-processor-implementation-interface-java
+      #
+      # To specify which record processor to create, pass a block to your
+      # distribtued consumer that returns a new record processor. This block
+      # may (nay, WILL) be called from a background thread so make sure that
+      # it's thread-safe.
+      #
+      # Telekinesis provides a BaseProcessor that implements no-op versions
+      # of all of the required methods to make writing quick processors easier
+      # and a Block processor that executes the specified block every time
+      # `process_records` is called.
+      #
+      # To write a stream tailer, you might use Block as follows:
+      #
+      #     Telekinesis::Consumer::DistributedConsumer.new(config) do
+      #       Telekinesis::Consumer::Block do |records, _|
+      #         records.each {|r| puts r}
+      #       end
+      #     end
+      #
+      def initialize(config, &block)
+        raise ArgumentError, "No block given!" unless block_given?
+        kcl_config = self.class.build_config(config)
+        @under = com.kickstarter.jruby.Telekinesis.new_worker(kcl_config, &block)
+      end
+
+      # Return the underlying KCL worker. It's a java.lang.Runnable.
+      def as_runnable
+        @under
+      end
+
+      # Start the KCL worker. If background is set to `true`, the worker is
+      # started in its own JRuby Thread and the Thread is returned. Otherwise,
+      # starts in the current thread and returns nil.
+      def run(background = false)
+        if background
+          Thread.new { @under.run }
+        else
+          @under.run
+        end
+      end
+
+      protected
+
+      def self.build_config(config)
+        creds_hash = config.fetch(:credentials, {})
+        credentials_provider = Telekinesis::Aws::JavaClientAdapter.build_credentials_provider(creds_hash)
+
+        # App and Stream are mandatory.
+        app, stream = [:app, :stream].map do |k|
+          raise ArgumentError, "#{k} is required" unless config.include?(k)
+          config[k]
+        end
+        # Use this host as the worker_id by default.
+        worker_id = config.fetch(:worker_id, `hostname`.chomp)
+
+        KinesisClientLibConfiguration.new(app, stream, credentials_provider, worker_id).tap do |kcl_config|
+          config.fetch(:options, {}).each do |k, v|
+            # Handle initial position in stream separately. It's the only option
+            # that requires a value conversion.
+            if k.to_s == 'initial_position_in_stream'
+              kcl_config.with_initial_position_in_stream(InitialPositionInStream.value_of(v))
+            else
+              setter = "with_#{k}".to_sym
+              if kcl_config.respond_to?(setter)
+                kcl_config.send(setter, v)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/telekinesis/consumer.rb

@@ -0,0 +1,3 @@
+require "telekinesis/consumer/distributed_consumer"
+require "telekinesis/consumer/base_processor"
+require "telekinesis/consumer/block"

data/lib/telekinesis/java_util.rb

@@ -0,0 +1,46 @@
+module Telekinesis
+  module JavaUtil
+    java_import java.util.concurrent.locks.ReentrantReadWriteLock
+
+    # Sugar around java.util.concurrent.ReentrantReadWriteLock so that it's
+    # easy to use with blocks.
+    #
+    # e.g.
+    #
+    # lock = ReentrantReadWriteLock.new
+    # some_value = 12345
+    #
+    # # In a reader thread
+    # lock.read_lock do
+    #  # Read some data! This won't block any other calls to read_lock, but will
+    #  # block if another thread is in a section guarded by write_lock.
+    # end
+    #
+    # # In a writer thread
+    # lock.write_lock do
+    #   # Write some data! This is exclusive with *any* other code guarded by
+    #   # either read_lock or write_lock.
+    # end
+    class ReadWriteLock
+      def initialize(fair = false)
+        lock = ReentrantReadWriteLock.new(fair)
+        @read = lock.read_lock
+        @write = lock.write_lock
+      end
+
+      def read_lock
+        @read.lock_interruptibly
+        yield
+      ensure
+        @read.unlock
+      end
+
+      def write_lock
+        @write.lock_interruptibly
+        yield
+      ensure
+        @write.unlock
+      end
+    end
+  end
+end

data/lib/telekinesis/logging/java_logging.rb

@@ -0,0 +1,18 @@
+require "logger"
+require "telekinesis/logging/ruby_logger_handler"
+
+module Telekinesis
+  module Logging
+    java_import java.util.logging.Logger
+    java_import java.util.logging.LogManager
+
+    def self.capture_java_logging(logger)
+      LogManager.log_manager.reset
+      Logger.get_logger("").add_handler(RubyLoggerHandler.create(logger))
+    end
+
+    def self.disable_java_logging
+      LogManager.log_manager.reset
+    end
+  end
+end

data/lib/telekinesis/logging/ruby_logger_handler.rb

@@ -0,0 +1,54 @@
+module Telekinesis
+  module Logging
+    java_import java.util.logging.Level
+    java_import java.util.logging.Handler
+
+    # A java logging Handler that delegates to a Ruby logger. The name of the
+    # j.u.l. logger is used as the progname argument to Logger.add.
+    #
+    # The translation between j.u.l. serverity levels and Ruby Logger levels
+    # isn't exact.
+    class RubyLoggerHandler < Handler
+      # NOTE: Since this class overrides a Java class, we have to use the Java
+      # constructor and set the logger after instantiation. (Overriding in
+      # JRuby is weird). Use this method to create a new logger that delegates
+      # to the passed logger.
+      def self.create(logger)
+        new.tap do |s|
+          s.set_logger(logger)
+        end
+      end
+
+      SEVERITY = {
+        # NOTE: There's no Java equivalent of FATAL.
+        Level::SEVERE => Logger::ERROR,
+        Level::WARNING => Logger::WARN,
+        Level::INFO => Logger::INFO,
+        Level::CONFIG => Logger::INFO,
+        Level::FINE=> Logger::DEBUG,
+        Level::FINER=> Logger::DEBUG,
+        Level::FINEST=> Logger::DEBUG,
+      }
+
+      def set_logger(l)
+        @logger = l
+      end
+
+      def close
+        @logger.close
+      end
+
+      # Ruby's logger has no flush method.
+      def flush; end
+
+      def publish(log_record)
+        message = if log_record.thrown.nil?
+          log_record.message
+        else
+          "#{log_record.message}: #{log_record.thrown}"
+        end
+        @logger.add(SEVERITY[log_record.level], message, log_record.logger_name)
+      end
+    end
+  end
+end