rdkafka 0.28.0 → 0.29.0

data/lib/rdkafka/consumer/topic_partition_list.rb

@@ -118,7 +118,12 @@ module Rdkafka
             else
               elem[:offset]
             end
-            partition = Partition.new(elem[:partition], offset, elem[:err])
+            partition = Partition.new(
+              elem[:partition],
+              offset,
+              elem[:err],
+              elem[:metadata].null? ? nil : elem[:metadata].read_string(elem[:metadata_size])
+            )
             partitions.push(partition)
             data[elem[:topic]] = partitions
           end
@@ -137,33 +142,50 @@ module Rdkafka
       def to_native_tpl
         tpl = Rdkafka::Bindings.rd_kafka_topic_partition_list_new(count)
 
-        @data.each do |topic, partitions|
-          if partitions
-            partitions.each do |p|
-              Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
-                tpl,
-                topic,
-                p.partition
-              )
-
-              if p.offset
-                offset = p.offset.is_a?(Time) ? p.offset.to_f * 1_000 : p.offset
-
-                Rdkafka::Bindings.rd_kafka_topic_partition_list_set_offset(
+        begin
+          @data.each do |topic, partitions|
+            if partitions
+              partitions.each do |p|
+                ref = Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
                   tpl,
                   topic,
-                  p.partition,
-                  offset
+                  p.partition
                 )
+
+                if p.metadata
+                  part = Rdkafka::Bindings::TopicPartition.new(ref)
+                  str_ptr = FFI::MemoryPointer.from_string(p.metadata)
+                  # The metadata string is owned by librdkafka once handed over and released here:
+                  # https://github.com/confluentinc/librdkafka/blob/e03d3bb91ed92a38f38d9806b8d8deffe78a1de5/src/rdkafka_partition.c#L2682C18-L2682C18
+                  str_ptr.autorelease = false
+                  part[:metadata] = str_ptr
+                  part[:metadata_size] = p.metadata.bytesize
+                end
+
+                if p.offset
+                  offset = p.offset.is_a?(Time) ? p.offset.to_f * 1_000 : p.offset
+
+                  Rdkafka::Bindings.rd_kafka_topic_partition_list_set_offset(
+                    tpl,
+                    topic,
+                    p.partition,
+                    offset
+                  )
+                end
               end
+            else
+              Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
+                tpl,
+                topic,
+                Rdkafka::Bindings::RD_KAFKA_PARTITION_UA
+              )
             end
-          else
-            Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
-              tpl,
-              topic,
-              Rdkafka::Bindings::RD_KAFKA_PARTITION_UA
-            )
           end
+        rescue Exception
+          # The native list is caller-owned until it is handed to librdkafka. If population raises
+          # partway (e.g. an invalid partition), destroy it here so it is not leaked.
+          Rdkafka::Bindings.rd_kafka_topic_partition_list_destroy(tpl)
+          raise
         end
 
         tpl

data/lib/rdkafka/consumer.rb

@@ -14,14 +14,45 @@ module Rdkafka
     include Enumerable
     include Helpers::Time
     include Helpers::OAuth
+    include Helpers::Metadata
 
     # @private
     # @param native_kafka [NativeKafka] wrapper around the native Kafka consumer handle
     def initialize(native_kafka)
       @native_kafka = native_kafka
+      # Single-element holder shared with the GC finalizer so it can destroy the lazily created
+      # consumer queue without capturing `self` (capturing the consumer in its own finalizer would
+      # pin it and prevent it from ever being collected).
+      @consumer_queue_holder = []
+
+      # Makes sure the consumer is closed (consumer queue destroyed and native client destroyed)
+      # before it gets GCed by Ruby.
+      ObjectSpace.define_finalizer(self, self.class.finalizer(native_kafka, @consumer_queue_holder))
+    end
+
+    # Builds the GC finalizer for a consumer. It mirrors {#close}: close the consumer, destroy the
+    # consumer-queue reference, then destroy the native client. The default `NativeKafka#finalizer`
+    # went straight to `rd_kafka_destroy`, leaving the consumer-queue reference (from
+    # `rd_kafka_queue_get_consumer`, taken by `poll_batch`) dangling - which can make
+    # `rd_kafka_destroy` block inside the finalizer (process hang at GC/shutdown) or leak the handle.
+    #
+    # @private
+    # @param native_kafka [NativeKafka] the wrapped native client
+    # @param queue_holder [Array] single-element holder carrying the consumer queue pointer (or empty)
+    # @return [Proc] finalizer proc that must not reference the consumer instance
+    def self.finalizer(native_kafka, queue_holder)
+      proc do
+        next if native_kafka.closed?
+
+        native_kafka.synchronize do |inner|
+          Rdkafka::Bindings.rd_kafka_consumer_close(inner)
+
+          queue = queue_holder[0]
+          Rdkafka::Bindings.rd_kafka_queue_destroy(queue) if queue
+        end
 
