Add offsets_for_times method on consumer (#229) (#265)

* Add offsets_for_times method on consumer (#229)

* Add raw_timestamp attr_reader to Message class

* Add offsets_for_times method to Consumer class

* Make sure offsets_for_times call is blocking

* Fix @param typo

* make timeout configurable for offsets_for_times

* Remove publication of raw_timestamp

* Convert offset to ms timestamp if Time

* Fix typo

* align inner format

---------

Co-authored-by: Tim Flapper <[email protected]>
Co-authored-by: Tomasz Pajor <[email protected]>

Author: 3 people

lib/rdkafka/bindings.rb

@@ -188,6 +188,7 @@ class TopicPartitionList < FFI::Struct
     attach_function :rd_kafka_pause_partitions, [:pointer, :pointer], :int, blocking: true
     attach_function :rd_kafka_resume_partitions, [:pointer, :pointer], :int, blocking: true
     attach_function :rd_kafka_seek, [:pointer, :int32, :int64, :int], :int, blocking: true
+    attach_function :rd_kafka_offsets_for_times, [:pointer, :pointer, :int], :int, blocking: true
 
     # Headers
     attach_function :rd_kafka_header_get_all, [:pointer, :size_t, :pointer, :pointer, SizePtr], :int

lib/rdkafka/consumer.rb

@@ -406,6 +406,39 @@ def seek(message)
       end
     end
 
+    # Lookup offset for the given partitions by timestamp.
+    #
+    # @param list [TopicPartitionList] The TopicPartitionList with timestamps instead of offsets
+    #
+    # @raise [RdKafkaError] When the OffsetForTimes lookup fails
+    #
+    # @return [TopicPartitionList]
+    def offsets_for_times(list, timeout_ms = 1000)
+      closed_consumer_check(__method__)
+
+      if !list.is_a?(TopicPartitionList)
+        raise TypeError.new("list has to be a TopicPartitionList")
+      end
+
+      tpl = list.to_native_tpl
+
+      response = @native_kafka.with_inner do |inner|
+        Rdkafka::Bindings.rd_kafka_offsets_for_times(
+          inner,
+          tpl,
+          timeout_ms # timeout
+        )
+      end
+
+      if response != 0
+        raise Rdkafka::RdkafkaError.new(response)
+      end
+
+      TopicPartitionList.from_native_tpl(tpl)
+    ensure
+      Rdkafka::Bindings.rd_kafka_topic_partition_list_destroy(tpl) if tpl
+    end
+
     # Manually commit the current offsets of this consumer.
     #
     # To use this set `enable.auto.commit`to `false` to disable automatic triggering

lib/rdkafka/consumer/topic_partition_list.rb

@@ -142,11 +142,13 @@ def to_native_tpl
               )
 
               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,
-                  p.offset
+                  offset
                 )
               end
             end

spec/rdkafka/consumer/topic_partition_list_spec.rb

@@ -219,5 +219,24 @@
 
       expect(list).to eq other
     end
+
+    it "should create a native list with timetamp offsets if offsets are Time" do
+      list = Rdkafka::Consumer::TopicPartitionList.new.tap do |list|
+        list.add_topic_and_partitions_with_offsets("topic", 0 => Time.at(1505069646, 250_000))
+      end
+
+      tpl = list.to_native_tpl
+
+      compare_list = Rdkafka::Consumer::TopicPartitionList.new.tap do |list|
+        list.add_topic_and_partitions_with_offsets(
+          "topic",
+          0 => (Time.at(1505069646, 250_000).to_f * 1000).floor
+        )
+      end
+
+      native_list = Rdkafka::Consumer::TopicPartitionList.from_native_tpl(tpl)
+
+      expect(native_list).to eq compare_list
+    end
   end
 end

spec/rdkafka/consumer_spec.rb

@@ -953,6 +953,69 @@ def new_message
     end
   end
 
+  describe "#offsets_for_times" do
+    it "should raise when not TopicPartitionList" do
+      expect { consumer.offsets_for_times([]) }.to raise_error(TypeError)
+    end
+
+    it "should raise an error when offsets_for_times fails" do
+      tpl = Rdkafka::Consumer::TopicPartitionList.new
+
+      expect(Rdkafka::Bindings).to receive(:rd_kafka_offsets_for_times).and_return(7)
+
+      expect { consumer.offsets_for_times(tpl) }.to raise_error(Rdkafka::RdkafkaError)
+    end
+
+    context "when subscribed" do
+      let(:timeout) { 1000 }
+
+      before do
+        consumer.subscribe("consume_test_topic")
+
+        # 1. partitions are assigned
+        wait_for_assignment(consumer)
+        expect(consumer.assignment).not_to be_empty
+
+        # 2. eat unrelated messages
+        while(consumer.poll(timeout)) do; end
+      end
+
+      after { consumer.unsubscribe }
+
+      def send_one_message(val)
+        producer.produce(
+          topic:     "consume_test_topic",
+          payload:   "payload #{val}",
+          key:       "key 0",
+          partition: 0
+        ).wait
+      end
+
+      it "returns a TopicParticionList with updated offsets" do
+        send_one_message("a")
+        send_one_message("b")
+        send_one_message("c")
+
+        consumer.poll(timeout)
+        message = consumer.poll(timeout)
+        consumer.poll(timeout)
+
+        tpl = Rdkafka::Consumer::TopicPartitionList.new.tap do |list|
+          list.add_topic_and_partitions_with_offsets(
+            "consume_test_topic",
+            [
+              [0, message.timestamp]
+            ]
+          )
+        end
+
+        tpl_response = consumer.offsets_for_times(tpl)
+
+        expect(tpl_response.to_h["consume_test_topic"][0].offset).to eq message.offset
+      end
+    end
+  end
+
   describe "a rebalance listener" do
     let(:consumer) do
       config = rdkafka_consumer_config