Merge pull request karafka#196 from gvisokinskas/feature/handle_in_callback

Pass the delivery handle to the callback

Author: leonmaia

Diff for: lib/rdkafka/callbacks.rb

@@ -97,7 +97,7 @@ def self.call(_, message_ptr, opaque_ptr)
           delivery_handle[:pending] = false
           # Call delivery callback on opaque
           if opaque = Rdkafka::Config.opaques[opaque_ptr.to_i]
-            opaque.call_delivery_callback(Rdkafka::Producer::DeliveryReport.new(message[:partition], message[:offset], message[:err]))
+            opaque.call_delivery_callback(Rdkafka::Producer::DeliveryReport.new(message[:partition], message[:offset], message[:err]), delivery_handle)
           end
         end
       end

Diff for: lib/rdkafka/config.rb

@@ -278,8 +278,8 @@ class Opaque
     attr_accessor :producer
     attr_accessor :consumer_rebalance_listener
 
-    def call_delivery_callback(delivery_handle)
-      producer.call_delivery_callback(delivery_handle) if producer
+    def call_delivery_callback(delivery_report, delivery_handle)
+      producer.call_delivery_callback(delivery_report, delivery_handle) if producer
     end
 
     def call_on_partitions_assigned(consumer, list)

Diff for: lib/rdkafka/producer.rb

@@ -9,6 +9,12 @@ class Producer
     # @return [Proc, nil]
     attr_reader :delivery_callback
 
+    # @private
+    # Returns the number of arguments accepted by the callback, by default this is nil.
+    #
+    # @return [Integer, nil]
+    attr_reader :delivery_callback_arity
+
     # @private
     def initialize(client, partitioner_name)
       @client = client
@@ -19,14 +25,15 @@ def initialize(client, partitioner_name)
     end
 
     # Set a callback that will be called every time a message is successfully produced.
-    # The callback is called with a {DeliveryReport}
+    # The callback is called with a {DeliveryReport} and {DeliveryHandle}
     #
     # @param callback [Proc, #call] The callback
     #
     # @return [nil]
     def delivery_callback=(callback)
       raise TypeError.new("Callback has to be callable") unless callback.respond_to?(:call)
       @delivery_callback = callback
+      @delivery_callback_arity = arity(callback)
     end
 
     # Close this producer and wait for the internal poll queue to empty.
@@ -151,8 +158,17 @@ def produce(topic:, payload: nil, key: nil, partition: nil, partition_key: nil,
     end
 
     # @private
-    def call_delivery_callback(delivery_handle)
-      @delivery_callback.call(delivery_handle) if @delivery_callback
+    def call_delivery_callback(delivery_report, delivery_handle)
+      return unless @delivery_callback
+
+      args = [delivery_report, delivery_handle].take(@delivery_callback_arity)
+      @delivery_callback.call(*args)
+    end
+
+    def arity(callback)
+      return callback.arity if callback.respond_to?(:arity)
+
+      callback.method(:call).arity
     end
 
     def closed_producer_check(method)

Diff for: spec/rdkafka/producer_spec.rb

@@ -49,6 +49,27 @@
         # Callback should have been called
         expect(@callback_called).to be true
       end
+
+      it "should provide handle" do
+        @callback_handle = nil
+
+        producer.delivery_callback = lambda { |_, handle| @callback_handle = handle }
+
+        # Produce a message
+        handle = producer.produce(
+          topic:   "produce_test_topic",
+          payload: "payload",
+          key:     "key"
+        )
+
+        # Wait for it to be delivered
+        handle.wait(max_wait_timeout: 15)
+
+        # Join the producer thread.
+        producer.close
+
+        expect(handle).to be @callback_handle
+      end
     end
 
     context "with a callable object" do
@@ -93,6 +114,36 @@ def call(report)
         expect(called_report.first.partition).to eq 1
         expect(called_report.first.offset).to be >= 0
       end
+
+      it "should provide handle" do
+        callback_handles = []
+        callback = Class.new do
+          def initialize(callback_handles)
+            @callback_handles = callback_handles
+          end
+
+          def call(_, handle)
+            @callback_handles << handle
+          end
+        end
+        producer.delivery_callback = callback.new(callback_handles)
+
+        # Produce a message
+        handle = producer.produce(
+          topic:   "produce_test_topic",
+          payload: "payload",
+          key:     "key"
+        )
+
+        # Wait for it to be delivered
+        handle.wait(max_wait_timeout: 15)
+
+        # Join the producer thread.
+        producer.close
+
+        # Callback should have been called
+        expect(handle).to be callback_handles.first
+      end
     end
 
     it "should not accept a callback that's not callable" do