[Chore] Refresh the topics when reading messages

Author: lmduc

lib/kafka/client.rb

@@ -336,6 +336,9 @@ def async_producer(delivery_interval: 0, delivery_threshold: 0, max_queue_size:
     # @param fetcher_max_queue_size [Integer] max number of items in the fetch queue that
     #   are stored for further processing. Note, that each item in the queue represents a
     #   response from a single broker.
+    # @param refresh_topic_interval [Integer] interval of refreshing the topic list.
+    #   If it is 0, the topic list won't be refreshed (default)
+    #   If it is n (n > 0), the topic list will be refreshed every n seconds
     # @return [Consumer]
     def consumer(
         group_id:,
@@ -345,7 +348,8 @@ def consumer(
         offset_commit_threshold: 0,
         heartbeat_interval: 10,
         offset_retention_time: nil,
-        fetcher_max_queue_size: 100
+        fetcher_max_queue_size: 100,
+        refresh_topic_interval: 0
     )
       cluster = initialize_cluster
 
@@ -399,6 +403,7 @@ def consumer(
         fetcher: fetcher,
         session_timeout: session_timeout,
         heartbeat: heartbeat,
+        refresh_topic_interval: refresh_topic_interval
       )
     end
 

lib/kafka/consumer.rb

@@ -44,7 +44,7 @@ module Kafka
   #
   class Consumer
 
-    def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manager:, session_timeout:, heartbeat:)
+    def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manager:, session_timeout:, heartbeat:, refresh_topic_interval: 0)
       @cluster = cluster
       @logger = TaggedLogger.new(logger)
       @instrumenter = instrumenter
@@ -53,6 +53,7 @@ def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manage
       @session_timeout = session_timeout
       @fetcher = fetcher
       @heartbeat = heartbeat
