Resurrect the fix provided by @ale7714 in #495

Author: klippx

lib/kafka/compressor.rb

@@ -32,10 +32,12 @@ def compress(message_set)
       return message_set if @codec.nil? || message_set.size < @threshold
 
       compressed_data = compress_data(message_set)
+      offset = message_set.messages.map(&:offset).max
 
       wrapper_message = Protocol::Message.new(
         value: compressed_data,
         codec_id: @codec.codec_id,
+        offset: offset
       )
 
       Protocol::MessageSet.new(messages: [wrapper_message])

lib/kafka/protocol/message.rb

@@ -56,7 +56,22 @@ def decompress
         # For some weird reason we need to cut out the first 20 bytes.
         data = codec.decompress(value)
         message_set_decoder = Decoder.from_string(data)
-        MessageSet.decode(message_set_decoder)
+        message_set = MessageSet.decode(message_set_decoder)
+
+        base_offset = offset - message_set.size + 1
+
+        # The contained messages need to have their offset corrected.
+        messages = message_set.messages.each_with_index.map do |message, i|
+          Message.new(
+            offset: base_offset + i,
+            value: message.value,
+            key: message.key,
+            create_time: message.create_time,
+            codec_id: message.codec_id
+          )
+        end
+
+        MessageSet.new(messages: messages)
       end
 
       def self.decode(decoder)

spec/compressor_spec.rb

@@ -5,8 +5,8 @@
     it "encodes and decodes compressed messages" do
       compressor = Kafka::Compressor.new(codec_name: :snappy, threshold: 1, instrumenter: instrumenter)
 
-      message1 = Kafka::Protocol::Message.new(value: "hello1", offset: -1)
-      message2 = Kafka::Protocol::Message.new(value: "hello2", offset: 0)
+      message1 = Kafka::Protocol::Message.new(value: "hello1")
+      message2 = Kafka::Protocol::Message.new(value: "hello2")
 
       message_set = Kafka::Protocol::MessageSet.new(messages: [message1, message2])
       compressed_message_set = compressor.compress(message_set)
@@ -22,10 +22,28 @@
       # When decoding a compressed message, the offsets are calculated relative to that
       # of the container message. The broker will set the offset in normal operation,
       # but at the client-side we set it to -1.
-      expect(messages.map(&:offset)).to eq [-1, 0]
+      expect(messages.map(&:offset)).to eq [-2, -1]
     end
 
-    it "encodes and decodes compressed messages" do
+    it "sets offsets correctly for compressed messages with relative offsets" do
+      compressor = Kafka::Compressor.new(codec_name: :snappy, threshold: 1, instrumenter: instrumenter)
+
+      message1 = Kafka::Protocol::Message.new(value: "hello1", offset: 0)
+      message2 = Kafka::Protocol::Message.new(value: "hello2", offset: 1)
+      message3 = Kafka::Protocol::Message.new(value: "hello3", offset: 772043)
+
+      message_set = Kafka::Protocol::MessageSet.new(messages: [message1, message2, message3])
+      compressed_message_set = compressor.compress(message_set)
+      data = Kafka::Protocol::Encoder.encode_with(compressed_message_set)
+
+      decoder = Kafka::Protocol::Decoder.from_string(data)
+      decoded_message_set = Kafka::Protocol::MessageSet.decode(decoder)
+      messages = decoded_message_set.messages
+
+      expect(messages.map(&:offset)).to eq [772041, 772042, 772043]
+    end
+
+    it "keeps the offsets for compressed messages with provided offsets" do
       compressor = Kafka::Compressor.new(codec_name: :snappy, threshold: 1, instrumenter: instrumenter)
 
       message1 = Kafka::Protocol::Message.new(
@@ -45,9 +63,9 @@
       compressed_message_set = compressor.compress(message_set)
 
       data = Kafka::Protocol::Encoder.encode_with(compressed_message_set)
+
       decoder = Kafka::Protocol::Decoder.from_string(data)
-      decoded_message = Kafka::Protocol::Message.decode(decoder)
-      decoded_message_set = decoded_message.decompress
+      decoded_message_set = Kafka::Protocol::MessageSet.decode(decoder)
       messages = decoded_message_set.messages
 
       expect(messages.map(&:offset)).to eq [772033, 772034, 772035, 772036, 772037, 772038, 772039, 772040, 772041, 772042, 772043]