Fix compressed offset bug

lib/kafka/compressor.rb

@@ -28,14 +28,15 @@ def initialize(codec_name: nil, threshold: 1, instrumenter:)
 
  # @param message_set [Protocol::MessageSet]
  # @return [Protocol::MessageSet]
- def compress(message_set)
+ def compress(message_set, offset: -1)
  return message_set if @codec.nil? || message_set.size < @threshold
 
  compressed_data = compress_data(message_set)
 
  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

@@ -58,10 +58,14 @@ def decompress
  message_set_decoder = Decoder.from_string(data)
  message_set = MessageSet.decode(message_set_decoder)
 
+ max_relative_offset = message_set.messages.last.offset
+ return message_set if max_relative_offset == offset
+
  # The contained messages need to have their offset corrected.
+ base_offset = offset - max_relative_offset
  messages = message_set.messages.each_with_index.map do |message, i|
  Message.new(
- offset: offset + i,
+ offset: message.offset + base_offset,
  value: message.value,
  key: message.key,
  create_time: message.create_time,

spec/compressor_spec.rb

@@ -2,29 +2,6 @@
  describe ".compress" do
  let(:instrumenter) { Kafka::Instrumenter.new(client_id: "test") }
 
- 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")
- message2 = Kafka::Protocol::Message.new(value: "hello2")
-
- message_set = Kafka::Protocol::MessageSet.new(messages: [message1, message2])
- 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
- messages = decoded_message_set.messages
-
- expect(messages.map(&:value)).to eq ["hello1", "hello2"]
-
- # 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]
- end
-
  it "only compresses the messages if there are at least the configured threshold" do
  compressor = Kafka::Compressor.new(codec_name: :snappy, threshold: 3, instrumenter: instrumenter)
 

spec/protocol/message_set_spec.rb

@@ -59,4 +59,59 @@
  Kafka::Protocol::MessageSet.decode(decoder)
  }.to raise_exception(Kafka::MessageTooLargeToRead)
  end
+
+ describe '.decode' do
+ let(:instrumenter) { Kafka::Instrumenter.new(client_id: "test") }
+ let(:compressor) { Kafka::Compressor.new(codec_name: :snappy, threshold: 1, instrumenter: instrumenter) }
+
+ def encode(messages: [], wrapper_message_offset: -1)
+ message_set = Kafka::Protocol::MessageSet.new(messages: messages)
+ compressed_message_set = compressor.compress(message_set, offset: wrapper_message_offset)
+ Kafka::Protocol::Encoder.encode_with(compressed_message_set)
+ end
+
+ def decode(data)
+ decoder = Kafka::Protocol::Decoder.from_string(data)
+ Kafka::Protocol::MessageSet
+ .decode(decoder)
+ .messages
+ end
+
+ 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: 2)
+
+ data = encode(messages: [message1, message2, message3], wrapper_message_offset: 1000)
+ messages = decode(data)
+
+ expect(messages.map(&:offset)).to eq [998, 999, 1000]
+ end
+
+ it "sets offsets correctly for compressed messages with relative offsets on a compacted topic" do
+ message1 = Kafka::Protocol::Message.new(value: "hello1", offset: 0)
+ message2 = Kafka::Protocol::Message.new(value: "hello2", offset: 2)
+ message3 = Kafka::Protocol::Message.new(value: "hello3", offset: 3)
+
+ data = encode(messages: [message1, message2, message3], wrapper_message_offset: 1000)
+ messages = decode(data)
+
+ expect(messages.map(&:offset)).to eq [997, 999, 1000]
+ end
+
+ it "keeps the predefined offsets for messages delivered in 0.9 format" do
+ compressor = Kafka::Compressor.new(codec_name: :snappy, threshold: 1, instrumenter: instrumenter)
+
+ message1 = Kafka::Protocol::Message.new(value: "hello1", offset: 997)
+ message2 = Kafka::Protocol::Message.new(value: "hello2", offset: 999)
+ message3 = Kafka::Protocol::Message.new(value: "hello3", offset: 1000)
+
+ data = encode(messages: [message1, message2, message3], wrapper_message_offset: 1000)
+ messages = decode(data)
+
+ expect(messages.map(&:offset)).to eq [997, 999, 1000]
+ end
+ end
 end