ruby-kafka 0.1.7 → 0.2.0

data/lib/kafka/consumer_group.rb

@@ -0,0 +1,172 @@
+require "set"
+require "kafka/round_robin_assignment_strategy"
+
+module Kafka
+  class ConsumerGroup
+    attr_reader :assigned_partitions
+
+    def initialize(cluster:, logger:, group_id:, session_timeout:)
+      @cluster = cluster
+      @logger = logger
+      @group_id = group_id
+      @session_timeout = session_timeout
+      @member_id = ""
+      @generation_id = nil
+      @members = {}
+      @topics = Set.new
+      @assigned_partitions = {}
+      @assignment_strategy = RoundRobinAssignmentStrategy.new(cluster: @cluster)
+    end
+
+    def subscribe(topic)
+      @topics.add(topic)
+      @cluster.add_target_topics([topic])
+    end
+
+    def member?
+      !@generation_id.nil?
+    end
+
+    def join
+      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..."
+      @coordinator = nil
+      retry
+    end
+
+    def leave
+      @logger.info "[#{@member_id}] Leaving group `#{@group_id}`"
+      coordinator.leave_group(group_id: @group_id, member_id: @member_id)
+    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,
+      )
+
+      response.topics.each do |topic, partitions|
+        partitions.each do |partition, error_code|
+          Protocol.handle_error(error_code)
+        end
+      end
+    rescue UnknownMemberId
+      @logger.error "Kicked out of group; rejoining"
+      join
+      retry
+    rescue IllegalGeneration
+      @logger.error "Illegal generation #{@generation_id}; rejoining group"
+      join
+      retry
+    end
+
+    def heartbeat
+      @logger.info "[#{@member_id}] Sending heartbeat..."
+
+      response = coordinator.heartbeat(
+        group_id: @group_id,
+        generation_id: @generation_id,
+        member_id: @member_id,
+      )
+
+      Protocol.handle_error(response.error_code)
+    rescue ConnectionError => e
+      @logger.error "Connection error while sending heartbeat; rejoining"
+      join
+    rescue UnknownMemberId
+      @logger.error "Kicked out of group; rejoining"
+      join
+    rescue RebalanceInProgress
+      @logger.error "Group is rebalancing; rejoining"
+      join
+    end
+
+    private
+
+    def join_group
+      @logger.info "Joining group `#{@group_id}`"
+
+      response = coordinator.join_group(
+        group_id: @group_id,
+        session_timeout: @session_timeout,
+        member_id: @member_id,
+      )
+
+      Protocol.handle_error(response.error_code)
+
+      @generation_id = response.generation_id
+      @member_id = response.member_id
+      @leader_id = response.leader_id
+      @members = response.members
+
+      @logger.info "[#{@member_id}] 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 = nil
+      sleep 1
+
+      retry
+    end
+
+    def group_leader?
+      @member_id == @leader_id
+    end
+
+    def synchronize
+      @logger.info "[#{@member_id}] Synchronizing group"
+
+      group_assignment = {}
+
+      if group_leader?
+        @logger.info "[#{@member_id}] Chosen as leader of group `#{@group_id}`"
+
+        group_assignment = @assignment_strategy.assign(
+          members: @members.keys,
+          topics: @topics,
+        )
+      end
+
+      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 "[#{@member_id}] Partitions assigned for `#{topic}`: #{assigned_partitions.join(', ')}"
+      end
+
+      @assigned_partitions.replace(response.member_assignment.topics)
+    end
+
+    def coordinator
+      @coordinator ||= @cluster.get_group_coordinator(group_id: @group_id)
+    rescue GroupCoordinatorNotAvailable
+      @logger.error "Group coordinator not available for group `#{@group_id}`"
+
+      sleep 1
+
+      retry
+    end
+  end
+end

data/lib/kafka/fetch_operation.rb

@@ -16,7 +16,7 @@ module Kafka
   #     operation.execute
   #
   class FetchOperation
-    def initialize(cluster:, logger:, min_bytes:, max_wait_time:)
+    def initialize(cluster:, logger:, min_bytes: 1, max_wait_time: 5)
       @cluster = cluster
       @logger = logger
       @min_bytes = min_bytes
