karafka 2.0.0.beta2 → 2.0.0.beta5

data/lib/karafka/base_consumer.rb

@@ -10,44 +10,74 @@ module Karafka
     attr_accessor :messages
     # @return [Karafka::Connection::Client] kafka connection client
     attr_accessor :client
-    # @return [Karafka::TimeTrackers::Pause] current topic partition pause tracker
-    attr_accessor :pause_tracker
+    # @return [Karafka::Processing::Coordinator] coordinator
+    attr_accessor :coordinator
     # @return [Waterdrop::Producer] producer instance
     attr_accessor :producer
 
+    # Can be used to run preparation code
+    #
+    # @private
+    # @note This should not be used by the end users as it is part of the lifecycle of things but
+    #   not as part of the public api. This can act as a hook when creating non-blocking
+    #   consumers and doing other advanced stuff
+    def on_before_consume; end
+
     # Executes the default consumer flow.
     #
+    # @return [Boolean] true if there was no exception, otherwise false.
+    #
     # @note We keep the seek offset tracking, and use it to compensate for async offset flushing
     #   that may not yet kick in when error occurs. That way we pause always on the last processed
     #   message.
     def on_consume
       Karafka.monitor.instrument('consumer.consumed', caller: self) do
         consume
-
-        pause_tracker.reset
-
-        # Mark as consumed only if manual offset management is not on
-        next if topic.manual_offset_management
-
-        # We use the non-blocking one here. If someone needs the blocking one, can implement it
-        # with manual offset management
-        mark_as_consumed(messages.last)
       end
+
+      @coordinator.consumption(self).success!
     rescue StandardError => e
+      @coordinator.consumption(self).failure!
+
       Karafka.monitor.instrument(
         'error.occurred',
         error: e,
         caller: self,
         type: 'consumer.consume.error'
       )
+    ensure
+      # We need to decrease number of jobs that this coordinator coordinates as it has finished
+      @coordinator.decrement
+    end
+
+    # @private
+    # @note This should not be used by the end users as it is part of the lifecycle of things but
+    #   not as part of the public api.
+    def on_after_consume
+      return if revoked?
+
+      if @coordinator.success?
+        coordinator.pause_tracker.reset
+
+        # Mark as consumed only if manual offset management is not on
+        return if topic.manual_offset_management?
 
-      pause(@seek_offset || messages.first.offset)
+        # We use the non-blocking one here. If someone needs the blocking one, can implement it
+        # with manual offset management
+        mark_as_consumed(messages.last)
+      else
+        pause(@seek_offset || messages.first.offset)
+      end
     end
 
     # Trigger method for running on shutdown.
     #
     # @private
     def on_revoked
+      coordinator.revoke
+
+      resume
+
       Karafka.monitor.instrument('consumer.revoked', caller: self) do
         revoked
       end
@@ -76,31 +106,8 @@ module Karafka
       )
     end
 
-    # Can be used to run preparation code
-    #
-    # @private
-    # @note This should not be used by the end users as it is part of the lifecycle of things but
-    #   not as part of the public api. This can act as a hook when creating non-blocking
-    #   consumers and doing other advanced stuff
-    def on_prepared
-      Karafka.monitor.instrument('consumer.prepared', caller: self) do
-        prepared
-      end
-    rescue StandardError => e
-      Karafka.monitor.instrument(
-        'error.occurred',
-        error: e,
-        caller: self,
-        type: 'consumer.prepared.error'
-      )
-    end
-
     private
 
-    # Method that gets called in the blocking flow allowing to setup any type of resources or to
-    # send additional commands to Kafka before the proper execution starts.
-    def prepared; end
-
     # Method that will perform business logic and on data received from Kafka (it will consume
     #   the data)
     # @note This method needs bo be implemented in a subclass. We stub it here as a failover if
@@ -120,21 +127,40 @@ module Karafka
     # Marks message as consumed in an async way.
     #
     # @param message [Messages::Message] last successfully processed message.
+    # @return [Boolean] true if we were able to mark the offset, false otherwise. False indicates
+    #   that we were not able and that we have lost the partition.
+    #
     # @note We keep track of this offset in case we would mark as consumed and got error when
     #   processing another message. In case like this we do not pause on the message we've already
     #   processed but rather at the next one. This applies to both sync and async versions of this
     #   method.
     def mark_as_consumed(message)
