ruby-kafka-temp-fork 0.0.1

data/lib/kafka/consumer_group.rb

@@ -0,0 +1,231 @@
+# frozen_string_literal: true
+
+require "set"
+require "kafka/consumer_group/assignor"
+require "kafka/round_robin_assignment_strategy"
+
+module Kafka
+  class ConsumerGroup
+    attr_reader :assigned_partitions, :generation_id, :group_id
+
+    def initialize(cluster:, logger:, group_id:, session_timeout:, rebalance_timeout:, retention_time:, instrumenter:, assignment_strategy:)
+      @cluster = cluster
+      @logger = TaggedLogger.new(logger)
+      @group_id = group_id
+      @session_timeout = session_timeout
+      @rebalance_timeout = rebalance_timeout
+      @instrumenter = instrumenter
+      @member_id = ""
+      @generation_id = nil
+      @members = {}
+      @topics = Set.new
+      @assigned_partitions = {}
+      @assignor = Assignor.new(
+        cluster: cluster,
+        strategy: assignment_strategy || RoundRobinAssignmentStrategy.new
+      )
+      @retention_time = retention_time
+    end
+
+    def subscribe(topic)
+      @topics.add(topic)
+      @cluster.add_target_topics([topic])
+    end
+
+    def subscribed_partitions
+      @assigned_partitions.select { |topic, _| @topics.include?(topic) }
+    end
+
+    def assigned_to?(topic, partition)
+      subscribed_partitions.fetch(topic, []).include?(partition)
+    end
+
+    def member?
+      !@generation_id.nil?
+    end
+
+    def join
+      if @topics.empty?
+        raise Kafka::Error, "Cannot join group without at least one topic subscription"
+      end
+
+      join_group
+      synchronize
+    rescue NotCoordinatorForGroup
+      @logger.error "Failed to find coordinator for group `#{@group_id}`; retrying..."
+      sleep 1
+      @coordinator = nil
+      retry
+    rescue ConnectionError
+      @logger.error "Connection error while trying to join group `#{@group_id}`; retrying..."
+      sleep 1
+      @cluster.mark_as_stale!
+      @coordinator = nil
+      retry
+    end
+
+    def leave
+      @logger.info "Leaving group `#{@group_id}`"
+
+      # Having a generation id indicates that we're a member of the group.
+      @generation_id = nil
+
+      @instrumenter.instrument("leave_group.consumer", group_id: @group_id) do
+        coordinator.leave_group(group_id: @group_id, member_id: @member_id)
+      end
+    rescue ConnectionError
+    end
+
+    def fetch_offsets
+      coordinator.fetch_offsets(
+        group_id: @group_id,
+        topics: @assigned_partitions,
+      )
+    end
+
+    def commit_offsets(offsets)
+      response = coordinator.commit_offsets(
+        group_id: @group_id,
+        member_id: @member_id,
+        generation_id: @generation_id,
+        offsets: offsets,
+        retention_time: @retention_time
+      )
+
+      response.topics.each do |topic, partitions|
+        partitions.each do |partition, error_code|
+          Protocol.handle_error(error_code)
+        end
+      end
+    rescue Kafka::Error => e
+      @logger.error "Error committing offsets: #{e}"
+      raise OffsetCommitError, e
+    end
+
+    def heartbeat
+      @logger.debug "Sending heartbeat..."
+
+      @instrumenter.instrument('heartbeat.consumer',
+                               group_id: @group_id,
+                               topic_partitions: @assigned_partitions) do
+
+        response = coordinator.heartbeat(
+          group_id: @group_id,
+          generation_id: @generation_id,
+          member_id: @member_id,
+        )
+
+        Protocol.handle_error(response.error_code)
+      end
+    rescue ConnectionError, UnknownMemberId, IllegalGeneration => e
+      @logger.error "Error sending heartbeat: #{e}"
+      raise HeartbeatError, e
+    rescue RebalanceInProgress => e
+      @logger.warn "Error sending heartbeat: #{e}"
+      raise HeartbeatError, e
+    rescue NotCoordinatorForGroup
+      @logger.error "Failed to find coordinator for group `#{@group_id}`; retrying..."
+      sleep 1
+      @coordinator = nil
+      retry
+    end
+
+    def to_s
+      "[#{@group_id}] {" + assigned_partitions.map { |topic, partitions|
+        partition_str = partitions.size > 5 ?
+                          "#{partitions[0..4].join(', ')}..." :
+                          partitions.join(', ')
+        "#{topic}: #{partition_str}"
+      }.join('; ') + '}:'
+    end
+
+    private
+
+    def join_group
+      @logger.info "Joining group `#{@group_id}`"
+
+      @instrumenter.instrument("join_group.consumer", group_id: @group_id) do
+        response = coordinator.join_group(
+          group_id: @group_id,
+          session_timeout: @session_timeout,
+          rebalance_timeout: @rebalance_timeout,
+          member_id: @member_id,
+          topics: @topics,
+          protocol_name: @assignor.protocol_name,
+          user_data: @assignor.user_data,
+        )
+
+        Protocol.handle_error(response.error_code)
+
+        @generation_id = response.generation_id
+        @member_id = response.member_id
+        @leader_id = response.leader_id
+        @members = response.members
+      end
+
+      @logger.info "Joined group `#{@group_id}` with member id `#{@member_id}`"
+    rescue UnknownMemberId
+      @logger.error "Failed to join group; resetting member id and retrying in 1s..."
+
+      @member_id = ""
+      sleep 1
+
+      retry
+    rescue CoordinatorLoadInProgress
+      @logger.error "Coordinator broker still loading, retrying in 1s..."
+
+      sleep 1
+
+      retry
+    end
+
+    def group_leader?
+      @member_id == @leader_id
+    end
+
+    def synchronize
+      group_assignment = {}
+
+      if group_leader?
+        @logger.info "Chosen as leader of group `#{@group_id}`"
+
+        topics = Set.new
+        @members.each do |_member, metadata|
+          metadata.topics.each { |t| topics.add(t) }
+        end
+
+        group_assignment = @assignor.assign(
+          members: @members,
+          topics: topics,
+        )
+      end
+
+      @instrumenter.instrument("sync_group.consumer", group_id: @group_id) do
+        response = coordinator.sync_group(
+          group_id: @group_id,
+          generation_id: @generation_id,
+          member_id: @member_id,
+          group_assignment: group_assignment,
+        )
+
+        Protocol.handle_error(response.error_code)
+
+        response.member_assignment.topics.each do |topic, assigned_partitions|
+          @logger.info "Partitions assigned for `#{topic}`: #{assigned_partitions.join(', ')}"
+        end
+
+        @assigned_partitions.replace(response.member_assignment.topics)
+      end
+    end
+
+    def coordinator
+      @coordinator ||= @cluster.get_group_coordinator(group_id: @group_id)
+    rescue CoordinatorNotAvailable
+      @logger.error "Group coordinator not available for group `#{@group_id}`"
+
+      sleep 1
+
+      retry
+    end
+  end
+end