+      @refresh_topic_interval = refresh_topic_interval
 
       @pauses = Hash.new {|h, k|
         h[k] = Hash.new {|h2, k2|
@@ -74,9 +75,11 @@ def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manage
       #   some already consumed messages
       @current_offsets = Hash.new { |h, k| h[k] = {} }
 
-      # Hash storing topics that are already being subscribed
-      # When subcribing to a new topic, if it's already being subscribed before, skip it
-      @subscribed_topics = Set.new
+      # Map storing subscribed topics with their configuration
+      @subscribed_topics = Concurrent::Map.new
+
+      # Set storing topics that matched topics in @subscribed_topics
+      @matched_topics = Set.new
 
       # Whether join_group must be executed again because new topics are added
       @join_group_for_new_topics = false
@@ -100,29 +103,16 @@ def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manage
     #   checkpoint.
     # @param max_bytes_per_partition [Integer] the maximum amount of data fetched
     #   from a single partition at a time.
-    # @param refresh_topic_interval [Integer] interval of refreshing the topic list.
-    #   Only work when topic_or_regex is a regex
-    #   If it is 0, the topic list won't be refreshed (default)
-    #   If it is n (n > 0), the topic list will be refreshed every n seconds
     # @return [nil]
-    def subscribe(topic_or_regex, default_offset: nil, start_from_beginning: true, max_bytes_per_partition: 1048576, refresh_topic_interval: 0)
+    def subscribe(topic_or_regex, default_offset: nil, start_from_beginning: true, max_bytes_per_partition: 1048576)
       default_offset ||= start_from_beginning ? :earliest : :latest
 
-      unless topic_or_regex.is_a?(Regexp)
-        subscribe_to_topic(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
-        return
-      end
-
-      subscribe_to_regex(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
-      if refresh_topic_interval > 0
-        @thread ||= Thread.new do
-          while true
-            subscribe_to_regex(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition) if running?
-            sleep refresh_topic_interval
-          end
-        end
-        @thread.abort_on_exception = true
-      end
+      @subscribed_topics[topic_or_regex] = {
+        default_offset: default_offset,
+        start_from_beginning: start_from_beginning,
+        max_bytes_per_partition: max_bytes_per_partition
+      }
+      scan_for_subscribing
 
       nil
     end
@@ -228,6 +218,7 @@ def each_message(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automatica
       )
 
       consumer_loop do
+        refresh_topic_list_if_enabled
         batches = fetch_batches
 
         batches.each do |batch|
@@ -316,6 +307,7 @@ def each_batch(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automaticall
       )
 
       consumer_loop do
+        refresh_topic_list_if_enabled
         batches = fetch_batches
 
         batches.each do |batch|
@@ -473,6 +465,7 @@ def make_final_offsets_commit!(attempts = 3)
 
     def join_group
       @join_group_for_new_topics = false
+
       old_generation_id = @group.generation_id
 
       @group.join
@@ -534,6 +527,14 @@ def resume_paused_partitions!
       end
     end
 
+    def refresh_topic_list_if_enabled
+      return if @refresh_topic_interval <= 0
+      return if @refreshed_at && @refreshed_at + @refresh_topic_interval > Time.now
+
+      scan_for_subscribing
+      @refreshed_at = Time.now
+    end
+
     def fetch_batches
       # Return early if the consumer has been stopped.
       return [] if shutting_down?
@@ -592,20 +593,35 @@ def clear_current_offsets(excluding: {})
       end
     end
 
+    def scan_for_subscribing
+      @subscribed_topics.keys.each do |topic_or_regex|
+        default_offset = @subscribed_topics[topic_or_regex][:default_offset]
+        start_from_beginning = @subscribed_topics[topic_or_regex][:start_from_beginning]
+        max_bytes_per_partition = @subscribed_topics[topic_or_regex][:max_bytes_per_partition]
+
+        if topic_or_regex.is_a?(Regexp)
+          subscribe_to_regex(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
+        else
+          subscribe_to_topic(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
+        end
+      end
+    end
+
     def subscribe_to_regex(topic_regex, default_offset, start_from_beginning, max_bytes_per_partition)
       cluster_topics.select { |topic| topic =~ topic_regex }.each do |topic|
         subscribe_to_topic(topic, default_offset, start_from_beginning, max_bytes_per_partition)
       end
     end
 
     def subscribe_to_topic(topic, default_offset, start_from_beginning, max_bytes_per_partition)
-      return if @subscribed_topics.include?(topic)
-      @subscribed_topics.add(topic)
+      return if @matched_topics.include?(topic)
+      @matched_topics.add(topic)
       @join_group_for_new_topics = true
 
       @group.subscribe(topic)
       @offset_manager.set_default_offset(topic, default_offset)
       @fetcher.subscribe(topic, max_bytes_per_partition: max_bytes_per_partition)
+      @cluster.mark_as_stale!
     end
 
     def cluster_topics

spec/consumer_spec.rb

@@ -137,6 +137,7 @@
     allow(cluster).to receive(:add_target_topics)
     allow(cluster).to receive(:disconnect)
     allow(cluster).to receive(:refresh_metadata_if_necessary!)
+    allow(cluster).to receive(:mark_as_stale!)
 
     allow(offset_manager).to receive(:commit_offsets)
     allow(offset_manager).to receive(:commit_offsets_if_necessary)

spec/functional/consumer_group_spec.rb

@@ -131,8 +131,8 @@
     received_messages = []
 
     kafka = Kafka.new(kafka_brokers, client_id: "test", logger: logger)
-    consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time)
-    consumer.subscribe(/#{topic_a}|#{topic_b}/, refresh_topic_interval: 2)
+    consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time, refresh_topic_interval: 1)
+    consumer.subscribe(/#{topic_a}|#{topic_b}/)
 
     thread = Thread.new do
       consumer.each_message do |message|
@@ -145,7 +145,7 @@
     end
     thread.abort_on_exception = true
 
-    sleep 2
+    sleep 1
     messages_b.each { |i| producer.produce(i.to_s, topic: topic_b) }
     producer.deliver_messages