karafka 2.0.0.alpha5 → 2.0.0.beta2

data/lib/karafka/connection/listener.rb

@@ -3,9 +3,13 @@
 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
+
       # @param subscription_group [Karafka::Routing::SubscriptionGroup]
       # @param jobs_queue [Karafka::Processing::JobsQueue] queue where we should push work
       # @return [Karafka::Connection::Listener] listener instance
@@ -15,6 +19,12 @@ module Karafka
         @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.
@@ -51,33 +61,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
+          poll_and_remap_messages
 
           Karafka.monitor.instrument(
             'connection.listener.fetch_loop.received',
             caller: self,
-            messages_buffer: messages_buffer
+            messages_buffer: @messages_buffer
           )
 
           # 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 distribute_revoke_lost_partitions_jobs
-
-          distribute_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
@@ -98,55 +130,88 @@ 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
-      def distribute_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 = @pauses_manager.fetch(topic, partition)
-            executor = @executors.fetch(topic, partition, pause)
-            @jobs_queue << Processing::Jobs::Revoked.new(executor)
+            pause_tracker = @pauses_manager.fetch(topic, partition)
+            executor = @executors.fetch(topic, partition, pause_tracker)
+            jobs << Processing::Jobs::Revoked.new(executor)
           end
         end
 
-        true
+        @scheduler.schedule_revocation(@jobs_queue, jobs)
       end
 
-      # Takes the messages per topic partition and enqueues processing jobs in threads.
+      # 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 << Processing::Jobs::Shutdown.new(executor)
+        end
+
+        @scheduler.schedule_shutdown(@jobs_queue, jobs)
+      end
+
+      # 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 distribute_partitions_jobs(messages_buffer)
-        messages_buffer.each do |topic, partition, messages|
+      # @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?
+
+        jobs = []
+
+        @messages_buffer.each do |topic, partition, messages|
           pause = @pauses_manager.fetch(topic, partition)
 
           next if pause.paused?
 
           executor = @executors.fetch(topic, partition, pause)
 