@@ -24,7 +24,7 @@ module Kafka
       @topics = {}
     end
 
-    def fetch_from_partition(topic, partition, offset:, max_bytes:)
+    def fetch_from_partition(topic, partition, offset: :latest, max_bytes: 1048576)
       if offset == :earliest
         offset = -2
       elsif offset == :latest

data/lib/kafka/produce_operation.rb

@@ -25,13 +25,12 @@ module Kafka
   # * `sent_message_count` – the number of messages that were successfully sent.
   #
   class ProduceOperation
-    def initialize(cluster:, buffer:, compression_codec:, compression_threshold:, required_acks:, ack_timeout:, logger:)
+    def initialize(cluster:, buffer:, compressor:, required_acks:, ack_timeout:, logger:)
       @cluster = cluster
       @buffer = buffer
       @required_acks = required_acks
       @ack_timeout = ack_timeout
-      @compression_codec = compression_codec
-      @compression_threshold = compression_threshold
+      @compressor = compressor
       @logger = logger
     end
 
@@ -78,11 +77,8 @@ module Kafka
           messages_for_topics = {}
 
           message_buffer.each do |topic, partition, messages|
-            message_set = Protocol::MessageSet.new(
-              messages: messages,
-              compression_codec: @compression_codec,
-              compression_threshold: @compression_threshold,
-            )
+            message_set = Protocol::MessageSet.new(messages: messages)
+            message_set = @compressor.compress(message_set)
 
             messages_for_topics[topic] ||= {}
             messages_for_topics[topic][partition] = message_set

data/lib/kafka/producer.rb

@@ -4,7 +4,7 @@ require "kafka/message_buffer"
 require "kafka/produce_operation"
 require "kafka/pending_message_queue"
 require "kafka/pending_message"
-require "kafka/compression"
+require "kafka/compressor"
 
 module Kafka
 
@@ -173,8 +173,11 @@ module Kafka
       @retry_backoff = retry_backoff
       @max_buffer_size = max_buffer_size
       @max_buffer_bytesize = max_buffer_bytesize
-      @compression_codec = Compression.find_codec(compression_codec)
-      @compression_threshold = compression_threshold
+
+      @compressor = Compressor.new(
+        codec_name: @compression_codec,
+        threshold: @compression_threshold,
+      )
 
       # The set of topics that are produced to.
       @target_topics = Set.new
@@ -303,8 +306,7 @@ module Kafka
         buffer: @buffer,
         required_acks: @required_acks,
         ack_timeout: @ack_timeout,
-        compression_codec: @compression_codec,
-        compression_threshold: @compression_threshold,
+        compressor: @compressor,
         logger: @logger,
       )
 

data/lib/kafka/protocol.rb

@@ -8,6 +8,13 @@ module Kafka
       1 => :fetch,
       2 => :list_offset,
       3 => :topic_metadata,
+      8 => :offset_commit,
+      9 => :offset_fetch,
+      10 => :group_coordinator,
+      11 => :join_group,
+      12 => :heartbeat,
+      13 => :leave_group,
+      14 => :sync_group,
     }
 
     ERRORS = {
@@ -23,11 +30,17 @@ module Kafka
       9 => ReplicaNotAvailable,
       10 => MessageSizeTooLarge,
       12 => OffsetMetadataTooLarge,
+      15 => GroupCoordinatorNotAvailable,
+      16 => NotCoordinatorForGroup,
       17 => InvalidTopic,
       18 => RecordListTooLarge,
       19 => NotEnoughReplicas,
       20 => NotEnoughReplicasAfterAppend,
       21 => InvalidRequiredAcks,
+      22 => IllegalGeneration,
+      25 => UnknownMemberId,
+      26 => InvalidSessionTimeout,
+      27 => RebalanceInProgress,
     }
 
     def self.handle_error(error_code)
@@ -54,3 +67,17 @@ require "kafka/protocol/fetch_request"
 require "kafka/protocol/fetch_response"
 require "kafka/protocol/list_offset_request"
 require "kafka/protocol/list_offset_response"