data/lib/kafka/consumer_group/assignor.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/member_assignment"
+
+module Kafka
+  class ConsumerGroup
+
+    # A consumer group partition assignor
+    class Assignor
+      Partition = Struct.new(:topic, :partition_id)
+
+      # @param cluster [Kafka::Cluster]
+      # @param strategy [Object] an object that implements #protocol_type,
+      #   #user_data, and #assign.
+      def initialize(cluster:, strategy:)
+        @cluster = cluster
+        @strategy = strategy
+      end
+
+      def protocol_name
+        @strategy.respond_to?(:protocol_name) ? @strategy.protocol_name : @strategy.class.to_s
+      end
+
+      def user_data
+        @strategy.user_data if @strategy.respond_to?(:user_data)
+      end
+
+      # Assign the topic partitions to the group members.
+      #
+      # @param members [Hash<String, Kafka::Protocol::JoinGroupResponse::Metadata>] a hash
+      #   mapping member ids to metadata.
+      # @param topics [Array<String>] topics
+      # @return [Hash<String, Kafka::Protocol::MemberAssignment>] a hash mapping member
+      #   ids to assignments.
+      def assign(members:, topics:)
+        topic_partitions = topics.flat_map do |topic|
+          begin
+            partition_ids = @cluster.partitions_for(topic).map(&:partition_id)
+          rescue UnknownTopicOrPartition
+            raise UnknownTopicOrPartition, "unknown topic #{topic}"
+          end
+          partition_ids.map {|partition_id| Partition.new(topic, partition_id) }
+        end
+
+        group_assignment = {}
+
+        members.each_key do |member_id|
+          group_assignment[member_id] = Protocol::MemberAssignment.new
+        end
+        @strategy.call(cluster: @cluster, members: members, partitions: topic_partitions).each do |member_id, partitions|
+          Array(partitions).each do |partition|
+            group_assignment[member_id].assign(partition.topic, [partition.partition_id])
+          end
+        end
+
+        group_assignment
+      rescue Kafka::LeaderNotAvailable
+        sleep 1
+        retry
+      end
+    end
+  end
+end