-      # Makes sure, that native kafka gets closed before it gets GCed by Ruby
-      ObjectSpace.define_finalizer(self, native_kafka.finalizer)
+        native_kafka.close
+      end
     end
 
     # Starts the native Kafka polling thread and kicks off the init polling
@@ -187,6 +218,7 @@ module Rdkafka
         if @consumer_queue
           Rdkafka::Bindings.rd_kafka_queue_destroy(@consumer_queue)
           @consumer_queue = nil
+          @consumer_queue_holder[0] = nil
         end
       end
 
@@ -422,6 +454,8 @@ module Rdkafka
     #
     # @raise [RdkafkaError] When getting the positions fails.
     def position(list = nil)
+      closed_consumer_check(__method__)
+
       if list.nil?
         list = assignment
       elsif !list.is_a?(TopicPartitionList)
@@ -490,8 +524,7 @@ module Rdkafka
       out = {}
 
       topic_partition_list.to_h.each do |topic, partitions|
-        # Query high watermarks for this topic's partitions
-        # and compare to the offset in the list.
+        # Query high watermarks for this topic's partitions and compare to the offset in the list.
         topic_out = {}
         partitions.each do |p|
           next if p.offset.nil?
@@ -509,11 +542,12 @@ module Rdkafka
 
     # Returns the ClusterId as reported in broker metadata.
     #
+    # @param timeout_ms [Integer] timeout in milliseconds to wait for the cluster id
     # @return [String, nil]
-    def cluster_id
+    def cluster_id(timeout_ms = Defaults::CONSUMER_CLUSTER_ID_TIMEOUT_MS)
       closed_consumer_check(__method__)
       @native_kafka.with_inner do |inner|
-        Rdkafka::Bindings.rd_kafka_clusterid(inner)
+        read_and_free_native_string(inner, Rdkafka::Bindings.rd_kafka_clusterid(inner, timeout_ms))
       end
     end
 
@@ -525,7 +559,7 @@ module Rdkafka
     def member_id
       closed_consumer_check(__method__)
       @native_kafka.with_inner do |inner|
-        Rdkafka::Bindings.rd_kafka_memberid(inner)
+        read_and_free_native_string(inner, Rdkafka::Bindings.rd_kafka_memberid(inner))
       end
     end
 
@@ -534,16 +568,34 @@ module Rdkafka
     # When using this `enable.auto.offset.store` should be set to `false` in the config.
     #
     # @param message [Rdkafka::Consumer::Message] The message which offset will be stored
+    # @param metadata [String, nil] commit metadata string to store alongside the offset
     # @return [nil]
     # @raise [RdkafkaError] When storing the offset fails
-    def store_offset(message)
+    def store_offset(message, metadata = nil)
       closed_consumer_check(__method__)
 
       list = TopicPartitionList.new
-      list.add_topic_and_partitions_with_offsets(
-        message.topic,
-        message.partition => message.offset + 1
-      )
+
+      # For metadata aware commits we build the partition reference directly to save on
+      # objects allocations
+      if metadata
+        list.add_topic_and_partitions_with_offsets(
+          message.topic,
+          [
+            Consumer::Partition.new(
+              message.partition,
+              message.offset + 1,
+              0,
+              metadata
+            )
+          ]
+        )
+      else
+        list.add_topic_and_partitions_with_offsets(
+          message.topic,
+          message.partition => message.offset + 1
+        )
+      end
 
       tpl = list.to_native_tpl
 
@@ -572,8 +624,7 @@ module Rdkafka
     end
 
     # Seek to a particular message by providing the topic, partition and offset.
-    # The next poll on the topic/partition will return the
-    # message at the given offset.
+    # The next poll on the topic/partition will return the message at the given offset.
     #
     # @param topic [String] The topic in which to seek
     # @param partition [Integer] The partition number to seek
@@ -641,8 +692,7 @@ module Rdkafka
 
     # Manually commit the current offsets of this consumer.
     #
