karafka 2.2.12 → 2.2.14

data/lib/karafka/admin.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module Karafka
-  # Simple admin actions that we can perform via Karafka on our Kafka cluster
+  # Admin actions that we can perform via Karafka on our Kafka cluster
   #
   # @note It always initializes a new admin instance as we want to ensure it is always closed
   #   Since admin actions are not performed that often, that should be ok.
@@ -137,6 +137,109 @@ module Karafka
         end
       end
 
+      # Moves the offset on a given consumer group and provided topic to the requested location
+      #
+      # @param consumer_group_id [String] id of the consumer group for which we want to move the
+      #   existing offset
+      # @param topics_with_partitions_and_offsets [Hash] Hash with list of topics and settings to
+      #   where to move given consumer. It allows us to move particular partitions or whole topics
+      #   if we want to reset all partitions to for example a point in time.
+      #
+      # @note This method should **not** be executed on a running consumer group as it creates a
+      #   "fake" consumer and uses it to move offsets.
+      #
+      # @example Move a single topic partition nr 1 offset to 100
+      #   Karafka::Admin.seek_consumer_group('group-id', { 'topic' => { 1 => 100 } })
+      #
+      # @example Move offsets on all partitions of a topic to 100
+      #   Karafka::Admin.seek_consumer_group('group-id', { 'topic' => 100 })
+      #
+      # @example Move offset to 5 seconds ago on partition 2
+      #   Karafka::Admin.seek_consumer_group('group-id', { 'topic' => { 2 => 5.seconds.ago } })
+      def seek_consumer_group(consumer_group_id, topics_with_partitions_and_offsets)
+        tpl_base = {}
+
+        # Normalize the data so we always have all partitions and topics in the same format
+        # That is in a format where we have topics and all partitions with their per partition
+        # assigned offsets
+        topics_with_partitions_and_offsets.each do |topic, partitions_with_offsets|
+          tpl_base[topic] = {}
+
+          if partitions_with_offsets.is_a?(Hash)
+            tpl_base[topic] = partitions_with_offsets
+          else
+            topic(topic)[:partition_count].times do |partition|
+              tpl_base[topic][partition] = partitions_with_offsets
+            end
+          end
+        end
+
+        tpl = Rdkafka::Consumer::TopicPartitionList.new
+        # In case of time based location, we need to to a pre-resolution, that's why we keep it
+        # separately
+        time_tpl = Rdkafka::Consumer::TopicPartitionList.new
+
+        # Distribute properly the offset type
+        tpl_base.each do |topic, partitions_with_offsets|
+          partitions_with_offsets.each do |partition, offset|
+            target = offset.is_a?(Time) ? time_tpl : tpl
+            target.add_topic_and_partitions_with_offsets(topic, [[partition, offset]])
+          end
+        end
+
+        # We set this that way so we can impersonate this consumer group and seek where we want
+        mapped_consumer_group_id = app_config.consumer_mapper.call(consumer_group_id)
+        settings = { 'group.id': mapped_consumer_group_id }
+
+        with_consumer(settings) do |consumer|
+          # If we have any time based stuff to resolve, we need to do it prior to commits
+          unless time_tpl.empty?
+            real_offsets = consumer.offsets_for_times(time_tpl)
+
+            real_offsets.to_h.each do |name, results|
+              results.each do |result|
+                raise(Errors::InvalidTimeBasedOffsetError) unless result
+
+                partition = result.partition
+
+                # Negative offset means we're beyond last message and we need to query for the
+                # high watermark offset to get the most recent offset and move there
+                if result.offset.negative?
+                  _, offset = consumer.query_watermark_offsets(name, result.partition)
+                else
+                  # If we get an offset, it means there existed a message close to this time
+                  # location
+                  offset = result.offset
+                end
+
+                # Since now we have proper offsets, we can add this to the final tpl for commit
+                tpl.add_topic_and_partitions_with_offsets(name, [[partition, offset]])
+              end
+            end
+          end
+
+          consumer.commit(tpl, false)
+        end
+      end
+
+      # Removes given consumer group (if exists)
+      #
+      # @param consumer_group_id [String] consumer group name without the mapper name (if any used)
+      #
+      # @note Please note, Karafka will apply the consumer group mapper on the provided consumer
+      #   group.
+      #
+      # @note This method should not be used on a running consumer group as it will not yield any
+      #   results.
+      def delete_consumer_group(consumer_group_id)
+        mapped_consumer_group_id = app_config.consumer_mapper.call(consumer_group_id)
+
+        with_admin do |admin|
+          handler = admin.delete_group(mapped_consumer_group_id)
+          handler.wait(max_wait_timeout: app_config.admin.max_wait_time)
+        end
+      end
+
       # Fetches the watermark offsets for a given topic partition
       #
       # @param name [String, Symbol] topic name
