ruby-kafka-aws-iam 1.4.1

data/Rakefile

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require "bundler/gem_tasks"
+require "rspec/core/rake_task"
+
+RSpec::Core::RakeTask.new(:spec)
+
+task default: :spec

data/benchmarks/message_encoding.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "kafka"
+
+ready "message serialization" do
+  before do
+    message = Kafka::Protocol::Message.new(
+      value: "hello",
+      key: "world",
+    )
+
+    @io = StringIO.new
+    encoder = Kafka::Protocol::Encoder.new(@io)
+    message.encode(encoder)
+
+    @decoder = Kafka::Protocol::Decoder.new(@io)
+  end
+
+  go "decoding" do
+    @io.rewind
+    Kafka::Protocol::Message.decode(@decoder)
+  end
+end

data/bin/console

@@ -0,0 +1,8 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "kafka"
+
+require "pry"
+Pry.start

data/bin/setup

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -euo pipefail
+IFS=$'\n\t'
+
+bundle install

data/docker-compose.yml

@@ -0,0 +1,39 @@
+version: '2'
+services:
+  zookeeper:
+    image: wurstmeister/zookeeper
+    ports:
+      - "2181:2181"
+  kafka1:
+    image: wurstmeister/kafka:0.11.0.1
+    ports:
+      - "9092:9092"
+    environment:
+      KAFKA_BROKER_ID: 1
+      KAFKA_ADVERTISED_HOST_NAME: 192.168.99.100
+      KAFKA_ADVERTISED_PORT: 9092
+      KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock
+  kafka2:
+    image: wurstmeister/kafka:0.11.0.1
+    ports:
+      - "9093:9092"
+    environment:
+      KAFKA_BROKER_ID: 2
+      KAFKA_ADVERTISED_HOST_NAME: 192.168.99.100
+      KAFKA_ADVERTISED_PORT: 9093
+      KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock
+  kafka3:
+    image: wurstmeister/kafka:0.11.0.1
+    ports:
+      - "9094:9092"
+    environment:
+      KAFKA_BROKER_ID: 3
+      KAFKA_ADVERTISED_HOST_NAME: 192.168.99.100
+      KAFKA_ADVERTISED_PORT: 9094
+      KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock

data/examples/consumer-group.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+$LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
+
+require "kafka"
+
+logger = Logger.new(STDOUT)
+brokers = ENV.fetch("KAFKA_BROKERS", "localhost:9092").split(",")
+
+# Make sure to create this topic in your Kafka cluster or configure the
+# cluster to auto-create topics.
+topic = "text"
+
+kafka = Kafka.new(
+  seed_brokers: brokers,
+  client_id: "test",
+  socket_timeout: 20,
+  logger: logger,
+)
+
+consumer = kafka.consumer(group_id: "test")
+consumer.subscribe(topic)
+
+trap("TERM") { consumer.stop }
+trap("INT") { consumer.stop }
+
+begin
+  consumer.each_message do |message|
+  end
+rescue Kafka::ProcessingError => e
+  warn "Got #{e.cause}"
+  consumer.pause(e.topic, e.partition, timeout: 20)
+
+  retry
+end