-    # To use this set `enable.auto.commit`to `false` to disable automatic triggering
-    # of commits.
+    # To use this set `enable.auto.commit`to `false` to disable automatic triggering of commits.
     #
     # If `enable.auto.offset.store` is set to `true` the offset of the last consumed
     # message for every partition is used. If set to `false` you can use {store_offset} to
@@ -841,8 +891,17 @@ module Rdkafka
             next
           end
 
-          results << Rdkafka::Consumer::Message.new(native_message)
-          Rdkafka::Bindings.rd_kafka_message_destroy(ptr)
+          begin
+            results << Rdkafka::Consumer::Message.new(native_message)
+          rescue Rdkafka::RdkafkaError => e
+            # A message that fails to build (e.g. a header read error) is surfaced inline as an
+            # error event rather than discarding the whole batch - including the messages already
+            # built - and raising, which silently lost them once their offsets had been stored.
+            results << e
+          ensure
+            Rdkafka::Bindings.rd_kafka_message_destroy(ptr)
+          end
+
           i += 1
         end
       ensure
@@ -907,8 +966,17 @@ module Rdkafka
             next
           end
 
-          results << Rdkafka::Consumer::Message.new(native_message)
-          Rdkafka::Bindings.rd_kafka_message_destroy(ptr)
+          begin
+            results << Rdkafka::Consumer::Message.new(native_message)
+          rescue Rdkafka::RdkafkaError => e
+            # A message that fails to build (e.g. a header read error) is surfaced inline as an
+            # error event rather than discarding the whole batch - including the messages already
+            # built - and raising, which silently lost them once their offsets had been stored.
+            results << e
+          ensure
+            Rdkafka::Bindings.rd_kafka_message_destroy(ptr)
+          end
+
           i += 1
         end
       ensure
@@ -994,12 +1062,34 @@ module Rdkafka
     def closed_consumer_check(method)
       raise Rdkafka::ClosedConsumerError.new(method) if closed?
     end
+    alias_method :closed_check, :closed_consumer_check
+
+    # Reads a librdkafka-allocated string and frees the underlying native buffer.
+    #
+    # librdkafka returns heap-allocated, caller-owned strings from functions like
+    # `rd_kafka_clusterid`/`rd_kafka_memberid`. The buffer must be released with
+    # `rd_kafka_mem_free`, otherwise it leaks on every call.
+    #
+    # @param inner [FFI::Pointer] the native client handle used to allocate the string
+    # @param ptr [FFI::Pointer] the native string pointer (may be null)
+    # @return [String, nil] the copied Ruby string, or nil when the pointer is null
+    def read_and_free_native_string(inner, ptr)
+      return nil if ptr.null?
+
+      ptr.read_string
+    ensure
+      Rdkafka::Bindings.rd_kafka_mem_free(inner, ptr) unless ptr.null?
+    end
 
     # Returns the consumer queue pointer, lazily initialized
     # @return [FFI::Pointer] consumer queue handle
     def consumer_queue
       @consumer_queue ||= @native_kafka.with_inner do |inner|
-        Rdkafka::Bindings.rd_kafka_queue_get_consumer(inner)
+        queue = Rdkafka::Bindings.rd_kafka_queue_get_consumer(inner)
+        # Share the pointer with the finalizer so it is destroyed even if the consumer is GC'd
+        # without an explicit close.
+        @consumer_queue_holder[0] = queue
+        queue
       end
     end
 

data/lib/rdkafka/defaults.rb

@@ -23,6 +23,10 @@ module Rdkafka
     # @see Consumer#committed
     CONSUMER_COMMITTED_TIMEOUT_MS = 2_000
 
+    # Default timeout for fetching the cluster id
+    # @see Consumer#cluster_id
+    CONSUMER_CLUSTER_ID_TIMEOUT_MS = 1_000
+
     # Default timeout for querying watermark offsets
     # @see Consumer#query_watermark_offsets
     CONSUMER_QUERY_WATERMARK_TIMEOUT_MS = 1_000
@@ -91,6 +95,16 @@ module Rdkafka
     # @see Metadata#initialize
     METADATA_RETRY_BACKOFF_BASE_MS = 100
 
+    # Maximum backoff time between metadata retries. Caps the exponential backoff so a long retry
+    # sequence against an unhealthy cluster cannot block the calling thread for minutes.
+    # @see Metadata#initialize
+    METADATA_RETRY_BACKOFF_MAX_MS = 1_000
+
+    # Soft wall-clock budget for the whole metadata retry loop; past it (and past
+    # METADATA_MIN_ATTEMPTS) the loop stops so a synchronous fetch cannot block the caller for long
+    # @see Metadata#initialize
+    METADATA_RETRY_BUDGET_MS = 5_000
+
     # Cache settings (in milliseconds)
 
     # Default time-to-live for cached partition counts
