Consumer timeout (#617)

* Consumer timeout

* Make it possible to spec consumer_timeout

* Close consuemr timeout check fiber when vhost is closed

spec/server_spec.cr

@@ -1106,4 +1106,14 @@ describe LavinMQ::Server do
       end
     end
   end
+
+  it "supports consumer timeouts" do
+    with_channel do |ch|
+      q = ch.queue("", exclusive: true, args: AMQP::Client::Arguments.new({"x-consumer-timeout": 100}))
+      q.publish_confirm("a")
+      expect_raises(AMQP::Client::Channel::ClosedException, /PRECONDITION/) do
+        q.subscribe(no_ack: false, tag: "c", block: true) { }
+      end
+    end
+  end
 end

spec/spec_helper.cr

@@ -23,6 +23,7 @@ LavinMQ::Config.instance.tap do |cfg|
   cfg.amqps_port = AMQPS_PORT
   cfg.http_port = HTTP_PORT
   cfg.segment_size = 512 * 1024
+  cfg.consumer_timeout_loop_interval = 1
 end
 
 # have to be required after config

src/lavinmq/client/channel.cr

@@ -48,7 +48,8 @@ module LavinMQ
         tag : UInt64,
         queue : Queue,
         sp : SegmentPosition,
-        consumer : Consumer?
+        consumer : Consumer?,
+        delivered_at : Time::Span
 
       def details_tuple
         {
@@ -643,11 +644,29 @@ module LavinMQ
       protected def next_delivery_tag(queue : Queue, sp, no_ack, consumer) : UInt64
         @unack_lock.synchronize do
           tag = @delivery_tag &+= 1
-          @unacked.push Unack.new(tag, queue, sp, consumer) unless no_ack
+          @unacked.push Unack.new(tag, queue, sp, consumer, RoughTime.monotonic) unless no_ack
           tag
         end
       end
 
+      # Iterate over all unacked messages and see if any has been unacked longer than the queue's consumer timeout
+      def check_consumer_timeout
+        @unack_lock.synchronize do
+          queues = Set(Queue).new # only check first delivered message per queue
+          @unacked.each do |unack|
+            if queues.add? unack.queue
+              if timeout = unack.queue.consumer_timeout
+                unacked_ms = RoughTime.monotonic - unack.delivered_at
+                if unacked_ms > timeout.milliseconds
+                  send AMQP::Frame::Channel::Close.new(@id, 406_u16, "PRECONDITION_FAILED - consumer timeout", 60_u16, 20_u16)
+                  break
+                end
+              end
+            end
+          end
+        end
+      end
+
       def has_capacity? : Bool
         return true if @global_prefetch_count.zero?
         prefetch_limit = @global_prefetch_count

src/lavinmq/config.cr

@@ -48,6 +48,8 @@ module LavinMQ
     property replication_bind : String? = nil
     property replication_port = 5679
     property max_deleted_definitions = 8192 # number of deleted queues, unbinds etc that compacts the definitions file
+    property consumer_timeout : UInt64? = nil
+    property consumer_timeout_loop_interval = 60 # seconds
     @@instance : Config = self.new
 
     def self.instance : LavinMQ::Config
@@ -103,6 +105,7 @@ module LavinMQ
         when "free_disk_min"           then @free_disk_min = v.to_i64
         when "free_disk_warn"          then @free_disk_warn = v.to_i64
         when "max_deleted_definitions" then @max_deleted_definitions = v.to_i
+        when "consumer_timeout"        then @consumer_timeout = v.to_u64
         else
           STDERR.puts "WARNING: Unrecognized configuration 'main/#{config}'"
         end

src/lavinmq/queue/queue.cr

@@ -53,6 +53,8 @@ module LavinMQ
     getter? paused = false
     getter paused_change = Channel(Bool).new
 
+    getter consumer_timeout : UInt64? = Config.instance.consumer_timeout
+
     @consumers_empty_change = Channel(Bool).new
 
     private def queue_expire_loop
@@ -226,6 +228,10 @@ module LavinMQ
           @vhost.upstreams.try &.link(v.as_s, self)
         when "federation-upstream-set"
           @vhost.upstreams.try &.link_set(v.as_s, self)
+        when "consumer-timeout"
+          unless @consumer_timeout.try &.< v.as_i64
+            @consumer_timeout = v.as_i64.to_u64
+          end
         end
       end
       @policy = policy
@@ -260,6 +266,8 @@ module LavinMQ
       validate_positive("x-delivery-limit", @delivery_limit)
       @reject_on_overflow = parse_header("x-overflow", String) == "reject-publish"
       @single_active_consumer_queue = parse_header("x-single-active-consumer", Bool) == true
+      @consumer_timeout = parse_header("x-consumer-timeout", Int).try &.to_u64
+      validate_positive("x-consumer-timeout", @consumer_timeout)
     end
 
     private macro parse_header(header, type)

src/lavinmq/vhost.cr

@@ -58,6 +58,19 @@ module LavinMQ
       @shovels = ShovelStore.new(self)
       @upstreams = Federation::UpstreamStore.new(self)
       load!
+      spawn check_consumer_timeouts_loop
+    end
+
+    private def check_consumer_timeouts_loop
+      loop do
+        sleep Config.instance.consumer_timeout_loop_interval
+        return if @closed
+        @connections.each do |c|
+          c.channels.each_value do |ch|
+            ch.check_consumer_timeout
+          end
+        end
+      end
     end
 
     def max_connections=(value : Int32) : Nil