-          @jobs_queue << Processing::Jobs::Consume.new(executor, messages)
+          jobs << Processing::Jobs::Consume.new(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,16 +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
 
-      # @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|
@@ -20,11 +30,35 @@ module Karafka
         end
       end
 
-      # Iterates over aggregated data providing messages per topic partition.
+      # 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?
+
+        # 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
+
+        raw_messages_buffer.each do |topic, partition, messages|
+          @size += messages.count
+
+          ktopic = @subscription_group.topics.find(topic)
+
+          @groups[topic][partition] = messages.map do |message|
+            Messages::Builders::Message.call(
+              message,
+              ktopic,
+              received_at
+            )
+          end
+        end
+      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
+      # @yieldparam [Array<Karafka::Messages::Message>] messages from a given topic partition
       def each
         @groups.each do |topic, partitions|
           partitions.each do |partition, messages|
@@ -33,24 +67,17 @@ 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
+      # @return [Boolean] is the buffer empty or does it contain any messages
+      def empty?
+        @size.zero?
       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.
+      private
+
+      # Clears the buffer completely
       def clear
         @size = 0
-        @groups.each_value(&:clear)
+        @groups.clear
       end
     end
   end

data/lib/karafka/connection/pauses_manager.rb

@@ -12,11 +12,11 @@ module Karafka
         end
       end
 
-      # Creates or fetches pause of a given topic partition.
+      # Creates or fetches pause tracker of a given topic partition.
       #
       # @param topic [String] topic name
       # @param partition [Integer] partition number
-      # @return [Karafka::TimeTrackers::Pause] pause instance
+      # @return [Karafka::TimeTrackers::Pause] pause tracker instance
       def fetch(topic, partition)
         @pauses[topic][partition] ||= TimeTrackers::Pause.new(
           timeout: Karafka::App.config.pause_timeout,

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/connection/rebalance_manager.rb

@@ -9,35 +9,50 @@ module Karafka
     #
     # @note Since this does not happen really often, we try to stick with same objects for the
     #   empty states most of the time, so we don't create many objects during the manager life
+    #
+    # @note Internally in the rebalance manager we have a notion of lost partitions. Partitions
+    #   that are lost, are those that got revoked but did not get re-assigned back. We do not
+    #   expose this concept outside and we normalize to have them revoked, as it is irrelevant
+    #   from the rest of the code perspective as only those that are lost are truly revoked.
     class RebalanceManager
+      # Empty array for internal usage not to create new objects
+      EMPTY_ARRAY = [].freeze
+
+      private_constant :EMPTY_ARRAY
+
       # @return [RebalanceManager]
       def initialize
-        @assigned = {}
-        @revoked = {}
+        @assigned_partitions = {}
+        @revoked_partitions = {}
+        @lost_partitions = {}
       end
 
-      # @return [Hash<String, Array<Integer>>] hash where the keys are the names of topics for
-      #   which we've got new partitions assigned and array with ids of the partitions as the value
-      # @note Once assigned partitions are fetched, the state will be reset since the callbacks
-      #   for new assigned partitions are set only during a state change
-      def assigned_partitions
-        return @assigned if @assigned.empty?
-
-        result = @assigned.dup
-        @assigned.clear
-        result
+      # Resets the rebalance manager state
+      # This needs to be done before each polling loop as during the polling, the state may be
+      # changed
+      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 Once revoked partitions are fetched, the state will be reset since the callbacks
-      #   for new revoked partitions are set only during a state change
+      # @note We do not consider as lost topics and partitions that got revoked and assigned
       def revoked_partitions
-        return @revoked if @revoked.empty?
+        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
+      end
 
-        result = @revoked.dup
-        @revoked.clear
-        result
+      # @return [Boolean] true if any partitions were revoked
+      def revoked_partitions?
+        !revoked_partitions.empty?
       end
 
       # Callback that kicks in inside of rdkafka, when new partitions are assigned.
@@ -46,7 +61,7 @@ module Karafka
       # @param _ [Rdkafka::Consumer]
       # @param partitions [Rdkafka::Consumer::TopicPartitionList]
       def on_partitions_assigned(_, partitions)
-        @assigned = partitions.to_h.transform_values { |part| part.map(&:partition) }
+        @assigned_partitions = partitions.to_h.transform_values { |part| part.map(&:partition) }
       end
 
       # Callback that kicks in inside of rdkafka, when partitions are revoked.
@@ -55,7 +70,7 @@ module Karafka
       # @param _ [Rdkafka::Consumer]
       # @param partitions [Rdkafka::Consumer::TopicPartitionList]
       def on_partitions_revoked(_, partitions)
-        @revoked = partitions.to_h.transform_values { |part| part.map(&:partition) }
+        @revoked_partitions = partitions.to_h.transform_values { |part| part.map(&:partition) }
       end
     end
   end

data/lib/karafka/contracts/config.rb

@@ -25,6 +25,7 @@ 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
@@ -32,6 +33,7 @@ module Karafka
           required(:routing_builder)
           required(:status)
           required(:process)
+          required(:scheduler)
           required(:subscription_groups_builder)
         end
       end
@@ -52,6 +54,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

data/lib/karafka/helpers/async.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Karafka
+  module Helpers
+    # Allows a given class to run async in a separate thread. Provides also few methods we may
+    # want to use to control the underlying thread
+    #
+    # @note Thread running code needs to manage it's own exceptions. If they leak out, they will
+    #   abort thread on exception.
+    module Async
+      class << self
+        # Adds forwardable to redirect thread-based control methods to the underlying thread that
+        # runs the async operations
+        #
+        # @param base [Class] class we're including this module in
+        def included(base)
+          base.extend ::Forwardable
+
+          base.def_delegators :@thread, :join, :terminate, :alive?
+        end
+      end
+
+      # Runs the `#call` method in a new thread
+      def async_call
+        @thread = Thread.new do
+          Thread.current.abort_on_exception = true
+
+          call
+        end
+      end
+    end
+  end
+end

data/lib/karafka/instrumentation/monitor.rb

@@ -22,7 +22,8 @@ module Karafka
         app.stopping
         app.stopped
 
-        consumer.consume
+        consumer.prepared
+        consumer.consumed
         consumer.revoked
         consumer.shutdown
 

data/lib/karafka/messages/batch_metadata.rb

@@ -13,10 +13,33 @@ module Karafka
       :deserializer,
       :partition,
       :topic,
+      :created_at,
       :scheduled_at,
-      :consumption_lag,
-      :processing_lag,
+      :processed_at,
       keyword_init: true
-    )
+    ) do
+      # This lag describes how long did it take for a message to be consumed from the moment it was
+      # created
+      def consumption_lag
+        time_distance_in_ms(processed_at, created_at)
+      end
+
+      # This lag describes how long did a batch have to wait before it was picked up by one of the
+      # workers
+      def processing_lag
+        time_distance_in_ms(processed_at, scheduled_at)
+      end
+
+      private
+
+      # Computes time distance in between two times in ms
+      #
+      # @param time1 [Time]
+      # @param time2 [Time]
+      # @return [Integer] distance in between two times in ms
+      def time_distance_in_ms(time1, time2)
+        ((time1 - time2) * 1_000).round
+      end
+    end
   end
 end