@@ -102,5 +116,10 @@ module Rdkafka
     # Maximum number of metadata fetch retry attempts
     # @see Metadata#initialize
     METADATA_MAX_RETRIES = 10
+
+    # Minimum metadata fetch attempts before the retry budget may end the loop, so a slow broker
+    # (whose requests each consume the full timeout) still gets a few tries
+    # @see Metadata#initialize
+    METADATA_MIN_ATTEMPTS = 3
   end
 end

data/lib/rdkafka/helpers/metadata.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Rdkafka
+  module Helpers
+    # Shared `#metadata` implementation for Admin, Consumer and Producer.
+    #
+    # `rd_kafka_metadata()` is handle-agnostic in librdkafka - it works on any `rd_kafka_t`
+    # (consumer, producer or admin) - so this Ruby-level implementation is identical across all
+    # three client types. Includers must provide a private `#closed_check(method)` that raises
+    # their own `Closed*Error`.
+    module Metadata
+      # Performs the metadata request using this client
+      #
+      # @param topic_name [String, nil] metadata about particular topic or all if nil
+      # @param timeout_ms [Integer] metadata request timeout
+      # @return [Rdkafka::Metadata] requested metadata
+      def metadata(topic_name = nil, timeout_ms = Defaults::METADATA_TIMEOUT_MS)
+        closed_check(__method__)
+
+        @native_kafka.with_inner do |inner|
+          # Must stay fully qualified: this module is itself named `Metadata`, so a bare
+          # `Metadata.new` here would resolve to `Rdkafka::Helpers::Metadata` (no `.new`) instead
+          # of this class.
+          Rdkafka::Metadata.new(inner, topic_name, timeout_ms)
+        end
+      end
+    end
+  end
+end

data/lib/rdkafka/metadata.rb

@@ -23,12 +23,57 @@ module Rdkafka
     # @param timeout_ms [Integer] timeout in milliseconds
     # @raise [RdkafkaError] when metadata fetch fails
     def initialize(native_client, topic_name = nil, timeout_ms = Defaults::METADATA_TIMEOUT_MS)
-      attempt ||= 0
-      attempt += 1
-
-      native_topic = if topic_name
-        Rdkafka::Bindings.rd_kafka_topic_new(native_client, topic_name, nil)
+      attempt = 0
+      deadline = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) +
+        Defaults::METADATA_RETRY_BUDGET_MS / 1_000.0
+
+      begin
+        attempt += 1
+        fetch_metadata(native_client, topic_name, timeout_ms)
+      rescue ::Rdkafka::RdkafkaError => e
+        raise unless RETRIED_ERRORS.include?(e.code)
+        raise if attempt > Defaults::METADATA_MAX_RETRIES
+
+        # Stop once the wall-clock retry budget is spent, but only after at least
+        # METADATA_MIN_ATTEMPTS tries so a slow broker (whose requests each consume the full
+        # timeout) still gets a few tries rather than being cut off after one or two.
+        raise if attempt >= Defaults::METADATA_MIN_ATTEMPTS &&
+          ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) >= deadline
+
+        # Exponential backoff between attempts, capped so a long retry sequence cannot block for
+        # minutes. The request timeout (`timeout_ms`) is intentionally left unchanged: it used to be
+        # overwritten with the backoff value, which shrank the first retries below the configured
+        # timeout (near-guaranteeing another timeout) and then inflated later ones to ~100s.
+        backoff_ms = [
+          (2**attempt) * Defaults::METADATA_RETRY_BACKOFF_BASE_MS,
+          Defaults::METADATA_RETRY_BACKOFF_MAX_MS
+        ].min
+
+        sleep(backoff_ms / 1_000.0)
+
+        retry
       end
+    end
+
+    private
+
+    # Performs a single metadata fetch attempt, freeing this attempt's native resources.
+    #
+    # Kept separate from {#initialize} so each retried attempt frees its own `native_topic` and
+    # metadata struct. `retry` re-enters `initialize`'s `begin` without running an `ensure` placed
+    # there, so doing the cleanup per attempt here is what prevents the per-retry native leak. The
+    # metadata struct is only allocated on success, so it is read and destroyed only after the
+    # result has been validated (avoids destroying a NULL/garbage pointer on a failed fetch).
+    #
+    # @param native_client [FFI::Pointer] pointer to the native Kafka client
+    # @param topic_name [String, nil] specific topic to fetch metadata for, or nil for all topics
+    # @param timeout_ms [Integer] timeout in milliseconds
+    # @raise [RdkafkaError] when the metadata fetch fails
+    def fetch_metadata(native_client, topic_name, timeout_ms)
+      native_topic = nil
+      metadata_ptr = nil
+
+      native_topic = Rdkafka::Bindings.rd_kafka_topic_new(native_client, topic_name, nil) if topic_name
 
       ptr = FFI::MemoryPointer.new(:pointer)
 