data/examples/firehose-consumer.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+$LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
+
+require "kafka"
+require "dotenv"
+
+Dotenv.load
+
+KAFKA_CLIENT_CERT = ENV.fetch("KAFKA_CLIENT_CERT")
+KAFKA_CLIENT_CERT_KEY = ENV.fetch("KAFKA_CLIENT_CERT_KEY")
+KAFKA_SERVER_CERT = ENV.fetch("KAFKA_SERVER_CERT")
+KAFKA_URL = ENV.fetch("KAFKA_URL")
+KAFKA_BROKERS = KAFKA_URL
+KAFKA_TOPIC = "test-messages"
+
+NUM_THREADS = 4
+
+queue = Queue.new
+
+threads = NUM_THREADS.times.map do |worker_id|
+  Thread.new do
+    logger = Logger.new($stderr)
+    logger.level = Logger::INFO
+
+    logger.formatter = proc {|severity, datetime, progname, msg|
+      "[#{worker_id}] #{severity.ljust(5)} -- #{msg}\n"
+    }
+
+    kafka = Kafka.new(
+      seed_brokers: KAFKA_BROKERS,
+      logger: logger,
+      connect_timeout: 30,
+      socket_timeout: 30,
+      ssl_client_cert: KAFKA_CLIENT_CERT,
+      ssl_client_cert_key: KAFKA_CLIENT_CERT_KEY,
+      ssl_ca_cert: KAFKA_SERVER_CERT,
+    )
+
+    consumer = kafka.consumer(group_id: "firehose")
+    consumer.subscribe(KAFKA_TOPIC)
+
+    i = 0
+    consumer.each_message do |message|
+      i += 1
+
+      if i % 1000 == 0
+        queue << i
+        i = 0
+      end
+
+      sleep 0.01
+    end
+  end
+end
+
+threads.each {|t| t.abort_on_exception = true }
+
+received_messages = 0
+
+loop do
+  received_messages += queue.pop
+  puts "===> Received #{received_messages} messages"
+end

data/examples/firehose-producer.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+$LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
+
+require "kafka"
+require "dotenv"
+
+Dotenv.load
+
+KAFKA_CLIENT_CERT = ENV.fetch("KAFKA_CLIENT_CERT")
+KAFKA_CLIENT_CERT_KEY = ENV.fetch("KAFKA_CLIENT_CERT_KEY")
+KAFKA_SERVER_CERT = ENV.fetch("KAFKA_SERVER_CERT")
+KAFKA_URL = ENV.fetch("KAFKA_URL")
+KAFKA_BROKERS = KAFKA_URL
+KAFKA_TOPIC = "test-messages"
+
+NUM_THREADS = 20
+
+threads = NUM_THREADS.times.map do
+  Thread.new do
+    logger = Logger.new($stderr)
+    logger.level = Logger::INFO
+
+    kafka = Kafka.new(
+      seed_brokers: KAFKA_BROKERS,
+      logger: logger,
+      ssl_client_cert: KAFKA_CLIENT_CERT,
+      ssl_client_cert_key: KAFKA_CLIENT_CERT_KEY,
+      ssl_ca_cert: KAFKA_SERVER_CERT,
+    )
+
+    producer = kafka.async_producer(
+      delivery_interval: 1,
+      max_queue_size: 5_000,
+      max_buffer_size: 10_000,
+    )
+
+    begin
+      loop do
+        producer.produce(rand.to_s, key: rand.to_s, topic: KAFKA_TOPIC)
+      end
+    rescue Kafka::BufferOverflow
+      logger.error "Buffer overflow, backing off for 1s"
+      sleep 1
+      retry
+    ensure
+      producer.shutdown
+    end
+  end
+end
+
+threads.each {|t| t.abort_on_exception = true }
+
+threads.map(&:join)

data/examples/simple-consumer.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+# Consumes lines from a Kafka partition and writes them to STDOUT.
+#
+# You need to define the environment variable KAFKA_BROKERS for this
+# to work, e.g.
+#
+#     export KAFKA_BROKERS=localhost:9092
+#
+
+$LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
+
+require "kafka"
+
+# We don't want log output to clutter the console. Replace `StringIO.new`
+# with e.g. `$stderr` if you want to see what's happening under the hood.
+logger = Logger.new(StringIO.new)
+
+brokers = ENV.fetch("KAFKA_BROKERS").split(",")
+
+# Make sure to create this topic in your Kafka cluster or configure the
+# cluster to auto-create topics.
+topic = "text"
+
+kafka = Kafka.new(
+  seed_brokers: brokers,
+  client_id: "simple-consumer",
+  socket_timeout: 20,
+  logger: logger,
+)
+
+kafka.each_message(topic: topic) do |message|
+  puts message.value
+end

