karafka 2.0.0.beta1 → 2.0.0.beta4

data/lib/karafka/connection/client.rb

@@ -30,7 +30,7 @@ module Karafka
         @mutex = Mutex.new
         @closed = false
         @subscription_group = subscription_group
-        @buffer = MessagesBuffer.new
+        @buffer = RawMessagesBuffer.new
         @rebalance_manager = RebalanceManager.new
         @kafka = build_consumer
         # Marks if we need to offset. If we did not store offsets, we should not commit the offset
@@ -86,8 +86,7 @@ module Karafka
       # @param message [Karafka::Messages::Message]
       def store_offset(message)
         @mutex.synchronize do
-          @offsetting = true
-          @kafka.store_offset(message)
+          internal_store_offset(message)
         end
       end
 
@@ -104,14 +103,7 @@ module Karafka
       def commit_offsets(async: true)
         @mutex.lock
 
-        return unless @offsetting
-
-        @kafka.commit(nil, async)
-        @offsetting = false
-      rescue Rdkafka::RdkafkaError => e
-        return if e.code == :no_offset
-
-        raise e
+        internal_commit_offsets(async: async)
       ensure
         @mutex.unlock
       end
@@ -128,7 +120,11 @@ module Karafka
       #
       # @param message [Messages::Message, Messages::Seek] message to which we want to seek to
       def seek(message)
+        @mutex.lock
+
         @kafka.seek(message)
+      ensure
+        @mutex.unlock
       end
 
       # Pauses given partition and moves back to last successful offset processed.
@@ -144,15 +140,17 @@ module Karafka
         # Do not pause if the client got closed, would not change anything
         return if @closed
 
+        pause_msg = Messages::Seek.new(topic, partition, offset)
+
+        internal_commit_offsets(async: false)
+
         tpl = topic_partition_list(topic, partition)
 
         return unless tpl
 
         @kafka.pause(tpl)
 
-        pause_msg = Messages::Seek.new(topic, partition, offset)
-
-        seek(pause_msg)
+        @kafka.seek(pause_msg)
       ensure
         @mutex.unlock
       end
@@ -166,6 +164,11 @@ module Karafka
 
         return if @closed
 
+        # Always commit synchronously offsets if any when we resume
+        # This prevents resuming without offset in case it would not be committed prior
+        # We can skip performance penalty since resuming should not happen too often
+        internal_commit_offsets(async: false)
+
         tpl = topic_partition_list(topic, partition)
 
         return unless tpl
@@ -187,6 +190,7 @@ module Karafka
       # Marks given message as consumed.
       #
       # @param [Karafka::Messages::Message] message that we want to mark as processed
+      # @return [Boolean] true if successful. False if we no longer own given partition
       # @note This method won't trigger automatic offsets commits, rather relying on the offset
       #   check-pointing trigger that happens with each batch processed
       def mark_as_consumed(message)
@@ -196,8 +200,10 @@ module Karafka
       # Marks a given message as consumed and commits the offsets in a blocking way.
       #
       # @param [Karafka::Messages::Message] message that we want to mark as processed
+      # @return [Boolean] true if successful. False if we no longer own given partition
       def mark_as_consumed!(message)
-        mark_as_consumed(message)
+        return false unless mark_as_consumed(message)
+
         commit_offsets!
       end
 
@@ -214,11 +220,44 @@ module Karafka
 
       private
 
+      # When we cannot store an offset, it means we no longer own the partition
+      #
+      # Non thread-safe offset storing method
+      # @param message [Karafka::Messages::Message]
+      # @return [Boolean] true if we could store the offset (if we still own the partition)
+      def internal_store_offset(message)
+        @offsetting = true
+        @kafka.store_offset(message)
+        true
+      rescue Rdkafka::RdkafkaError => e
+        return false if e.code == :assignment_lost
+        return false if e.code == :state
+
+        raise e
+      end
+
+      # Non thread-safe message committing method
+      # @param async [Boolean] should the commit happen async or sync (async by default)
+      # @return [Boolean] true if offset commit worked, false if we've lost the assignment
+      def internal_commit_offsets(async: true)
+        return true unless @offsetting
+
+        @kafka.commit(nil, async)
+        @offsetting = false
+
+        true
+      rescue Rdkafka::RdkafkaError => e
+        return false if e.code == :assignment_lost
+        return false if e.code == :no_offset
+
+        raise e
+      end
+
       # Commits the stored offsets in a sync way and closes the consumer.
       def close