@@ -264,7 +367,7 @@ module Karafka
       # @param settings [Hash] extra settings for config (if needed)
       # @return [::Rdkafka::Config] rdkafka config
       def config(type, settings)
-        group_id = app_config.consumer_mapper.call(
+        mapped_admin_group_id = app_config.consumer_mapper.call(
           app_config.admin.group_id
         )
 
@@ -272,8 +375,11 @@ module Karafka
           .kafka
           .then(&:dup)
           .merge(app_config.admin.kafka)
+          .tap { |config| config[:'group.id'] = mapped_admin_group_id }
+          # We merge after setting the group id so it can be altered if needed
+          # In general in admin we only should alter it when we need to impersonate a given
+          # consumer group or do something similar
           .merge!(settings)
-          .tap { |config| config[:'group.id'] = group_id }
           .then { |config| Karafka::Setup::AttributesMap.public_send(type, config) }
           .then { |config| ::Rdkafka::Config.new(config) }
       end

data/lib/karafka/app.rb

@@ -36,6 +36,13 @@ module Karafka
       # Just a nicer name for the consumer groups
       alias routes consumer_groups
 
+      # Returns current assignments of this process. Both topics and partitions
+      #
+      # @return [Hash<Karafka::Routing::Topic, Array<Integer>>]
+      def assignments
+        Instrumentation::AssignmentsTracker.instance.current
+      end
+
       # Allow for easier status management via `Karafka::App` by aliasing status methods here
       Status::STATES.each do |state, transition|
         class_eval <<~RUBY, __FILE__, __LINE__ + 1

data/lib/karafka/base_consumer.rb

@@ -34,16 +34,9 @@ module Karafka
     # @note This should not be used by the end users as it is part of the lifecycle of things and
     #   not as a part of the public api. This should not perform any extensive operations as it is
     #   blocking and running in the listener thread.
-    def on_before_enqueue
+    def on_before_schedule_consume
       @used = true
-      handle_before_enqueue
-    rescue StandardError => e
-      Karafka.monitor.instrument(
-        'error.occurred',
-        error: e,
-        caller: self,
-        type: 'consumer.before_enqueue.error'
-      )
+      handle_before_schedule_consume
     end
 
     # Can be used to run preparation code in the worker
@@ -59,13 +52,6 @@ module Karafka
       # We run this after the full metadata setup, so we can use all the messages information
       # if needed
       handle_before_consume
-    rescue StandardError => e
-      Karafka.monitor.instrument(
-        'error.occurred',
-        error: e,
-        caller: self,
-        type: 'consumer.before_consume.error'
-      )
     end
 
     # Executes the default consumer flow.
@@ -94,13 +80,13 @@ module Karafka
     #   not as part of the public api.
     def on_after_consume
       handle_after_consume
-    rescue StandardError => e
-      Karafka.monitor.instrument(
-        'error.occurred',
-        error: e,
-        caller: self,
-        type: 'consumer.after_consume.error'
-      )
+    end
+
+    # Can be used to run code prior to scheduling of idle execution
+    #
+    # @private
+    def on_before_schedule_idle
+      handle_before_schedule_idle
     end
 
     # Trigger method for running on idle runs without messages
@@ -108,13 +94,13 @@ module Karafka
     # @private
     def on_idle
       handle_idle
-    rescue StandardError => e
-      Karafka.monitor.instrument(
-        'error.occurred',
-        error: e,
-        caller: self,
-        type: 'consumer.idle.error'
-      )
+    end
+
+    # Can be used to run code prior to scheduling of revoked execution
+    #
+    # @private
+    def on_before_schedule_revoked
+      handle_before_schedule_revoked
     end
 
     # Trigger method for running on partition revocation.
@@ -131,6 +117,13 @@ module Karafka
       )
     end
 
+    # Can be used to run code prior to scheduling of revoked execution
+    #
+    # @private
+    def on_before_schedule_shutdown
+      handle_before_schedule_shutdown
+    end
+
     # Trigger method for running on shutdown.
     #
     # @private

data/lib/karafka/connection/client.rb

@@ -45,10 +45,7 @@ module Karafka
         @buffer = RawMessagesBuffer.new
         @tick_interval = ::Karafka::App.config.internal.tick_interval
         @rebalance_manager = RebalanceManager.new(@subscription_group.id)