data/examples/simple-producer.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+# Reads lines from STDIN, writing them to Kafka.
+#
+# You need to define the environment variable KAFKA_BROKERS for this
+# to work, e.g.
+#
+#     export KAFKA_BROKERS=localhost:9092
+#
+
+$LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
+
+require "kafka"
+
+logger = Logger.new($stderr)
+brokers = ENV.fetch("KAFKA_BROKERS")
+
+# Make sure to create this topic in your Kafka cluster or configure the
+# cluster to auto-create topics.
+topic = "text"
+
+kafka = Kafka.new(
+  seed_brokers: brokers,
+  client_id: "simple-producer",
+  logger: logger,
+)
+
+producer = kafka.producer
+
+begin
+  $stdin.each_with_index do |line, index|
+    producer.produce(line, topic: topic)
+
+    # Send messages for every 10 lines.
+    producer.deliver_messages if index % 10 == 0
+  end
+ensure
+  # Make sure to send any remaining messages.
+  producer.deliver_messages
+
+  producer.shutdown
+end

data/examples/ssl-producer.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+# Reads lines from STDIN, writing them to Kafka.
+
+$LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
+
+require "kafka"
+
+logger = Logger.new($stderr)
+brokers = ENV.fetch("KAFKA_BROKERS")
+
+# Make sure to create this topic in your Kafka cluster or configure the
+# cluster to auto-create topics.
+topic = "page-visits"
+
+ssl_context = OpenSSL::SSL::SSLContext.new
+ssl_context.set_params(
+  cert: OpenSSL::X509::Certificate.new(ENV.fetch("KAFKA_CLIENT_CERT")),
+  key: OpenSSL::PKey::RSA.new(ENV.fetch("KAFKA_CLIENT_CERT_KEY")),
+)
+
+kafka = Kafka.new(
+  seed_brokers: brokers,
+  client_id: "ssl-producer",
+  logger: logger,
+  ssl: true,
+  ssl_context: ssl_context,
+)
+
+producer = kafka.producer
+
+begin
+  $stdin.each_with_index do |line, index|
+    producer.produce(line, topic: topic)
+
+    # Send messages for every 10 lines.
+    producer.deliver_messages if index % 10 == 0
+  end
+ensure
+  # Make sure to send any remaining messages.
+  producer.deliver_messages
+
+  producer.shutdown
+end

data/lib/kafka/async_producer.rb