-      client.mark_as_consumed(message)
+      unless client.mark_as_consumed(message)
+        coordinator.revoke
+
+        return false
+      end
+
       @seek_offset = message.offset + 1
+
+      true
     end
 
     # Marks message as consumed in a sync way.
     #
     # @param message [Messages::Message] last successfully processed message.
+    # @return [Boolean] true if we were able to mark the offset, false otherwise. False indicates
+    #   that we were not able and that we have lost the partition.
     def mark_as_consumed!(message)
-      client.mark_as_consumed!(message)
+      unless client.mark_as_consumed!(message)
+        coordinator.revoke
+
+        return false
+      end
+
       @seek_offset = message.offset + 1
+
+      true
     end
 
     # Pauses processing on a given offset for the current topic partition
@@ -144,23 +170,20 @@ module Karafka
     # @param timeout [Integer, nil] how long in milliseconds do we want to pause or nil to use the
     #   default exponential pausing strategy defined for retries
     def pause(offset, timeout = nil)
+      timeout ? coordinator.pause_tracker.pause(timeout) : coordinator.pause_tracker.pause
+
       client.pause(
         messages.metadata.topic,
         messages.metadata.partition,
         offset
       )
-
-      timeout ? pause_tracker.pause(timeout) : pause_tracker.pause
     end
 
     # Resumes processing of the current topic partition
     def resume
-      client.resume(
-        messages.metadata.topic,
-        messages.metadata.partition
-      )
-
-      pause_tracker.expire
+      # This is sufficient to expire a partition pause, as with it will be resumed by the listener
+      # thread before the next poll.
+      coordinator.pause_tracker.expire
     end
 
     # Seeks in the context of current topic and partition
@@ -175,5 +198,12 @@ module Karafka
         )
       )
     end
+
+    # @return [Boolean] true if partition was revoked from the current consumer
+    # @note We know that partition got revoked because when we try to mark message as consumed,
+    #   unless if is successful, it will return false
+    def revoked?
+      coordinator.revoked?
+    end
   end
 end

data/lib/karafka/connection/client.rb

@@ -36,6 +36,12 @@ module Karafka
         # Marks if we need to offset. If we did not store offsets, we should not commit the offset
         # position as it will crash rdkafka
         @offsetting = false
+        # We need to keep track of what we have paused for resuming
+        # In case we loose partition, we still need to resume it, otherwise it won't be fetched
+        # again if we get reassigned to it later on. We need to keep them as after revocation we
+        # no longer may be able to fetch them from Kafka. We could build them but it is easier
+        # to just keep them here and use if needed when cannot be obtained
+        @paused_tpls = Hash.new { |h, k| h[k] = {} }
       end
 
       # Fetches messages within boundaries defined by the settings (time, size, topics, etc).
@@ -45,12 +51,13 @@ module Karafka
       # @note This method should not be executed from many threads at the same time
       def batch_poll
         time_poll = TimeTrackers::Poll.new(@subscription_group.max_wait_time)
-        time_poll.start
 
         @buffer.clear
         @rebalance_manager.clear
 
         loop do
+          time_poll.start
+
           # Don't fetch more messages if we do not have any time left
           break if time_poll.exceeded?
           # Don't fetch more messages if we've fetched max as we've wanted
@@ -69,7 +76,11 @@ module Karafka
           # If partition revocation happens, we need to remove messages from revoked partitions
           # as well as ensure we do not have duplicated due to the offset reset for partitions
           # that we got assigned
-          remove_revoked_and_duplicated_messages if @rebalance_manager.revoked_partitions?
+          # We also do early break, so the information about rebalance is used as soon as possible
+          if @rebalance_manager.changed?
+            remove_revoked_and_duplicated_messages
+            break
+          end
 
           # Finally once we've (potentially) removed revoked, etc, if no messages were returned
           # we can break.
@@ -86,8 +97,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 +114,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 +131,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 +151,21 @@ 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)
+
+        # Here we do not use our cached tpls because we should not try to pause something we do
+        # not own anymore.
         tpl = topic_partition_list(topic, partition)
 
         return unless tpl
 
-        @kafka.pause(tpl)
+        @paused_tpls[topic][partition] = tpl
 
-        pause_msg = Messages::Seek.new(topic, partition, offset)
+        @kafka.pause(tpl)
 
-        seek(pause_msg)
+        @kafka.seek(pause_msg)
       ensure
         @mutex.unlock
       end
@@ -166,9 +179,18 @@ module Karafka
 
         return if @closed
 