-        @rebalance_callback = Instrumentation::Callbacks::Rebalance.new(
-          @subscription_group.id,
-          @subscription_group.consumer_group.id
-        )
+        @rebalance_callback = Instrumentation::Callbacks::Rebalance.new(@subscription_group)
         @events_poller = Helpers::IntervalRunner.new { events_poll }
         @kafka = build_consumer
         # There are few operations that can happen in parallel from the listener threads as well
@@ -309,12 +306,18 @@ module Karafka
 
       # Closes and resets the client completely.
       def reset
-        close
-
-        @events_poller.reset
-        @closed = false
-        @paused_tpls.clear
-        @kafka = build_consumer
+        Karafka.monitor.instrument(
+          'client.reset',
+          caller: self,
+          subscription_group: @subscription_group
+        ) do
+          close
+
+          @events_poller.reset
+          @closed = false
+          @paused_tpls.clear
+          @kafka = build_consumer
+        end
       end
 
       # Runs a single poll on the main queue and consumer queue ignoring all the potential errors

data/lib/karafka/connection/consumer_group_coordinator.rb

@@ -16,7 +16,7 @@ module Karafka
     class ConsumerGroupCoordinator
       # @param group_size [Integer] number of separate subscription groups in a consumer group
       def initialize(group_size)
-        @shutdown_lock = Mutex.new
+        @shutdown_mutex = Mutex.new
         @group_size = group_size
         @finished = Set.new
       end
@@ -30,12 +30,12 @@ module Karafka
       # @return [Boolean] can we start shutdown on a given listener
       # @note If true, will also obtain a lock so no-one else will be closing the same time we do
       def shutdown?
-        finished? && @shutdown_lock.try_lock
+        finished? && @shutdown_mutex.try_lock
       end
 
       # Unlocks the shutdown lock
       def unlock
-        @shutdown_lock.unlock if @shutdown_lock.owned?
+        @shutdown_mutex.unlock if @shutdown_mutex.owned?
       end
 
       # Marks given listener as finished

data/lib/karafka/connection/listener.rb

@@ -23,8 +23,9 @@ module Karafka
       # @param consumer_group_coordinator [Karafka::Connection::ConsumerGroupCoordinator]
       # @param subscription_group [Karafka::Routing::SubscriptionGroup]
       # @param jobs_queue [Karafka::Processing::JobsQueue] queue where we should push work
+      # @param scheduler [Karafka::Processing::Scheduler] scheduler we want to use
       # @return [Karafka::Connection::Listener] listener instance
-      def initialize(consumer_group_coordinator, subscription_group, jobs_queue)
+      def initialize(consumer_group_coordinator, subscription_group, jobs_queue, scheduler)
         proc_config = ::Karafka::App.config.internal.processing
 
         @id = SecureRandom.hex(6)
@@ -36,8 +37,7 @@ module Karafka
         @executors = Processing::ExecutorsBuffer.new(@client, subscription_group)
         @jobs_builder = proc_config.jobs_builder
         @partitioner = proc_config.partitioner_class.new(subscription_group)
-        # We reference scheduler here as it is much faster than fetching this each time
-        @scheduler = proc_config.scheduler
+        @scheduler = scheduler
         @events_poller = Helpers::IntervalRunner.new { @client.events_poll }
         # 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
@@ -45,6 +45,8 @@ module Karafka
         @messages_buffer = MessagesBuffer.new(subscription_group)
         @mutex = Mutex.new
         @stopped = false
+
+        @jobs_queue.register(@subscription_group.id)
       end
 
       # Runs the main listener fetch loop.
@@ -230,7 +232,7 @@ module Karafka
             # here. In cases like this, we do not run a revocation job
             @executors.find_all(topic, partition).each do |executor|
               job = @jobs_builder.revoked(executor)
-              job.before_enqueue
+              job.before_schedule
               jobs << job
             end
 
@@ -243,7 +245,7 @@ module Karafka
           end
         end
 
-        @scheduler.schedule_revocation(@jobs_queue, jobs)
+        @scheduler.on_schedule_revocation(jobs)
       end
 
       # Enqueues the shutdown jobs for all the executors that exist in our subscription group
@@ -252,11 +254,11 @@ module Karafka
 
         @executors.each do |executor|
           job = @jobs_builder.shutdown(executor)
-          job.before_enqueue
+          job.before_schedule
           jobs << job
         end
 
-        @scheduler.schedule_shutdown(@jobs_queue, jobs)
+        @scheduler.on_schedule_shutdown(jobs)
       end
 
       # Polls messages within the time and amount boundaries defined in the settings and then
@@ -296,14 +298,17 @@ module Karafka
           end
         end
 
