Drop msg on requeue if above delivery_limit (#996)

Drops messages directly on requeue if they have been redelivered more than x-delivery-limit times, instead of waiting for the next time we try to deliver them. Also drops messages when a new policy is applied and removes the check on delivery-count on delivery.

Fixes #976

Author: viktorerlingsson

spec/policies_spec.cr

@@ -204,6 +204,25 @@ describe LavinMQ::VHost do
     end
   end
 
+  it "should drop messages if above delivery-limit" do
+    with_amqp_server do |s|
+      defs = {"delivery-limit" => JSON::Any.new(0_i64)} of String => JSON::Any
+      with_channel(s) do |ch|
+        args = AMQP::Client::Arguments.new
+        args["x-delivery-limit"] = 5
+        q = ch.queue("delivery-limit", exclusive: true, args: args)
+        q.publish_confirm "m1"
+        q.publish_confirm "m2"
+        q.get(no_ack: false).not_nil!.reject(requeue: true)
+        ch.queue_declare("delivery-limit", passive: true)[:message_count].should eq 2
+        s.vhosts["/"].add_policy("delivery-limit", "^.*$", "all", defs, 12_i8)
+        sleep 10.milliseconds
+        ch.queue_declare("delivery-limit", passive: true)[:message_count].should eq 1
+        s.vhosts["/"].delete_policy("delivery-limit")
+      end
+    end
+  end
+
   describe "with max-length-bytes policy applied" do
     it "should replace with max-length" do
       with_amqp_server do |s|

spec/server_spec.cr

@@ -959,7 +959,7 @@ describe LavinMQ::Server do
     with_amqp_server do |s|
       with_channel(s) do |ch|
         args = AMQP::Client::Arguments.new
-        args["x-delivery-limit"] = 2
+        args["x-delivery-limit"] = 1
         q = ch.queue("delivery_limit", args: args)
         q.publish "m1"
         msg = q.get(no_ack: false).not_nil!
@@ -969,7 +969,6 @@ describe LavinMQ::Server do
         msg.properties.headers.not_nil!["x-delivery-count"].as(Int32).should eq 1
         msg.reject(requeue: true)
         Fiber.yield
-        q.get(no_ack: false).should be_nil
         s.vhosts["/"].queues["delivery_limit"].empty?.should be_true
       end
     end

src/lavinmq/amqp/queue/queue.cr

@@ -236,6 +236,7 @@ module LavinMQ::AMQP
         when "delivery-limit"
           unless @delivery_limit.try &.< v.as_i64
             @delivery_limit = v.as_i64
+            drop_redelivered
           end
         when "federation-upstream"
           @vhost.upstreams.try &.link(v.as_s, self)
@@ -514,6 +515,27 @@ module LavinMQ::AMQP
       end
     end
 
+    private def drop_redelivered : Nil
+      counter = 0
+      if limit = @delivery_limit
+        @msg_store_lock.synchronize do
+          loop do
+            env = @msg_store.first? || break
+            delivery_count = @deliveries.fetch(env.segment_position, 0) || break
+            break unless delivery_count > limit
+            env = @msg_store.shift? || break
+            @log.debug { "Over delivery limit, drop sp=#{env.segment_position}" }
+            expire_msg(env, :delivery_limit)
+            counter &+= 1
+            if counter >= 16 * 1024
+              Fiber.yield
+              counter = 0
+            end
+          end
+        end
+      end
+    end
+
     private def time_to_message_expiration : Time::Span?
       env = @msg_store_lock.synchronize { @msg_store.first? } || return
       @log.debug { "Checking if message #{env.message} has to be expired" }
@@ -789,15 +811,10 @@ module LavinMQ::AMQP
     end
 
     private def with_delivery_count_header(env) : Envelope?
-      if limit = @delivery_limit
+      if @delivery_limit
         sp = env.segment_position
         headers = env.message.properties.headers || AMQP::Table.new
         delivery_count = @deliveries.fetch(sp, 0)
-        # @log.debug { "Delivery count: #{delivery_count} Delivery limit: #{@delivery_limit}" }
-        if delivery_count >= limit
-          expire_msg(env, :delivery_limit)
-          return nil
-        end
         headers["x-delivery-count"] = delivery_count if delivery_count > 0 # x-delivery-count not included in first delivery
         @deliveries[sp] = delivery_count + 1
         env.message.properties.headers = headers
@@ -834,6 +851,11 @@ module LavinMQ::AMQP
         if has_expired?(sp, requeue: true) # guarantee to not deliver expired messages
           expire_msg(sp, :expired)
         else
+          if delivery_limit = @delivery_limit
+            if @deliveries.fetch(sp, 0) > delivery_limit
+              return expire_msg(sp, :delivery_limit)
+            end
+          end
           @msg_store_lock.synchronize do
             @msg_store.requeue(sp)
           end