-        tpl = topic_partition_list(topic, partition)
+        # 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)
+
+        # If we were not able, let's try to reuse the one we have (if we have)
+        tpl = topic_partition_list(topic, partition) || @paused_tpls[topic][partition]
 
         return unless tpl
+        # If we did not have it, it means we never paused this partition, thus no resume should
+        # happen in the first place
+        return unless @paused_tpls[topic].delete(partition)
 
         @kafka.resume(tpl)
       ensure
@@ -187,6 +209,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 +219,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
 
@@ -208,17 +233,51 @@ module Karafka
         @mutex.synchronize do
           @closed = false
           @offsetting = false
+          @paused_tpls.clear
           @kafka = build_consumer
         end
       end
 
       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 +286,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 +340,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
@@ -323,7 +389,7 @@ module Karafka
       # we are no longer responsible in a given process for processing those messages and they
       # should have been picked up by a different process.
       def remove_revoked_and_duplicated_messages
-        @rebalance_manager.revoked_partitions.each do |topic, partitions|
+        @rebalance_manager.lost_partitions.each do |topic, partitions|
           partitions.each do |partition|
             @buffer.delete(topic, partition)
           end

data/lib/karafka/connection/listener.rb

@@ -10,17 +10,23 @@ module Karafka
     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
-        @pauses_manager = PausesManager.new
+        @jobs_builder = ::Karafka::App.config.internal.processing.jobs_builder
+        @coordinators = Processing::CoordinatorsBuffer.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
+        @scheduler = ::Karafka::App.config.internal.processing.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.
@@ -62,16 +68,20 @@ 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.
-          poll_and_remap_messages
-
           Karafka.monitor.instrument(
             'connection.listener.fetch_loop.received',
             caller: self,
             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
+
+          # This will ensure, that in the next poll, we continue processing (if we get them back)
+          # partitions that we have paused
+          resume_assigned_partitions
 
           # 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
@@ -80,6 +90,9 @@ module Karafka
           build_and_schedule_revoke_lost_partitions_jobs
 
           # We wait only on jobs from our subscription group. Other groups are independent.
+          # This will block on revoked jobs until they are finished. Those are not meant to last
+          # long and should not have any bigger impact on the system. Doing this in a blocking way
+          # simplifies the overall design and prevents from race conditions
           wait
 
           build_and_schedule_consumption_jobs
@@ -130,7 +143,7 @@ module Karafka
 
       # Resumes processing of partitions that were paused due to an error.
       def resume_paused_partitions
-        @pauses_manager.resume do |topic, partition|
+        @coordinators.resume do |topic, partition|
           @client.resume(topic, partition)
         end
       end
@@ -146,9 +159,23 @@ module Karafka
 
         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 << Processing::Jobs::Revoked.new(executor)
+            # We revoke the coordinator here, so we do not have to revoke it in the revoke job
+            # itself (this happens prior to scheduling those jobs)
+            @coordinators.revoke(topic, partition)
+
+            # There may be a case where we have lost partition of which data we have never
+            # processed (if it was assigned and revoked really fast), thus we may not have it
+            # here. In cases like this, we do not run a revocation job
+            @executors.find_all(topic, partition).each do |executor|
+              jobs << @jobs_builder.revoked(executor)
+            end
+
+            # We need to remove all the executors of a given topic partition that we have lost, so
+            # next time we pick up it's work, new executors kick in. This may be needed especially
+            # for LRJ where we could end up with a race condition
+            # This revocation needs to happen after the jobs are scheduled, otherwise they would
+            # be scheduled with new executors instead of old
+            @executors.revoke(topic, partition)
           end
         end
 
@@ -160,7 +187,7 @@ module Karafka
         jobs = []
 
         @executors.each do |_, _, executor|
-          jobs << Processing::Jobs::Shutdown.new(executor)
+          jobs << @jobs_builder.shutdown(executor)
         end
 
         @scheduler.schedule_shutdown(@jobs_queue, jobs)
@@ -177,6 +204,17 @@ module Karafka
         )
       end
 
+      # Revoked partition needs to be resumed if we were processing them earlier. This will do
+      # nothing to things that we are planning to process. Without this, things we get
+      # re-assigned would not be polled.
+      def resume_assigned_partitions
+        @client.rebalance_manager.assigned_partitions.each do |topic, partitions|
+          partitions.each do |partition|
+            @client.resume(topic, partition)
+          end
+        end
+      end
+
       # Takes the messages per topic partition and enqueues processing jobs in threads using
       # given scheduler.
       def build_and_schedule_consumption_jobs