-        jobs.each(&:before_enqueue)
+        jobs.each(&:before_schedule)
 
-        @scheduler.schedule_consumption(@jobs_queue, jobs)
+        @scheduler.on_schedule_consumption(jobs)
       end
 
       # Waits for all the jobs from a given subscription group to finish before moving forward
       def wait
-        @jobs_queue.wait(@subscription_group.id) { @events_poller.call }
+        @jobs_queue.wait(@subscription_group.id) do
+          @events_poller.call
+          @scheduler.on_manage
+        end
       end
 
       # Waits without blocking the polling
@@ -319,6 +324,8 @@ module Karafka
       def wait_pinging(wait_until:, after_ping: -> {})
         until wait_until.call
           @client.ping
+          @scheduler.on_manage
+
           after_ping.call
           sleep(0.2)
         end
@@ -334,6 +341,7 @@ module Karafka
         # resetting.
         @jobs_queue.wait(@subscription_group.id)
         @jobs_queue.clear(@subscription_group.id)
+        @scheduler.on_clear(@subscription_group.id)
         @events_poller.reset
         @client.reset
         @coordinators.reset

data/lib/karafka/connection/listeners_batch.rb

@@ -11,6 +11,10 @@ module Karafka
       # @param jobs_queue [JobsQueue]
       # @return [ListenersBatch]
       def initialize(jobs_queue)
+        # We need one scheduler for all the listeners because in case of complex schedulers, they
+        # should be able to distribute work whenever any work is done in any of the listeners
+        scheduler = App.config.internal.processing.scheduler_class.new(jobs_queue)
+
         @coordinators = []
 
         @batch = App.subscription_groups.flat_map do |_consumer_group, subscription_groups|
@@ -24,7 +28,8 @@ module Karafka
             Connection::Listener.new(
               consumer_group_coordinator,
               subscription_group,
-              jobs_queue
+              jobs_queue,
+              scheduler
             )
           end
         end

data/lib/karafka/contracts/config.rb

@@ -73,7 +73,8 @@ module Karafka
 
         nested(:processing) do
           required(:jobs_builder) { |val| !val.nil? }
-          required(:scheduler) { |val| !val.nil? }
+          required(:jobs_queue_class) { |val| !val.nil? }
+          required(:scheduler_class) { |val| !val.nil? }
           required(:coordinator_class) { |val| !val.nil? }
           required(:partitioner_class) { |val| !val.nil? }
           required(:strategy_selector) { |val| !val.nil? }

data/lib/karafka/instrumentation/assignments_tracker.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module Karafka
+  module Instrumentation
+    # Keeps track of active assignments and materializes them by returning the routing topics
+    # with appropriate partitions that are assigned at a given moment
+    #
+    # It is auto-subscribed as part of Karafka itself.
+    #
+    # It is not heavy from the computational point of view, as it only operates during rebalances.
+    #
+    # We keep assignments as flat topics structure because we can go from topics to both
+    # subscription and consumer groups if needed.
+    class AssignmentsTracker
+      include Singleton
+
+      def initialize
+        @mutex = Mutex.new
+        @assignments = Hash.new { |hash, key| hash[key] = [] }
+      end
+
+      # Returns all the active/current assignments of this given process
+      #
+      # @return [Hash<Karafka::Routing::Topic, Array<Integer>>]
+      #
+      # @note Keep in mind, that those assignments can change any time, especially when working
+      #   with multiple consumer groups or subscription groups.
+      #
+      # @note We return a copy because we modify internals and we do not want user to tamper with
+      #   the data accidentally
+      def current
+        assignments = {}
+
+        @assignments.each do |topic, partitions|
+          assignments[topic] = partitions.dup.freeze
+        end
+
+        assignments.freeze
+      end
+
+      # Clears all the assignments
+      def clear
+        @mutex.synchronize do
+          @assignments.clear
+        end
+      end
+
+      # When client is under reset due to critical issues, remove all of its assignments as we will
+      #   get a new set of assignments
+      # @param event [Karafka::Core::Monitoring::Event]
+      def on_client_reset(event)
+        sg = event[:subscription_group]
+
+        @mutex.synchronize do
+          @assignments.delete_if do |topic, _partitions|
+            topic.subscription_group.id == sg.id
+          end
+        end
+      end
+
+      # Removes partitions from the current assignments hash
+      #
+      # @param event [Karafka::Core::Monitoring::Event]
+      def on_rebalance_partitions_revoked(event)
+        sg = event[:subscription_group]
+
+        @mutex.synchronize do
+          event[:tpl].to_h.each do |topic, partitions|
+            topic = sg.topics.find(topic)
+
+            @assignments[topic] -= partitions.map(&:partition)
+            @assignments[topic].sort!
+            # Remove completely topics for which we do not have any assignments left
+            @assignments.delete_if { |_topic, cur_partitions| cur_partitions.empty? }
+          end
+        end
+      end
+
+      # # Adds partitions to the current assignments hash
+      #
+      # @param event [Karafka::Core::Monitoring::Event]
+      def on_rebalance_partitions_assigned(event)
+        sg = event[:subscription_group]
+
+        @mutex.synchronize do
+          event[:tpl].to_h.each do |topic, partitions|
+            topic = sg.topics.find(topic)
+
+            @assignments[topic] += partitions.map(&:partition)
+            @assignments[topic].sort!
+          end
+        end
+      end
+    end
+  end
+end