data/lib/kafka/crc32_hash.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require "zlib"
+
+module Kafka
+  class Crc32Hash
+
+    # crc32 is supported natively
+    def load; end
+
+    def hash(value)
+      Zlib.crc32(value)
+    end
+  end
+end

data/lib/kafka/datadog.rb

@@ -0,0 +1,420 @@
+# frozen_string_literal: true
+
+begin
+  require "datadog/statsd"
+rescue LoadError
+  $stderr.puts "In order to report Kafka client metrics to Datadog you need to install the `dogstatsd-ruby` gem."
+  raise
+end
+
+require "active_support/subscriber"
+
+module Kafka
+
+  # Reports operational metrics to a Datadog agent using the modified Statsd protocol.
+  #
+  #     require "kafka/datadog"
+  #
+  #     # Default is "ruby_kafka".
+  #     Kafka::Datadog.namespace = "custom-namespace"
+  #
+  #     # Default is "127.0.0.1".
+  #     Kafka::Datadog.host = "statsd.something.com"
+  #
+  #     # Default is 8125.
+  #     Kafka::Datadog.port = 1234
+  #
+  # Once the file has been required, no further configuration is needed – all operational
+  # metrics are automatically emitted.
+  module Datadog
+    STATSD_NAMESPACE = "ruby_kafka"
+
+    class << self
+      def statsd
+        @statsd ||= ::Datadog::Statsd.new(host, port, namespace: namespace, tags: tags, socket_path: socket_path)
+      end
+
+      def statsd=(statsd)
+        clear
+        @statsd = statsd
+      end
+
+      def host
+        @host
+      end
+
+      def host=(host)
+        @host = host
+        clear
+      end
+
+      def port
+        @port
+      end
+
+      def port=(port)
+        @port = port
+        clear
+      end
+
+      def socket_path
+        @socket_path
+      end
+
+      def socket_path=(socket_path)
+        @socket_path = socket_path
+        clear
+      end
+
+      def namespace
+        @namespace ||= STATSD_NAMESPACE
+      end
+
+      def namespace=(namespace)
+        @namespace = namespace
+        clear
+      end
+
+      def tags
+        @tags ||= []
+      end
+
+      def tags=(tags)
+        @tags = tags
+        clear
+      end
+
+      private
+
+      def clear
+        @statsd && @statsd.close
+        @statsd = nil
+      end
+    end
+
+    class StatsdSubscriber < ActiveSupport::Subscriber
+      private
+
+      %w[increment histogram count timing gauge].each do |type|
+        define_method(type) do |*args, **kwargs|
+          emit(type, *args, **kwargs)
+        end
+      end
+
+      def emit(type, *args, tags: {})
+        tags = tags.map {|k, v| "#{k}:#{v}" }.to_a
+
+        Kafka::Datadog.statsd.send(type, *args, tags: tags)
+      end
+    end
+
+    class ConnectionSubscriber < StatsdSubscriber
+      def request(event)
+        client = event.payload.fetch(:client_id)
+        api = event.payload.fetch(:api, "unknown")
+        request_size = event.payload.fetch(:request_size, 0)
+        response_size = event.payload.fetch(:response_size, 0)
+        broker = event.payload.fetch(:broker_host)
+
+        tags = {
+          client: client,
+          api: api,
+          broker: broker
+        }
+
+        timing("api.latency", event.duration, tags: tags)
+        increment("api.calls", tags: tags)
+
+        histogram("api.request_size", request_size, tags: tags)
+        histogram("api.response_size", response_size, tags: tags)
+
+        if event.payload.key?(:exception)
+          increment("api.errors", tags: tags)
+        end
+      end
+
+      attach_to "connection.kafka"
+    end
+
+    class ConsumerSubscriber < StatsdSubscriber
+      def process_message(event)
+        offset = event.payload.fetch(:offset)
+        offset_lag = event.payload.fetch(:offset_lag)
+        create_time = event.payload.fetch(:create_time)
+        time_lag = create_time && ((Time.now - create_time) * 1000).to_i
+
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+          topic: event.payload.fetch(:topic),
+          partition: event.payload.fetch(:partition),
+        }
+
+        if event.payload.key?(:exception)
+          increment("consumer.process_message.errors", tags: tags)
+        else
+          timing("consumer.process_message.latency", event.duration, tags: tags)
+          increment("consumer.messages", tags: tags)
+        end
+
+        gauge("consumer.offset", offset, tags: tags)
+        gauge("consumer.lag", offset_lag, tags: tags)
+
+        # Not all messages have timestamps.
+        if time_lag
+          gauge("consumer.time_lag", time_lag, tags: tags)
+        end
+      end
+
+      def process_batch(event)
+        offset = event.payload.fetch(:last_offset)
+        messages = event.payload.fetch(:message_count)
+        create_time = event.payload.fetch(:last_create_time)
+        time_lag = create_time && ((Time.now - create_time) * 1000).to_i
+
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+          topic: event.payload.fetch(:topic),
+          partition: event.payload.fetch(:partition),
+        }
+
+        if event.payload.key?(:exception)
+          increment("consumer.process_batch.errors", tags: tags)
+        else
+          timing("consumer.process_batch.latency", event.duration, tags: tags)
+          count("consumer.messages", messages, tags: tags)
+        end
+
+        gauge("consumer.offset", offset, tags: tags)
+
+        if time_lag
+          gauge("consumer.time_lag", time_lag, tags: tags)
+        end
+      end
+
+      def fetch_batch(event)
+        lag = event.payload.fetch(:offset_lag)
+        batch_size = event.payload.fetch(:message_count)
+
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+          topic: event.payload.fetch(:topic),
+          partition: event.payload.fetch(:partition),
+        }
+
+        histogram("consumer.batch_size", batch_size, tags: tags)
+        gauge("consumer.lag", lag, tags: tags)
+      end
+
+      def join_group(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+        }
+
+        timing("consumer.join_group", event.duration, tags: tags)
+
+        if event.payload.key?(:exception)
+          increment("consumer.join_group.errors", tags: tags)
+        end
+      end
+
+      def sync_group(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+        }
+
+        timing("consumer.sync_group", event.duration, tags: tags)
+
+        if event.payload.key?(:exception)
+          increment("consumer.sync_group.errors", tags: tags)
+        end
+      end
+
+      def leave_group(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+        }
+
+        timing("consumer.leave_group", event.duration, tags: tags)
+
+        if event.payload.key?(:exception)
+          increment("consumer.leave_group.errors", tags: tags)
+        end
+      end
+
+      def loop(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+        }
+
+        histogram("consumer.loop.duration", event.duration, tags: tags)
+      end
+
+      def pause_status(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+          topic: event.payload.fetch(:topic),
+          partition: event.payload.fetch(:partition),
+        }
+
+        duration = event.payload.fetch(:duration)
+
+        gauge("consumer.pause.duration", duration, tags: tags)
+      end
+
+      attach_to "consumer.kafka"
+    end
+
+    class ProducerSubscriber < StatsdSubscriber
+      def produce_message(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+        message_size = event.payload.fetch(:message_size)
+        buffer_size = event.payload.fetch(:buffer_size)
+        max_buffer_size = event.payload.fetch(:max_buffer_size)
+        buffer_fill_ratio = buffer_size.to_f / max_buffer_size.to_f
+        buffer_fill_percentage = buffer_fill_ratio * 100.0
+
+        tags = {
+          client: client,
+          topic: topic,
+        }
+
+        # This gets us the write rate.
+        increment("producer.produce.messages", tags: tags.merge(topic: topic))
+
+        # Information about typical/average/95p message size.
+        histogram("producer.produce.message_size", message_size, tags: tags.merge(topic: topic))
+
+        # Aggregate message size.
+        count("producer.produce.message_size.sum", message_size, tags: tags.merge(topic: topic))
+
+        # This gets us the avg/max buffer size per producer.
+        histogram("producer.buffer.size", buffer_size, tags: tags)
+
+        # This gets us the avg/max buffer fill ratio per producer.
+        histogram("producer.buffer.fill_ratio", buffer_fill_ratio, tags: tags)
+        histogram("producer.buffer.fill_percentage", buffer_fill_percentage, tags: tags)
+      end
+
+      def buffer_overflow(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          topic: event.payload.fetch(:topic),
+        }
+
+        increment("producer.produce.errors", tags: tags)
+      end
+
+      def deliver_messages(event)
+        client = event.payload.fetch(:client_id)
+        message_count = event.payload.fetch(:delivered_message_count)
+        attempts = event.payload.fetch(:attempts)
+
+        tags = {
+          client: client,
+        }
+
+        if event.payload.key?(:exception)
+          increment("producer.deliver.errors", tags: tags)
+        end
+
+        timing("producer.deliver.latency", event.duration, tags: tags)
+
+        # Messages delivered to Kafka:
+        count("producer.deliver.messages", message_count, tags: tags)
+
+        # Number of attempts to deliver messages:
+        histogram("producer.deliver.attempts", attempts, tags: tags)
+      end
+
+      def ack_message(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          topic: event.payload.fetch(:topic),
+        }
+
+        # Number of messages ACK'd for the topic.
+        increment("producer.ack.messages", tags: tags)
+
+        # Histogram of delay between a message being produced and it being ACK'd.
+        histogram("producer.ack.delay", event.payload.fetch(:delay), tags: tags)
+      end
+
+      def topic_error(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          topic: event.payload.fetch(:topic)
+        }
+
+        increment("producer.ack.errors", tags: tags)
+      end
+
+      attach_to "producer.kafka"
+    end
+
+    class AsyncProducerSubscriber < StatsdSubscriber
+      def enqueue_message(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+        queue_size = event.payload.fetch(:queue_size)
+        max_queue_size = event.payload.fetch(:max_queue_size)
+        queue_fill_ratio = queue_size.to_f / max_queue_size.to_f
+
+        tags = {
+          client: client,
+          topic: topic,
+        }
+
+        # This gets us the avg/max queue size per producer.
+        histogram("async_producer.queue.size", queue_size, tags: tags)
+
+        # This gets us the avg/max queue fill ratio per producer.
+        histogram("async_producer.queue.fill_ratio", queue_fill_ratio, tags: tags)
+      end
+
+      def buffer_overflow(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+          topic: event.payload.fetch(:topic),
+        }
+
+        increment("async_producer.produce.errors", tags: tags)
+      end
+
+      def drop_messages(event)
+        tags = {
+          client: event.payload.fetch(:client_id),
+        }
+
+        message_count = event.payload.fetch(:message_count)
+
+        count("async_producer.dropped_messages", message_count, tags: tags)
+      end
+
+      attach_to "async_producer.kafka"
+    end
+
+    class FetcherSubscriber < StatsdSubscriber
+      def loop(event)
+        queue_size = event.payload.fetch(:queue_size)
+
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+        }
+
+        gauge("fetcher.queue_size", queue_size, tags: tags)
+      end
+
+      attach_to "fetcher.kafka"
+    end
+  end
+end