@@ -185,13 +223,17 @@ module Karafka
         jobs = []
 
         @messages_buffer.each do |topic, partition, messages|
-          pause = @pauses_manager.fetch(topic, partition)
+          coordinator = @coordinators.find_or_create(topic, partition)
+
+          # Start work coordination for this topic partition
+          coordinator.start
 
-          next if pause.paused?
+          # Count the job we're going to create here
+          coordinator.increment
 
-          executor = @executors.fetch(topic, partition, pause)
+          executor = @executors.find_or_create(topic, partition, 0)
 
-          jobs << Processing::Jobs::Consume.new(executor, messages)
+          jobs << @jobs_builder.consume(executor, messages, coordinator)
         end
 
         @scheduler.schedule_consumption(@jobs_queue, jobs)
@@ -227,7 +269,7 @@ module Karafka
         @jobs_queue.wait(@subscription_group.id)
         @jobs_queue.clear(@subscription_group.id)
         @client.reset
-        @pauses_manager = PausesManager.new
+        @coordinators.reset
         @executors = Processing::ExecutorsBuffer.new(@client, @subscription_group)
       end
     end

data/lib/karafka/connection/pauses_manager.rb

@@ -25,6 +25,14 @@ module Karafka
         )
       end
 
+      # Revokes pause tracker for a given topic partition
+      #
+      # @param topic [String] topic name
+      # @param partition [Integer] partition number
+      def revoke(topic, partition)
+        @pauses[topic].delete(partition)
+      end
+
       # Resumes processing of partitions for which pause time has ended.
       #
       # @yieldparam [String] topic name

data/lib/karafka/connection/rebalance_manager.rb

@@ -18,13 +18,15 @@ module Karafka
       # Empty array for internal usage not to create new objects
       EMPTY_ARRAY = [].freeze
 
+      attr_reader :assigned_partitions, :revoked_partitions
+
       private_constant :EMPTY_ARRAY
 
       # @return [RebalanceManager]
       def initialize
         @assigned_partitions = {}
         @revoked_partitions = {}
-        @lost_partitions = {}
+        @changed = false
       end
 
       # Resets the rebalance manager state
@@ -33,26 +35,12 @@ module Karafka
       def clear
         @assigned_partitions.clear
         @revoked_partitions.clear
-        @lost_partitions.clear
-      end
-
-      # @return [Hash<String, Array<Integer>>] hash where the keys are the names of topics for
-      #   which we've lost partitions and array with ids of the partitions as the value
-      # @note We do not consider as lost topics and partitions that got revoked and assigned
-      def revoked_partitions
-        return @revoked_partitions if @revoked_partitions.empty?
-        return @lost_partitions unless @lost_partitions.empty?
-
-        @revoked_partitions.each do |topic, partitions|
-          @lost_partitions[topic] = partitions - @assigned_partitions.fetch(topic, EMPTY_ARRAY)
-        end
-
-        @lost_partitions
+        @changed = false
       end
 
-      # @return [Boolean] true if any partitions were revoked
-      def revoked_partitions?
-        !revoked_partitions.empty?
+      # @return [Boolean] indicates a state change in the partitions assignment
+      def changed?
+        @changed
       end
 
       # Callback that kicks in inside of rdkafka, when new partitions are assigned.
@@ -62,6 +50,7 @@ module Karafka
       # @param partitions [Rdkafka::Consumer::TopicPartitionList]
       def on_partitions_assigned(_, partitions)
         @assigned_partitions = partitions.to_h.transform_values { |part| part.map(&:partition) }
+        @changed = true
       end
 
       # Callback that kicks in inside of rdkafka, when partitions are revoked.
@@ -71,6 +60,18 @@ module Karafka
       # @param partitions [Rdkafka::Consumer::TopicPartitionList]
       def on_partitions_revoked(_, partitions)
         @revoked_partitions = partitions.to_h.transform_values { |part| part.map(&:partition) }
+        @changed = true
+      end
+
+      # We consider as lost only partitions that were taken away and not re-assigned back to us
+      def lost_partitions
+        lost_partitions = {}
+
+        revoked_partitions.each do |topic, partitions|
+          lost_partitions[topic] = partitions - assigned_partitions.fetch(topic, EMPTY_ARRAY)
+        end
+
+        lost_partitions
       end
     end
   end