@@ -42,24 +87,16 @@ module Rdkafka
       # Error Handling
       raise Rdkafka::RdkafkaError.new(result) unless result.zero?
 
-      metadata_from_native(ptr.read_pointer)
-    rescue ::Rdkafka::RdkafkaError => e
-      raise unless RETRIED_ERRORS.include?(e.code)
-      raise if attempt > Defaults::METADATA_MAX_RETRIES
+      # rd_kafka_metadata only allocates the struct on success, so we read the pointer to destroy
+      # only after the result has been confirmed successful.
+      metadata_ptr = ptr.read_pointer
 
-      backoff_factor = 2**attempt
-      timeout_ms = backoff_factor * Defaults::METADATA_RETRY_BACKOFF_BASE_MS
-
-      sleep(timeout_ms / 1000.0)
-
-      retry
+      metadata_from_native(metadata_ptr)
     ensure
-      Rdkafka::Bindings.rd_kafka_topic_destroy(native_topic) if topic_name
-      Rdkafka::Bindings.rd_kafka_metadata_destroy(ptr.read_pointer)
+      Rdkafka::Bindings.rd_kafka_topic_destroy(native_topic) if native_topic
+      Rdkafka::Bindings.rd_kafka_metadata_destroy(metadata_ptr) if metadata_ptr && !metadata_ptr.null?
     end
 
-    private
-
     # Extracts metadata from native pointer
     # @param ptr [FFI::Pointer] pointer to native metadata
     def metadata_from_native(ptr)
@@ -134,6 +171,36 @@ module Rdkafka
         :replicas, :pointer,
         :in_sync_replica_brokers, :int,
         :isrs, :pointer
+
+      # The base `#to_h` skips FFI pointer members, which would drop the replica and in-sync
+      # replica assignments entirely. We dereference those pointers here so the partition hash
+      # exposes the broker ids backing the partition (needed e.g. to plan replication changes).
+      #
+      # @return [Hash{Symbol => Integer, Array<Integer>}] partition metadata:
+      #   * +:partition_id+ (Integer) - partition id
+      #   * +:leader+ (Integer) - broker id of the partition leader
+      #   * +:replica_count+ (Integer) - number of assigned replicas
+      #   * +:in_sync_replica_brokers+ (Integer) - number of in-sync replicas
+      #   * +:replicas+ (Array<Integer>) - broker ids of the assigned replicas
+      #   * +:isrs+ (Array<Integer>) - broker ids of the in-sync replicas
+      def to_h
+        super.merge(
+          replicas: read_broker_ids(self[:replicas], self[:replica_count]),
+          isrs: read_broker_ids(self[:isrs], self[:in_sync_replica_brokers])
+        )
+      end
+
+      private
+
+      # Reads `count` broker ids (int32) from a replicas/isrs pointer.
+      # @param pointer [FFI::Pointer] pointer to the broker ids array
+      # @param count [Integer] number of broker ids to read
+      # @return [Array<Integer>] broker ids (empty when there are none)
+      def read_broker_ids(pointer, count)
+        return [] if count.zero? || pointer.null?
+
+        pointer.read_array_of_int32(count)
+      end
     end
   end
 end

data/lib/rdkafka/native_kafka.rb

@@ -55,8 +55,7 @@ module Rdkafka
         Rdkafka::Bindings.rd_kafka_poll(@inner, 0)
 
         if @run_polling_thread
-          # Start thread to poll client for delivery callbacks,
-          # not used in consumer.
+          # Start thread to poll client for delivery callbacks, not used in consumer.
           @polling_thread = Thread.new do
             loop do
               @poll_mutex.synchronize do
@@ -197,8 +196,7 @@ module Rdkafka
           # Indicate to polling thread that we're closing
           @polling_thread[:closing] = true
 
-          # Wait for the polling thread to finish up,
-          # this can be aborted in practice if this
+          # Wait for the polling thread to finish up, this can be aborted in practice if this
           # code runs from a finalizer.
           @polling_thread.join
         end

