Automatic offset tracking for stream queues

spec/stream_queue_spec.cr

@@ -1,6 +1,38 @@
 require "./spec_helper"
 require "./../src/lavinmq/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::StreamQueue do
   stream_queue_args = LavinMQ::AMQP::Table.new({"x-queue-type": "stream"})
 
@@ -203,4 +235,239 @@ describe LavinMQ::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::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each_with_index do |offset, i|
+          msg_store.update_consumer_offset(tag_prefix + i.to_s, offset)
+        end
+        msg_store.close
+        sleep 0.1
+
+        msg_store = LavinMQ::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 "only saves one entry per consumer tag" 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::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each do |offset|
+          msg_store.update_consumer_offset(consumer_tag, offset)
+        end
+        msg_store.close
+        sleep 0.1
+
+        msg_store = LavinMQ::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.last_offset_by_consumer_tag(consumer_tag).should eq offsets.last
+        msg_store.@consumer_offsets.size.should eq 15
+
+        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::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        offsets.each_with_index do |offset, i|
+          msg_store.update_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::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::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::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::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 = "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::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        2000.times do |i|
+          next if i == 0
+          msg_store.update_consumer_offset("#{consumer_tag_prefix}#{i}", i)
+        end
+        msg_store.close
+
+        msg_store = LavinMQ::StreamQueue::StreamQueueMessageStore.new(data_dir, nil)
+        msg_store.@consumer_offsets.size.should eq 34_875
+
+        2000.times do |i|
+          next if i == 0
+          msg_store.last_offset_by_consumer_tag("#{consumer_tag_prefix}#{i}").should eq i
+        end
+
+        msg_store.close
+      end
+    end
+  end
 end

src/lavinmq/client/channel/stream_consumer.cr

@@ -9,11 +9,13 @@ module LavinMQ
         property segment : UInt32
         property pos : UInt32
         getter requeued = Deque(SegmentPosition).new
+        @track_offset = false
 
         def initialize(@channel : Client::Channel, @queue : StreamQueue, frame : AMQP::Frame::Basic::Consume)
+          @tag = frame.consumer_tag
           validate_preconditions(frame)
           offset = frame.arguments["x-stream-offset"]?
-          @offset, @segment, @pos = stream_queue.find_offset(offset)
+          @offset, @segment, @pos = stream_queue.find_offset(offset, @tag, @track_offset)
           super
         end
 
@@ -34,7 +36,10 @@ module LavinMQ
             raise Error::PreconditionFailed.new("x-priority not supported on stream queues")
           end
           case frame.arguments["x-stream-offset"]?
-          when Nil, Int, Time, "first", "next", "last"
+          when Nil
+            @track_offset = true unless @tag.starts_with?("amq.ctag-")
+          when Int, Time, "first", "next", "last"
+            @track_offset = true if frame.arguments["x-stream-automatic-offset-tracking"]?
           else raise Error::PreconditionFailed.new("x-stream-offset must be an integer, a timestamp, 'first', 'next' or 'last'")
           end
         end
@@ -82,6 +87,11 @@ module LavinMQ
           @queue.as(StreamQueue)
         end
 
+        def ack(sp)
+          stream_queue.update_consumer_offset(@tag, @offset) if @track_offset
+          super
+        end
+
         def reject(sp, requeue : Bool)
           super
           if requeue

src/lavinmq/mfile.cr

@@ -229,9 +229,9 @@ class MFile < IO
     Bytes.new(buffer, @size, read_only: true)
   end
 
-  def to_slice(pos, size)
+  def to_slice(pos, size, read_only = true)
     raise IO::EOFError.new if pos + size > @size
-    Bytes.new(buffer + pos, size, read_only: true)
+    Bytes.new(buffer + pos, size, read_only: read_only)
   end
 
   def advise(advice : Advice, addr = buffer, offset = 0, length = @capacity) : Nil
@@ -261,4 +261,9 @@ class MFile < IO
     raise IO::Error.from_errno("pread") if cnt == -1
     cnt
   end
+
+  def rename(new_path : String) : Nil
+    File.rename @path, new_path
+    @path = new_path
+  end
 end

src/lavinmq/queue/stream_queue.cr

@@ -74,6 +74,10 @@ module LavinMQ
       end
     end
 
+    def update_consumer_offset(consumer_tag : String, offset : Int64) : Nil
+      stream_queue_msg_store.update_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