-        commit_offsets!
-
         @mutex.synchronize do
+          internal_commit_offsets(async: false)
+
           @closed = true
 
           # Remove callbacks runners that were registered
@@ -227,7 +266,8 @@ module Karafka
 
           @kafka.close
           @buffer.clear
-          @rebalance_manager.clear
+          # @note We do not clear rebalance manager here as we may still have revocation info here
+          # that we want to consider valid prior to running another reconnection
         end
       end
 
@@ -280,7 +320,13 @@ module Karafka
 
         time_poll.backoff
 
-        retry
+        # We return nil, so we do not restart until running the whole loop
+        # This allows us to run revocation jobs and other things and we will pick up new work
+        # next time after dispatching all the things that are needed
+        #
+        # If we would retry here, the client reset would become transparent and we would not have
+        # a chance to take any actions
+        nil
       end
 
       # Builds a new rdkafka consumer instance based on the subscription group configuration

data/lib/karafka/connection/listener.rb

@@ -3,20 +3,34 @@
 module Karafka
   module Connection
     # A single listener that listens to incoming messages from a single subscription group.
-    # It polls the messages and then enqueues. It also takes care of potential recovery from
+    # It polls the messages and then enqueues jobs. It also takes care of potential recovery from
     # critical errors by restarting everything in a safe manner.
+    #
+    # This is the heart of the consumption process.
     class Listener
+      include Helpers::Async
+
+      # Can be useful for logging
+      # @return [String] id of this listener
+      attr_reader :id
+
       # @param subscription_group [Karafka::Routing::SubscriptionGroup]
       # @param jobs_queue [Karafka::Processing::JobsQueue] queue where we should push work
       # @return [Karafka::Connection::Listener] listener instance
       def initialize(subscription_group, jobs_queue)
+        @id = SecureRandom.uuid
         @subscription_group = subscription_group
         @jobs_queue = jobs_queue
+        @jobs_builder = ::Karafka::App.config.internal.jobs_builder
         @pauses_manager = PausesManager.new
         @client = Client.new(@subscription_group)
         @executors = Processing::ExecutorsBuffer.new(@client, subscription_group)
         # We reference scheduler here as it is much faster than fetching this each time
         @scheduler = ::Karafka::App.config.internal.scheduler
+        # We keep one buffer for messages to preserve memory and not allocate extra objects
+        # We can do this that way because we always first schedule jobs using messages before we
+        # fetch another batch.
+        @messages_buffer = MessagesBuffer.new(subscription_group)
       end
 
       # Runs the main listener fetch loop.
@@ -53,33 +67,55 @@ module Karafka
           )
 
           resume_paused_partitions
-          # We need to fetch data before we revoke lost partitions details as during the polling
-          # the callbacks for tracking lost partitions are triggered. Otherwise we would be always
-          # one batch behind.
-          messages_buffer = @client.batch_poll
 
           Karafka.monitor.instrument(
             'connection.listener.fetch_loop.received',
             caller: self,
-            messages_buffer: messages_buffer
-          )
+            messages_buffer: @messages_buffer
+          ) do
+            # We need to fetch data before we revoke lost partitions details as during the polling
+            # the callbacks for tracking lost partitions are triggered. Otherwise we would be
+            # always one batch behind.
+            poll_and_remap_messages
+          end
 
           # If there were revoked partitions, we need to wait on their jobs to finish before
           # distributing consuming jobs as upon revoking, we might get assigned to the same
           # partitions, thus getting their jobs. The revoking jobs need to finish before
           # appropriate consumers are taken down and re-created