+require "kafka/protocol/group_coordinator_request"
+require "kafka/protocol/group_coordinator_response"
+require "kafka/protocol/join_group_request"
+require "kafka/protocol/join_group_response"
+require "kafka/protocol/sync_group_request"
+require "kafka/protocol/sync_group_response"
+require "kafka/protocol/leave_group_request"
+require "kafka/protocol/leave_group_response"
+require "kafka/protocol/heartbeat_request"
+require "kafka/protocol/heartbeat_response"
+require "kafka/protocol/offset_fetch_request"
+require "kafka/protocol/offset_fetch_response"
+require "kafka/protocol/offset_commit_request"
+require "kafka/protocol/offset_commit_response"

data/lib/kafka/protocol/consumer_group_protocol.rb

@@ -0,0 +1,17 @@
+module Kafka
+  module Protocol
+    class ConsumerGroupProtocol
+      def initialize(version: 0, topics:, user_data: nil)
+        @version = version
+        @topics = topics
+        @user_data = user_data
+      end
+
+      def encode(encoder)
+        encoder.write_int16(@version)
+        encoder.write_array(@topics) {|topic| encoder.write_string(topic) }
+        encoder.write_bytes(@user_data)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/group_coordinator_request.rb

@@ -0,0 +1,21 @@
+module Kafka
+  module Protocol
+    class GroupCoordinatorRequest
+      def initialize(group_id:)
+        @group_id = group_id
+      end
+
+      def api_key
+        10
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+      end
+
+      def response_class
+        GroupCoordinatorResponse
+      end
+    end
+  end
+end

data/lib/kafka/protocol/group_coordinator_response.rb

@@ -0,0 +1,25 @@
+module Kafka
+  module Protocol
+    class GroupCoordinatorResponse
+      attr_reader :error_code
+
+      attr_reader :coordinator_id, :coordinator_host, :coordinator_port
+
+      def initialize(error_code:, coordinator_id:, coordinator_host:, coordinator_port:)
+        @error_code = error_code
+        @coordinator_id = coordinator_id
+        @coordinator_host = coordinator_host
+        @coordinator_port = coordinator_port
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          coordinator_id: decoder.int32,
+          coordinator_host: decoder.string,
+          coordinator_port: decoder.int32,
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/heartbeat_request.rb

@@ -0,0 +1,25 @@
+module Kafka
+  module Protocol
+    class HeartbeatRequest
+      def initialize(group_id:, generation_id:, member_id:)
+        @group_id = group_id
+        @generation_id = generation_id
+        @member_id = member_id
+      end
+
+      def api_key
+        12
+      end
+
+      def response_class
+        HeartbeatResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@generation_id)
+        encoder.write_string(@member_id)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/heartbeat_response.rb

@@ -0,0 +1,15 @@
+module Kafka
+  module Protocol
+    class HeartbeatResponse
+      attr_reader :error_code
+
+      def initialize(error_code:)
+        @error_code = error_code
+      end
+
+      def self.decode(decoder)
+        new(error_code: decoder.int16)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/join_group_request.rb

@@ -0,0 +1,39 @@
+require "kafka/protocol/consumer_group_protocol"
+
+module Kafka
+  module Protocol
+    class JoinGroupRequest
+      PROTOCOL_TYPE = "consumer"
+
+      def initialize(group_id:, session_timeout:, member_id:, topics: [])
+        @group_id = group_id
+        @session_timeout = session_timeout * 1000 # Kafka wants ms.
+        @member_id = member_id || ""
+        @protocol_type = PROTOCOL_TYPE
+        @group_protocols = {
+          "standard" => ConsumerGroupProtocol.new(topics: ["test-messages"]),
+        }
+      end
+
+      def api_key
+        11
+      end
+
+      def response_class
+        JoinGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@session_timeout)
+        encoder.write_string(@member_id)
+        encoder.write_string(@protocol_type)
+
+        encoder.write_array(@group_protocols) do |name, metadata|
+          encoder.write_string(name)
+          encoder.write_bytes(Encoder.encode_with(metadata))
+        end
+      end
+    end
+  end
+end