@@ -0,0 +1,297 @@
+# frozen_string_literal: true
+
+require "thread"
+
+module Kafka
+
+  # A Kafka producer that does all its work in the background so as to not block
+  # the calling thread. Calls to {#deliver_messages} are asynchronous and return
+  # immediately.
+  #
+  # In addition to this property it's possible to define automatic delivery
+  # policies. These allow placing an upper bound on the number of buffered
+  # messages and the time between message deliveries.
+  #
+  # * If `delivery_threshold` is set to a value _n_ higher than zero, the producer
+  #   will automatically deliver its messages once its buffer size reaches _n_.
+  # * If `delivery_interval` is set to a value _n_ higher than zero, the producer
+  #   will automatically deliver its messages every _n_ seconds.
+  #
+  # By default, automatic delivery is disabled and you'll have to call
+  # {#deliver_messages} manually.
+  #
+  # ## Buffer Overflow and Backpressure
+  #
+  # The calling thread communicates with the background thread doing the actual
+  # work using a thread safe queue. While the background thread is busy delivering
+  # messages, new messages will be buffered in the queue. In order to avoid
+  # the queue growing uncontrollably in cases where the background thread gets
+  # stuck or can't follow the pace of the calling thread, there's a maximum
+  # number of messages that is allowed to be buffered. You can configure this
+  # value by setting `max_queue_size`.
+  #
+  # If you produce messages faster than the background producer thread can
+  # deliver them to Kafka you will eventually fill the producer's buffer. Once
+  # this happens, the background thread will stop popping messages off the
+  # queue until it can successfully deliver the buffered messages. The queue
+  # will therefore grow in size, potentially hitting the `max_queue_size` limit.
+  # Once this happens, calls to {#produce} will raise a {BufferOverflow} error.
+  #
+  # Depending on your use case you may want to slow down the rate of messages
+  # being produced or perhaps halt your application completely until the
+  # producer can deliver the buffered messages and clear the message queue.
+  #
+  # ## Example
+  #
+  #     producer = kafka.async_producer(
+  #       # Keep at most 1.000 messages in the buffer before delivering:
+  #       delivery_threshold: 1000,
+  #
+  #       # Deliver messages every 30 seconds:
+  #       delivery_interval: 30,
+  #     )
+  #
+  #     # There's no need to manually call #deliver_messages, it will happen
+  #     # automatically in the background.
+  #     producer.produce("hello", topic: "greetings")
+  #
+  #     # Remember to shut down the producer when you're done with it.
+  #     producer.shutdown
+  #
+  class AsyncProducer
+    # Initializes a new AsyncProducer.
+    #
+    # @param sync_producer [Kafka::Producer] the synchronous producer that should
+    #   be used in the background.
+    # @param max_queue_size [Integer] the maximum number of messages allowed in
+    #   the queue.
+    # @param delivery_threshold [Integer] if greater than zero, the number of
+    #   buffered messages that will automatically trigger a delivery.
+    # @param delivery_interval [Integer] if greater than zero, the number of
+    #   seconds between automatic message deliveries.
+    #
+    def initialize(sync_producer:, max_queue_size: 1000, delivery_threshold: 0, delivery_interval: 0, max_retries: -1, retry_backoff: 0, instrumenter:, logger:)
+      raise ArgumentError unless max_queue_size > 0
+      raise ArgumentError unless delivery_threshold >= 0
+      raise ArgumentError unless delivery_interval >= 0
+
+      @queue = Queue.new
+      @max_queue_size = max_queue_size
+      @instrumenter = instrumenter
+      @logger = TaggedLogger.new(logger)
+
+      @worker = Worker.new(
+        queue: @queue,
+        producer: sync_producer,
+        delivery_threshold: delivery_threshold,
+        max_retries: max_retries,
+        retry_backoff: retry_backoff,
+        instrumenter: instrumenter,
+        logger: logger
+      )
+
+      # The timer will no-op if the delivery interval is zero.
+      @timer = Timer.new(queue: @queue, interval: delivery_interval)
+
+      @thread_mutex = Mutex.new
+    end
+
+    # Produces a message to the specified topic.
+    #
+    # @see Kafka::Producer#produce
+    # @param (see Kafka::Producer#produce)
+    # @raise [BufferOverflow] if the message queue is full.
+    # @return [nil]
+    def produce(value, topic:, **options)
+      # We want to fail fast if `topic` isn't a String
+      topic = topic.to_str
+
+      ensure_threads_running!
+
+      if @queue.size >= @max_queue_size
+        buffer_overflow topic,
+          "Cannot produce to #{topic}, max queue size (#{@max_queue_size} messages) reached"
+      end
+
+      args = [value, **options.merge(topic: topic)]
+      @queue << [:produce, args]
+
+      @instrumenter.instrument("enqueue_message.async_producer", {
+        topic: topic,
+        queue_size: @queue.size,
+        max_queue_size: @max_queue_size,
+      })
+
+      nil
+    end
+
+    # Asynchronously delivers the buffered messages. This method will return
+    # immediately and the actual work will be done in the background.
+    #
+    # @see Kafka::Producer#deliver_messages
+    # @return [nil]
+    def deliver_messages
+      ensure_threads_running!
+
+      @queue << [:deliver_messages, nil]
+
+      nil
+    end
+
+    # Shuts down the producer, releasing the network resources used. This
+    # method will block until the buffered messages have been delivered.
+    #
+    # @see Kafka::Producer#shutdown
+    # @return [nil]
+    def shutdown
+      ensure_threads_running!
+
+      @timer_thread && @timer_thread.exit
+      @queue << [:shutdown, nil]
+      @worker_thread && @worker_thread.join
+
+      nil
+    end
+
+    private
+
+    def ensure_threads_running!
+      return if worker_thread_alive? && timer_thread_alive?
+
+      @thread_mutex.synchronize do
+        @worker_thread = Thread.new { @worker.run } unless worker_thread_alive?
+        @timer_thread = Thread.new { @timer.run } unless timer_thread_alive?
+      end
+    end
+
+    def worker_thread_alive?
+      !!@worker_thread && @worker_thread.alive?
+    end
+
+    def timer_thread_alive?
+      !!@timer_thread && @timer_thread.alive?
+    end
+
+    def buffer_overflow(topic, message)
+      @instrumenter.instrument("buffer_overflow.async_producer", {
+        topic: topic,
+      })
+
+      raise BufferOverflow, message
+    end
+
+    class Timer
+      def initialize(interval:, queue:)
+        @queue = queue
+        @interval = interval
+      end
+
+      def run
+        # Permanently sleep if the timer interval is zero.
+        Thread.stop if @interval.zero?
+
+        loop do
+          sleep(@interval)
+          @queue << [:deliver_messages, nil]
+        end
+      end
+    end
+
+    class Worker
+      def initialize(queue:, producer:, delivery_threshold:, max_retries: -1, retry_backoff: 0, instrumenter:, logger:)
+        @queue = queue
+        @producer = producer
+        @delivery_threshold = delivery_threshold
+        @max_retries = max_retries
+        @retry_backoff = retry_backoff
+        @instrumenter = instrumenter
+        @logger = TaggedLogger.new(logger)
+      end
+
+      def run
+        @logger.push_tags(@producer.to_s)
+        @logger.info "Starting async producer in the background..."
+
+        do_loop
+      rescue Exception => e
+        @logger.error "Unexpected Kafka error #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+        @logger.error "Async producer crashed!"
+      ensure
+        @producer.shutdown
+        @logger.pop_tags
+      end
+
+      private
+
+      def do_loop
+        loop do
+          begin
+            operation, payload = @queue.pop
+
+            case operation
+            when :produce
+              produce(payload[0], **payload[1])
+              deliver_messages if threshold_reached?
+            when :deliver_messages
+              deliver_messages
+            when :shutdown
+              begin
+                # Deliver any pending messages first.
+                @producer.deliver_messages
+              rescue Error => e
+                @logger.error("Failed to deliver messages during shutdown: #{e.message}")
+
+                @instrumenter.instrument("drop_messages.async_producer", {
+                  message_count: @producer.buffer_size + @queue.size,
+                })
+              end
+
+              # Stop the run loop.
+              break
+            else
+              raise "Unknown operation #{operation.inspect}"
+            end
+          end
+        end
+      rescue Kafka::Error => e
+        @logger.error "Unexpected Kafka error #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+        @logger.info "Restarting in 10 seconds..."
+
+        sleep 10
+        retry
+      end
+
+      def produce(value, **kwargs)
+        retries = 0
+        begin
+          @producer.produce(value, **kwargs)
+        rescue BufferOverflow => e
+          deliver_messages
+          if @max_retries == -1
+            retry
+          elsif retries < @max_retries
+            retries += 1
+            sleep @retry_backoff**retries
+            retry
+          else
+            @logger.error("Failed to asynchronously produce messages due to BufferOverflow")
+            @instrumenter.instrument("error.async_producer", { error: e })
+          end
+        end
+      end
+
+      def deliver_messages
+        @producer.deliver_messages
+      rescue DeliveryFailed, ConnectionError => e
+        # Failed to deliver messages -- nothing to do but log and try again later.
+        @logger.error("Failed to asynchronously deliver messages: #{e.message}")
+        @instrumenter.instrument("error.async_producer", { error: e })
+      end
+
+      def threshold_reached?
+        @delivery_threshold > 0 &&
+          @producer.buffer_size >= @delivery_threshold
+      end
+    end
+  end
+end