-          wait(@subscription_group) if schedule_revoke_lost_partitions_jobs
-
-          schedule_partitions_jobs(messages_buffer)
+          build_and_schedule_revoke_lost_partitions_jobs
 
           # We wait only on jobs from our subscription group. Other groups are independent.
-          wait(@subscription_group)
+          wait
+
+          build_and_schedule_consumption_jobs
+
+          wait
 
           # We don't use the `#commit_offsets!` here for performance reasons. This can be achieved
           # if needed by using manual offset management.
           @client.commit_offsets
         end
 
+        # If we are stopping we will no longer schedule any jobs despite polling.
+        # We need to keep polling not to exceed the `max.poll.interval` for long-running
+        # non-blocking jobs and we need to allow them to finish. We however do not want to
+        # enqueue any new jobs. It's worth keeping in mind that it is the end user responsibility
+        # to detect shutdown in their long-running logic or else Karafka will force shutdown
+        # after a while.
+        #
+        # We do not care about resuming any partitions or lost jobs as we do not plan to do
+        # anything with them as we're in the shutdown phase.
+        wait_with_poll
+
+        # We do not want to schedule the shutdown jobs prior to finishing all the jobs
+        # (including non-blocking) as there might be a long-running job with a shutdown and then
+        # we would run two jobs in parallel for the same executor and consumer. We do not want that
+        # as it could create a race-condition.
+        build_and_schedule_shutdown_jobs
+
+        wait_with_poll
+
         shutdown
 
         # This is on purpose - see the notes for this method
@@ -100,57 +136,86 @@ module Karafka
 
       # Resumes processing of partitions that were paused due to an error.
       def resume_paused_partitions
-        @pauses_manager.resume { |topic, partition| @client.resume(topic, partition) }
+        @pauses_manager.resume do |topic, partition|
+          @client.resume(topic, partition)
+        end
       end
 
       # Enqueues revoking jobs for partitions that were taken away from the running process.
-      # @return [Boolean] was there anything to revoke
-      # @note We do not use scheduler here as those jobs are not meant to be order optimized in
-      #   any way. Since they operate occasionally it is irrelevant.
-      def schedule_revoke_lost_partitions_jobs
+      def build_and_schedule_revoke_lost_partitions_jobs
         revoked_partitions = @client.rebalance_manager.revoked_partitions
 
-        return false if revoked_partitions.empty?
+        # Stop early to save on some execution and array allocation
+        return if revoked_partitions.empty?
+
+        jobs = []
 
         revoked_partitions.each do |topic, partitions|
           partitions.each do |partition|
             pause_tracker = @pauses_manager.fetch(topic, partition)
             executor = @executors.fetch(topic, partition, pause_tracker)
-            @jobs_queue << Processing::Jobs::Revoked.new(executor)
+            jobs << @jobs_builder.revoked(executor)
           end
         end
 
-        true
+        @scheduler.schedule_revocation(@jobs_queue, jobs)
+      end
+
+      # Enqueues the shutdown jobs for all the executors that exist in our subscription group
+      def build_and_schedule_shutdown_jobs
+        jobs = []
+
+        @executors.each do |_, _, executor|
+          jobs << @jobs_builder.shutdown(executor)
+        end
+
+        @scheduler.schedule_shutdown(@jobs_queue, jobs)
       end
 
-      # Takes the messages per topic partition and enqueues processing jobs in threads.
+      # Polls messages within the time and amount boundaries defined in the settings and then
+      # builds karafka messages based on the raw rdkafka messages buffer returned by the
+      # `#batch_poll` method.
       #
