Automatic offset tracking for stream queues

spec/stream_queue_spec.cr

@@ -1,6 +1,38 @@
 require "./spec_helper"
 require "./../src/lavinmq/amqp/queue"
 
+module StreamQueueSpecHelpers
+  def self.publish(s, queue_name, nr_of_messages)
+    args = {"x-queue-type": "stream"}
+    with_channel(s) do |ch|
+      q = ch.queue(queue_name, args: AMQP::Client::Arguments.new(args))
+      nr_of_messages.times { |i| q.publish "m#{i}" }
+    end
+  end
+
+  def self.consume_one(s, queue_name, c_tag, c_args = AMQP::Client::Arguments.new)
+    args = {"x-queue-type": "stream"}
+    with_channel(s) do |ch|
+      ch.prefetch 1
+      q = ch.queue(queue_name, args: AMQP::Client::Arguments.new(args))
+      msgs = Channel(AMQP::Client::DeliverMessage).new
+      q.subscribe(no_ack: false, tag: c_tag, args: c_args) do |msg|
+        msgs.send msg
+        msg.ack
+      end
+      msgs.receive
+    end
+  end
+
+  def self.offset_from_headers(headers)
+    if headers
+      headers["x-stream-offset"].as(Int64)
+    else
+      fail("No headers found")
+    end
+  end
+end
+
 describe LavinMQ::AMQP::StreamQueue do
   stream_queue_args = LavinMQ::AMQP::Table.new({"x-queue-type": "stream"})
 
@@ -203,4 +235,274 @@ describe LavinMQ::AMQP::StreamQueue do
       end
     end
   end