data/lib/rdkafka/producer/delivery_handle.rb

@@ -8,14 +8,14 @@ module Rdkafka
       layout :pending, :bool,
         :response, :int,
         :partition, :int,
-        :offset, :int64,
-        :topic_name, :pointer
+        :offset, :int64
 
       # @return [Object, nil] label set during message production or nil by default
       attr_accessor :label
 
       # @return [String] topic where we are trying to send the message
-      # We use this instead of reading from `topic_name` pointer to save on memory allocations
+      # Set in `#produce`, where the topic is known upfront. Keeping it as a Ruby attribute
+      # spares a per-message native string copy in the delivery callback.
       attr_accessor :topic
 
       # @return [String] the name of the operation (e.g. "delivery")

data/lib/rdkafka/producer/partitions_count_cache.rb

@@ -37,12 +37,11 @@ module Rdkafka
     #    contention in multi-threaded environments while ensuring data consistency.
     #
     # 6. Topic recreation handling
-    #    If a topic is deleted and recreated with fewer partitions, the cache will continue to
-    #    report the higher count until either the TTL expires or the process is restarted. This
-    #    design choice simplifies the implementation while relying on librdkafka's error handling
-    #    for edge cases. In production environments, topic recreation with different partition
-    #    counts is typically accompanied by application restarts to handle structural changes.
-    #    This also aligns with the previous cache implementation.
+    #    If a topic is deleted and recreated with fewer partitions, the cache keeps reporting the
+    #    higher count only until the entry's TTL expires. The first refresh after expiry performs
+    #    an authoritative metadata read and adopts the lower count. Within the TTL window a lower
+    #    value is still ignored, so a transient or racy lower read cannot clobber a correct higher
+    #    count.
     class PartitionsCountCache
       include Helpers::Time
 
@@ -91,28 +90,14 @@ module Rdkafka
         current_info = @counts[topic]
 
         if current_info.nil? || expired?(current_info[0])
+          # The cached entry is missing or expired, so the block performs an authoritative metadata
+          # read. We hand it to `set`, which adopts a higher count always and a lower count once the
+          # entry has expired (e.g. the topic was recreated with fewer partitions). We then return
+          # whatever `set` settled on so a concurrent refresh that wrote a higher value still wins.
           new_count = yield
+          set(topic, new_count)
 
-          if current_info.nil?
-            # No existing data, create a new entry with mutex
-            set(topic, new_count)
-
-            return new_count
-          else
-            current_count = current_info[1]
-
-            if new_count > current_count
-              # Higher value needs mutex to update both timestamp and count
-              set(topic, new_count)
-
-              return new_count
-            else
-              # Same or lower value, just update timestamp without mutex
-              refresh_timestamp(topic)
-
-              return current_count
-            end
-          end
+          return @counts[topic][1]
         end
 
         current_info[1]
@@ -133,8 +118,11 @@ module Rdkafka
         # First check outside mutex to avoid unnecessary locking
         current_info = @counts[topic]
 
-        # For lower values, we don't update count but might need to refresh timestamp
-        if current_info && new_count < current_info[1]
+        # Within the TTL window a lower value is treated as a stale/racy read and ignored, since
+        # partition counts only grow during normal operation. Once the entry has expired a lower
+        # value is an authoritative refresh (e.g. the topic was recreated with fewer partitions),
+        # so we fall through and adopt it below.
+        if current_info && new_count < current_info[1] && !expired?(current_info[0])
           refresh_timestamp(topic)
 
           return
@@ -148,17 +136,15 @@ module Rdkafka
           if current_info.nil?
             # Create new entry
             @counts[topic] = [monotonic_now_ms, new_count]
+          elsif new_count > current_info[1] || expired?(current_info[0])
+            # A higher count always wins; a lower count is accepted only when the existing entry
+            # has expired, so a concurrent fresh higher value (which reset the timestamp) is never
+            # clobbered by a stale lower one.
+            current_info[0] = monotonic_now_ms
+            current_info[1] = new_count
           else
-            current_count = current_info[1]
-
-            if new_count > current_count
-              # Update to higher count value
-              current_info[0] = monotonic_now_ms
-              current_info[1] = new_count
-            else
-              # Same or lower count, update timestamp only
-              current_info[0] = monotonic_now_ms
-            end
+            # Same or lower count within the TTL window: refresh the timestamp only
+            current_info[0] = monotonic_now_ms
           end
         end
       end