-      # @param messages_buffer [Karafka::Connection::MessagesBuffer] buffer with messages
-      def schedule_partitions_jobs(messages_buffer)
-        @scheduler.call(messages_buffer) do |topic, partition, messages|
-          pause = @pauses_manager.fetch(topic, partition)
+      # @note There are two buffers, one for raw messages and one for "built" karafka messages
+      def poll_and_remap_messages
+        @messages_buffer.remap(
+          @client.batch_poll
+        )
+      end
+
+      # Takes the messages per topic partition and enqueues processing jobs in threads using
+      # given scheduler.
+      def build_and_schedule_consumption_jobs
+        return if @messages_buffer.empty?
 
-          next if pause.paused?
+        jobs = []
 
-          executor = @executors.fetch(topic, partition, pause)
+        @messages_buffer.each do |topic, partition, messages|
+          pause_tracker = @pauses_manager.fetch(topic, partition)
 
-          @jobs_queue << Processing::Jobs::Consume.new(executor, messages)
+          executor = @executors.fetch(topic, partition, pause_tracker)
+
+          jobs << @jobs_builder.consume(executor, messages)
         end
+
+        @scheduler.schedule_consumption(@jobs_queue, jobs)
       end
 
       # Waits for all the jobs from a given subscription group to finish before moving forward
-      # @param subscription_group [Karafka::Routing::SubscriptionGroup]
-      def wait(subscription_group)
-        @jobs_queue.wait(subscription_group.id)
+      def wait
+        @jobs_queue.wait(@subscription_group.id)
+      end
+
+      # Waits without blocking the polling
+      # This should be used only when we no longer plan to use any incoming data and we can safely
+      # discard it
+      def wait_with_poll
+        @client.batch_poll until @jobs_queue.empty?(@subscription_group.id)
       end
 
       # Stops the jobs queue, triggers shutdown on all the executors (sync), commits offsets and
       # stops kafka client.
       def shutdown
-        @jobs_queue.close
-        # This runs synchronously, making sure we finish all the shutdowns before we stop the
-        # client.
-        @executors.shutdown
         @client.commit_offsets!
         @client.stop
       end

data/lib/karafka/connection/listeners_batch.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Karafka
+  module Connection
+    # Abstraction layer around listeners batch.
+    class ListenersBatch
+      include Enumerable
+
+      # @param jobs_queue [JobsQueue]
+      # @return [ListenersBatch]
+      def initialize(jobs_queue)
+        @batch = App.subscription_groups.map do |subscription_group|
+          Connection::Listener.new(subscription_group, jobs_queue)
+        end
+      end
+
+      # Iterates over available listeners and yields each listener
+      # @param block [Proc] block we want to run
+      def each(&block)
+        @batch.each(&block)
+      end
+    end
+  end
+end

data/lib/karafka/connection/messages_buffer.rb

@@ -2,20 +2,26 @@
 
 module Karafka
   module Connection
-    # Buffer for messages.
-    # When message is added to this buffer, it gets assigned to an array with other messages from
-    # the same topic and partition.
+    # Buffer used to build and store karafka messages built based on raw librdkafka messages.
     #
-    # @note This buffer is NOT threadsafe.
+    # Why do we have two buffers? `RawMessagesBuffer` is used to store raw messages and to handle
+    #   cases related to partition revocation and reconnections. It is "internal" to the listening
+    #   process. `MessagesBuffer` on the other hand is used to "translate" those raw messages that
+    #   we know that are ok into Karafka messages and to simplify further work with them.
+    #
+    # While it adds a bit of overhead, it makes conceptual things much easier and it adds only two
+    #   simple hash iterations over messages batch.
+    #
+    # @note This buffer is NOT thread safe. We do not worry about it as we do not use it outside
+    #   of the main listener loop. It can be cleared after the jobs are scheduled with messages
+    #   it stores, because messages arrays are not "cleared" in any way directly and their
+    #   reference stays.
     class MessagesBuffer
       attr_reader :size
 
-      extend Forwardable
-
-      def_delegators :@groups, :each
-
-      # @return [Karafka::Connection::MessagesBuffer] buffer instance
-      def initialize
+      # @param subscription_group [Karafka::Routing::SubscriptionGroup]
+      def initialize(subscription_group)
+        @subscription_group = subscription_group
         @size = 0
         @groups = Hash.new do |topic_groups, topic|
           topic_groups[topic] = Hash.new do |partition_groups, partition|