data/lib/karafka/instrumentation/callbacks/rebalance.rb

@@ -6,11 +6,10 @@ module Karafka
       # Callback that connects to the librdkafka rebalance callback and converts those events into
       # our internal events
       class Rebalance
-        # @param subscription_group_id [String] id of the current subscription group instance
-        # @param consumer_group_id [String] id of the current consumer group
-        def initialize(subscription_group_id, consumer_group_id)
-          @subscription_group_id = subscription_group_id
-          @consumer_group_id = consumer_group_id
+        # @param subscription_group [Karafka::Routes::SubscriptionGroup] subscription group for
+        #   which we want to manage rebalances
+        def initialize(subscription_group)
+          @subscription_group = subscription_group
         end
 
         # Publishes an event that partitions are going to be revoked.
@@ -53,8 +52,12 @@ module Karafka
           ::Karafka.monitor.instrument(
             "rebalance.#{name}",
             caller: self,
-            subscription_group_id: @subscription_group_id,
-            consumer_group_id: @consumer_group_id,
+            # We keep the id references here for backwards compatibility as some of the monitors
+            # may use the id references
+            subscription_group_id: @subscription_group.id,
+            subscription_group: @subscription_group,
+            consumer_group_id: @subscription_group.consumer_group.id,
+            consumer_group: @subscription_group.consumer_group,
             tpl: tpl
           )
         end

data/lib/karafka/instrumentation/logger_listener.rb

@@ -241,15 +241,6 @@ module Karafka
         when 'consumer.revoked.error'
           error "Consumer on revoked failed due to an error: #{error}"
           error details
-        when 'consumer.before_enqueue.error'
-          error "Consumer before enqueue failed due to an error: #{error}"
-          error details
-        when 'consumer.before_consume.error'
-          error "Consumer before consume failed due to an error: #{error}"
-          error details
-        when 'consumer.after_consume.error'
-          error "Consumer after consume failed due to an error: #{error}"
-          error details
         when 'consumer.idle.error'
           error "Consumer idle failed due to an error: #{error}"
           error details

data/lib/karafka/instrumentation/notifications.rb

@@ -20,6 +20,7 @@ module Karafka
         active_job.consume
         active_job.consumed
 
+        app.initializing
         app.initialized
         app.running
         app.quieting
@@ -30,26 +31,28 @@ module Karafka
 
         client.pause
         client.resume
+        client.reset
 
         connection.listener.before_fetch_loop
         connection.listener.fetch_loop
         connection.listener.fetch_loop.received
 
-        connection.client.poll.error
-        connection.client.unsubscribe.error
-
         rebalance.partitions_assign
         rebalance.partitions_assigned
         rebalance.partitions_revoke
         rebalance.partitions_revoked
 
+        consumer.before_schedule_consume
         consumer.consume
         consumer.consumed
         consumer.consuming.pause
         consumer.consuming.retry
+        consumer.before_schedule_idle
         consumer.idle
+        consumer.before_schedule_revoked
         consumer.revoke
         consumer.revoked
+        consumer.before_schedule_shutdown
         consumer.shutting_down
         consumer.shutdown
 

data/lib/karafka/instrumentation/vendors/datadog/logger_listener.rb

@@ -94,8 +94,8 @@ module Karafka
               error "Consumer consuming error: #{error}"
             when 'consumer.revoked.error'
               error "Consumer on revoked failed due to an error: #{error}"
-            when 'consumer.before_enqueue.error'
-              error "Consumer before enqueue failed due to an error: #{error}"
+            when 'consumer.before_schedule.error'
+              error "Consumer before schedule failed due to an error: #{error}"
             when 'consumer.before_consume.error'
               error "Consumer before consume failed due to an error: #{error}"
             when 'consumer.after_consume.error'