+
+  describe "Automatic consumer offset tracking" do
+    it "resumes from last offset on reconnect" do
+      queue_name = Random::Secure.hex
+      consumer_tag = Random::Secure.hex
+      offset = 3
+
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, offset + 1)
+        offset.times { StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag) }
+        sleep 0.1
+
+        # consume again, should start from last offset automatically
+        msg = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag)
+        StreamQueueSpecHelpers.offset_from_headers(msg.properties.headers).should eq offset + 1
+      end
+    end
+
+    it "reads offsets from file on init" do
+      queue_name = Random::Secure.hex
+      offsets = [84_i64, 24_i64, 1_i64, 100_i64, 42_i64]
+      tag_prefix = "ctag-"
+
+      with_amqp_server do |s|
+        vhost = s.vhosts["/"]
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+
+        data_dir = File.join(vhost.data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each_with_index do |offset, i|
+          msg_store.store_consumer_offset(tag_prefix + i.to_s, offset)
+        end
+        msg_store.close
+        sleep 0.1
+
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each_with_index do |offset, i|
+          msg_store.last_offset_by_consumer_tag(tag_prefix + i.to_s).should eq offset
+        end
+        msg_store.close
+      end
+    end
+
+    it "appends consumer tag file" do
+      queue_name = Random::Secure.hex
+      offsets = [84_i64, 24_i64, 1_i64, 100_i64, 42_i64]
+      consumer_tag = "ctag-1"
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each do |offset|
+          msg_store.store_consumer_offset(consumer_tag, offset)
+        end
+        bytesize = consumer_tag.bytesize + 1 + 8
+        msg_store.@consumer_offsets.size.should eq bytesize*5
+        msg_store.last_offset_by_consumer_tag(consumer_tag).should eq offsets.last
+        msg_store.close
+      end
+    end
+
+    it "compacts consumer tag file on restart" do
+      queue_name = Random::Secure.hex
+      offsets = [84_i64, 24_i64, 1_i64, 100_i64, 42_i64]
+      consumer_tag = "ctag-1"
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each do |offset|
+          msg_store.store_consumer_offset(consumer_tag, offset)
+        end
+        msg_store.close
+
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.last_offset_by_consumer_tag(consumer_tag).should eq offsets.last
+        bytesize = consumer_tag.bytesize + 1 + 8
+        msg_store.@consumer_offsets.size.should eq bytesize
+        msg_store.close
+      end
+    end
+
+    it "compacts consumer tag file when full" do
+      queue_name = Random::Secure.hex
+      offsets = [84_i64, 24_i64, 1_i64, 100_i64, 42_i64]
+      consumer_tag = Random::Secure.hex(32)
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        bytesize = consumer_tag.bytesize + 1 + 8
+
+        offsets = (LavinMQ::Config.instance.segment_size / bytesize).to_i32 + 1
+        offsets.times do |i|
+          msg_store.store_consumer_offset(consumer_tag, i)
+        end
+        msg_store.last_offset_by_consumer_tag(consumer_tag).should eq offsets - 1
+        msg_store.@consumer_offsets.size.should eq bytesize*2
+        msg_store.close
+      end
+    end
+
+    it "does not track offset if x-stream-offset is set" do
+      queue_name = Random::Secure.hex
+      consumer_tag = Random::Secure.hex
+      c_args = AMQP::Client::Arguments.new({"x-stream-offset": 0})
+
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 2)
+        msg = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag, c_args)
+        StreamQueueSpecHelpers.offset_from_headers(msg.properties.headers).should eq 1
+        sleep 0.1
+
+        # should consume the same message again since tracking was not saved from last consume
+        msg_2 = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag)
+        StreamQueueSpecHelpers.offset_from_headers(msg_2.properties.headers).should eq 1
+      end
+    end
+
+    it "should not use saved offset if x-stream-offset is set" do
+      queue_name = Random::Secure.hex
+      consumer_tag = Random::Secure.hex
+      c_args = AMQP::Client::Arguments.new({"x-stream-offset": 0})
+
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 2)
+
+        # get message without x-stream-offset, tracks offset
+        msg = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag)
+        StreamQueueSpecHelpers.offset_from_headers(msg.properties.headers).should eq 1
+        sleep 0.1
+
+        # consume with x-stream-offset set, should consume the same message again
+        msg_2 = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag, c_args)
+        StreamQueueSpecHelpers.offset_from_headers(msg_2.properties.headers).should eq 1
+      end
+    end
+
+    it "should use saved offset if x-stream-offset & x-stream-automatic-offset-tracking is set" do
+      queue_name = Random::Secure.hex
+      consumer_tag = Random::Secure.hex
+      c_args = AMQP::Client::Arguments.new({"x-stream-offset": 0, "x-stream-automatic-offset-tracking": true})
+
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 2)
+
+        # get message without x-stream-offset, tracks offset
+        msg = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag, c_args)
+        StreamQueueSpecHelpers.offset_from_headers(msg.properties.headers).should eq 1
+        sleep 0.1
+
+        # consume with x-stream-offset set, should consume the same message again
+        msg_2 = StreamQueueSpecHelpers.consume_one(s, queue_name, consumer_tag, c_args)
+        StreamQueueSpecHelpers.offset_from_headers(msg_2.properties.headers).should eq 2
+      end
+    end
+
+    it "cleanup_consumer_offsets removes outdated offset" do
+      queue_name = Random::Secure.hex
+      offsets = [84_i64, -10_i64]
+      tag_prefix = "ctag-"
+
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each_with_index do |offset, i|
+          msg_store.store_consumer_offset(tag_prefix + i.to_s, offset)
+        end
+        sleep 0.1
+        msg_store.cleanup_consumer_offsets
+        msg_store.close
+        sleep 0.1
+
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.last_offset_by_consumer_tag(tag_prefix + 1.to_s).should eq nil
+        msg_store.last_offset_by_consumer_tag(tag_prefix + 0.to_s).should eq offsets[0]
+        msg_store.close
+      end
+    end
+
+    it "runs cleanup when removing segment" do
+      consumer_tag = "ctag-1"
+      queue_name = Random::Secure.hex
+      args = {"x-queue-type": "stream", "x-max-length": 1}
+      msg_body = Bytes.new(LavinMQ::Config.instance.segment_size)
+
+      with_amqp_server do |s|
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+
+        with_channel(s) do |ch|
+          q = ch.queue(queue_name, args: AMQP::Client::Arguments.new(args))
+          q.publish_confirm msg_body
+        end
+
+        with_channel(s) do |ch|
+          ch.prefetch 1
+          q = ch.queue(queue_name, args: AMQP::Client::Arguments.new(args))
+          msgs = Channel(AMQP::Client::DeliverMessage).new
+          q.subscribe(no_ack: false, tag: consumer_tag) do |msg|
+            msgs.send msg
+            msg.ack
+          end
+          msgs.receive
+        end
+
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.last_offset_by_consumer_tag(consumer_tag).should eq 2
+
+        with_channel(s) do |ch|
+          q = ch.queue(queue_name, args: AMQP::Client::Arguments.new(args))
+          2.times { q.publish_confirm msg_body }
+        end
+
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.last_offset_by_consumer_tag(consumer_tag).should eq nil
+      end
+    end
+
+    it "does not track offset if c-tag is auto-generated" do
+      queue_name = Random::Secure.hex
+
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+        args = {"x-queue-type": "stream"}
+        c_tag = ""
+        with_channel(s) do |ch|
+          ch.prefetch 1
+          q = ch.queue(queue_name, args: AMQP::Client::Arguments.new(args))
+          msgs = Channel(AMQP::Client::DeliverMessage).new
+          c_tag = q.subscribe(no_ack: false) do |msg|
+            msgs.send msg
+            msg.ack
+          end
+          msgs.receive
+        end
+
+        sleep 0.1
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.last_offset_by_consumer_tag(c_tag).should eq nil
+      end
+    end
+
+    it "expands consumer offset file when needed" do
+      queue_name = Random::Secure.hex
+      consumer_tag_prefix = Random::Secure.hex(32)
+      with_amqp_server do |s|
+        StreamQueueSpecHelpers.publish(s, queue_name, 1)
+        data_dir = File.join(s.vhosts["/"].data_dir, Digest::SHA1.hexdigest queue_name)
+        msg_store = LavinMQ::AMQP::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        one_offset_bytesize = "#{consumer_tag_prefix}1000".bytesize + 1 + 8
+        offsets = (LavinMQ::Config.instance.segment_size / one_offset_bytesize).to_i32 + 1
+        bytesize = 0
+        offsets.times do |i|
+          consumer_tag = "#{consumer_tag_prefix}#{i + 1000}"
+          msg_store.store_consumer_offset(consumer_tag, i + 1000)
+          bytesize += consumer_tag.bytesize + 1 + 8
+        end
+        msg_store.@consumer_offsets.size.should eq bytesize
+        msg_store.@consumer_offsets.size.should be > LavinMQ::Config.instance.segment_size
+        offsets.times do |i|
+          msg_store.last_offset_by_consumer_tag("#{consumer_tag_prefix}#{i + 1000}").should eq i + 1000
+        end
+      end
+    end
+  end
 end

src/lavinmq/amqp/queue/stream_queue.cr

@@ -79,6 +79,10 @@ module LavinMQ::AMQP
       end
     end
 
+    def store_consumer_offset(consumer_tag : String, offset : Int64) : Nil
+      stream_queue_msg_store.store_consumer_offset(consumer_tag, offset)
+    end
+
     # yield the next message in the ready queue
     # returns true if a message was deliviered, false otherwise
     # if we encouncer an unrecoverable ReadError, close queue