@@ -24,64 +30,54 @@ module Karafka
         end
       end
 
-      # Adds a message to the buffer.
-      #
-      # @param message [Rdkafka::Consumer::Message] raw rdkafka message
-      # @return [Array<Rdkafka::Consumer::Message>] given partition topic sub-buffer array
-      def <<(message)
-        @size += 1
-        @groups[message.topic][message.partition] << message
-      end
+      # Remaps raw messages from the raw messages buffer to Karafka messages
+      # @param raw_messages_buffer [RawMessagesBuffer] buffer with raw messages
+      def remap(raw_messages_buffer)
+        clear unless @size.zero?
 
-      # Removes given topic and partition data out of the buffer
-      # This is used when there's a partition revocation
-      # @param topic [String] topic we're interested in
-      # @param partition [Integer] partition of which data we want to remove
-      def delete(topic, partition)
-        return unless @groups.key?(topic)
-        return unless @groups.fetch(topic).key?(partition)
+        # Since it happens "right after" we've received the messages, it is close enough it time
+        # to be used as the moment we received messages.
+        received_at = Time.now
 
-        topic_data = @groups.fetch(topic)
-        topic_data.delete(partition)
+        raw_messages_buffer.each do |topic, partition, messages|
+          @size += messages.count
 
-        recount!
+          ktopic = @subscription_group.topics.find(topic)
 
-        # If there are no more partitions to handle in a given topic, remove it completely
-        @groups.delete(topic) if topic_data.empty?
+          @groups[topic][partition] = messages.map do |message|
+            Messages::Builders::Message.call(
+              message,
+              ktopic,
+              received_at
+            )
+          end
+        end
       end
 
-      # Removes duplicated messages from the same partitions
-      # This should be used only when rebalance occurs, as we may get data again we already have
-      # due to the processing from the last offset. In cases like this, we may get same data
-      # again and we do want to ensure as few duplications as possible
-      def uniq!
-        @groups.each_value do |partitions|
-          partitions.each_value do |messages|
-            messages.uniq!(&:offset)
+      # Allows to iterate over all the topics and partitions messages
+      #
+      # @yieldparam [String] topic name
+      # @yieldparam [Integer] partition number
+      # @yieldparam [Array<Karafka::Messages::Message>] messages from a given topic partition
+      def each
+        @groups.each do |topic, partitions|
+          partitions.each do |partition, messages|
+            yield(topic, partition, messages)
           end
         end
-
-        recount!
       end
 
-      # Removes all the data from the buffer.
-      #
-      # @note We do not clear the whole groups hash but rather we clear the partition hashes, so
-      #   we save ourselves some objects allocations. We cannot clear the underlying arrays as they
-      #   may be used in other threads for data processing, thus if we would clear it, we could
-      #   potentially clear a raw messages array for a job that is in the jobs queue.
-      def clear
-        @size = 0
-        @groups.each_value(&:clear)
+      # @return [Boolean] is the buffer empty or does it contain any messages
+      def empty?
+        @size.zero?
       end
 
       private
 
-      # Updates the messages count if we performed any operations that could change the state
-      def recount!
-        @size = @groups.each_value.sum do |partitions|
-          partitions.each_value.map(&:count).sum
-        end
+      # Clears the buffer completely
+      def clear
+        @size = 0
+        @groups.clear
       end
     end
   end

data/lib/karafka/connection/raw_messages_buffer.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module Karafka
+  module Connection
+    # Buffer for raw librdkafka messages.
+    #
+    # When message is added to this buffer, it gets assigned to an array with other messages from
+    # the same topic and partition.
+    #
+    # @note This buffer is NOT threadsafe.
+    #
+    # @note We store data here in groups per topic partition to handle the revocation case, where
+    #    we may need to remove messages from a single topic partition.
+    class RawMessagesBuffer
+      attr_reader :size
+
+      # @return [Karafka::Connection::MessagesBuffer] buffer instance
+      def initialize
+        @size = 0
+        @groups = Hash.new do |topic_groups, topic|
+          topic_groups[topic] = Hash.new do |partition_groups, partition|
+            partition_groups[partition] = []
+          end
+        end
+      end
+
+      # Adds a message to the buffer.
+      #
+      # @param message [Rdkafka::Consumer::Message] raw rdkafka message
+      # @return [Array<Rdkafka::Consumer::Message>] given partition topic sub-buffer array
+      def <<(message)
+        @size += 1
+        @groups[message.topic][message.partition] << message
+      end
+
+      # Allows to iterate over all the topics and partitions messages
+      #
+      # @yieldparam [String] topic name
+      # @yieldparam [Integer] partition number
+      # @yieldparam [Array<Rdkafka::Consumer::Message>] topic partition aggregated results
+      def each
+        @groups.each do |topic, partitions|
+          partitions.each do |partition, messages|
+            yield(topic, partition, messages)
+          end
+        end
+      end
+
+      # Removes given topic and partition data out of the buffer
+      # This is used when there's a partition revocation
+      # @param topic [String] topic we're interested in
+      # @param partition [Integer] partition of which data we want to remove
+      def delete(topic, partition)
+        return unless @groups.key?(topic)
+        return unless @groups.fetch(topic).key?(partition)
+
+        topic_data = @groups.fetch(topic)
+        topic_data.delete(partition)
+
+        recount!
+
+        # If there are no more partitions to handle in a given topic, remove it completely
+        @groups.delete(topic) if topic_data.empty?
+      end
+
+      # Removes duplicated messages from the same partitions
+      # This should be used only when rebalance occurs, as we may get data again we already have
+      # due to the processing from the last offset. In cases like this, we may get same data
+      # again and we do want to ensure as few duplications as possible
+      def uniq!
+        @groups.each_value do |partitions|
+          partitions.each_value do |messages|
+            messages.uniq!(&:offset)
+          end
+        end
+
+        recount!
+      end
+
+      # Removes all the data from the buffer.
+      #
+      # @note We do not clear the whole groups hash but rather we clear the partition hashes, so
+      #   we save ourselves some objects allocations. We cannot clear the underlying arrays as they
+      #   may be used in other threads for data processing, thus if we would clear it, we could
+      #   potentially clear a raw messages array for a job that is in the jobs queue.
+      def clear
+        @size = 0
+        @groups.each_value(&:clear)
+      end
+
+      private
+
+      # Updates the messages count if we performed any operations that could change the state
+      def recount!
+        @size = @groups.each_value.sum do |partitions|
+          partitions.each_value.map(&:count).sum
+        end
+      end
+    end
+  end
+end

data/lib/karafka/contracts/config.rb

@@ -25,15 +25,17 @@ module Karafka
         required(:pause_max_timeout) { int? & gt?(0) }
         required(:pause_with_exponential_backoff).filled(:bool?)
         required(:shutdown_timeout) { int? & gt?(0) }
+        required(:max_wait_time) { int? & gt?(0) }
         required(:kafka).filled(:hash)
 
         # We validate internals just to be sure, that they are present and working
         required(:internal).schema do
           required(:routing_builder)
+          required(:subscription_groups_builder)
+          required(:jobs_builder)
           required(:status)
           required(:process)
           required(:scheduler)
-          required(:subscription_groups_builder)
         end
       end
 
@@ -53,6 +55,12 @@ module Karafka
           key(:pause_timeout).failure(:max_timeout_vs_pause_max_timeout)
         end
       end
+
+      rule(:shutdown_timeout, :max_wait_time) do
+        if values[:max_wait_time].to_i >= values[:shutdown_timeout].to_i
+          key(:shutdown_timeout).failure(:shutdown_timeout_vs_max_wait_time)
+        end
+      end
     end
   end
 end