issue #91, queue: full policy support

Allow to enable and disable deduplication via policy as well as
clearing existing ones.

Signed-off-by: Matteo Cafasso <[email protected]>

Author: noxdafox

lib/rabbitmq_message_deduplication/rabbit_message_deduplication_policy_event.ex

@@ -25,6 +25,7 @@ defmodule RabbitMQMessageDeduplication.PolicyEvent do
 
   alias :amqqueue, as: AMQQueue
   alias :gen_event, as: GenEvent
+  alias :rabbit_log, as: RabbitLog
   alias :rabbit_policy, as: RabbitPolicy
   alias :rabbit_amqqueue, as: RabbitQueue
   alias RabbitMQMessageDeduplication.Queue, as: DedupQueue
@@ -55,29 +56,39 @@ defmodule RabbitMQMessageDeduplication.PolicyEvent do
   def init(_), do: {:ok, []}
 
   @impl :gen_event
-  def handle_event({:event, :policy_set, policy, _, _}, state) do
-    case List.keyfind(policy[:definition], "x-message-deduplication", 0) do
-      {"x-message-deduplication", _} -> apply_to_queues()
-      nil -> 0
+  def handle_event({:event, :queue_policy_updated, policy, _, _}, state) do
+    status = case List.keyfind(policy[:definition], "x-message-deduplication", 0) do
+      {"x-message-deduplication", _} -> apply_to_queue(policy)
+      nil -> :ok
     end
 
-    {:ok, state}
+    {status, state}
   end
+
+  @impl :gen_event
+  def handle_event({:event, :queue_policy_cleared, policy, _, _}, state) do
+    {apply_to_queue(policy), state}
+  end
+
+  @impl :gen_event
   def handle_event(_, state), do: {:ok, state}
 
   @impl :gen_event
   def handle_call(_Request, state), do: {:ok, :not_understood, state}
 
   # Apply new policies to matching queues
-  defp apply_to_queues() do
-    for queue <- RabbitQueue.list() |> Enum.map(&RabbitPolicy.set/1) do
-      AMQQueue.get_pid(queue)
-      |> RabbitQueue.run_backing_queue(DedupQueue,
-                                       fn(_, state) ->
-                                         state
-                                         |> DedupQueue.dqstate(queue: queue)
-                                         |> DedupQueue.maybe_enable_dedup_queue()
-                                       end)
-    end
+  defp apply_to_queue(policy) do
+    {:ok, queue} = RabbitQueue.lookup(policy[:name])
+    queue = RabbitPolicy.set(queue)
+
+    RabbitLog.debug("Policy change for queue ~p ~n", [policy[:name]])
+
+    AMQQueue.get_pid(queue)
+    |> RabbitQueue.run_backing_queue(DedupQueue,
+                                     fn(_, state) ->
+                                       state
+                                       |> DedupQueue.dqstate(queue: queue)
+                                       |> DedupQueue.maybe_toggle_dedup_queue()
+                                     end)
   end
 end

lib/rabbitmq_message_deduplication/rabbit_message_deduplication_queue.ex

@@ -22,7 +22,7 @@ defmodule RabbitMQMessageDeduplication.Queue do
 
   """
 
-  import Record, only: [defrecord: 2, defrecord: 3, defrecordp: 2, extract: 2]
+  import Record, only: [defrecord: 2, defrecord: 3, extract: 2]
 
   require RabbitMQMessageDeduplication.Cache
   require RabbitMQMessageDeduplication.Common
@@ -55,7 +55,7 @@ defmodule RabbitMQMessageDeduplication.Queue do
     :P_basic, from_lib: "rabbit_common/include/rabbit_framing.hrl")
 
   defrecord :dqack, [:tag, :header]
-  defrecord :dqstate, [:queue, :queue_state, :dedup_enabled]
+  defrecord :dqstate, [:queue, :queue_state, dedup_enabled: false]
 
   # The passthrough macros call the underlying backing queue functions
   # The suffixes indicate the arity of the return values
@@ -149,7 +149,7 @@ defmodule RabbitMQMessageDeduplication.Queue do
 
   @impl :rabbit_backing_queue
   def init(queue, recovery, callback) do
-    state = maybe_enable_dedup_queue(dqstate(queue: queue))
+    state = maybe_toggle_dedup_queue(dqstate(queue: queue))
 
     passthrough1(state) do
       init(queue, recovery, callback)
@@ -164,10 +164,7 @@ defmodule RabbitMQMessageDeduplication.Queue do
   @impl :rabbit_backing_queue
   def delete_and_terminate(any, state = dqstate(queue: queue, queue_state: qs)) do
     if dedup_queue?(state) do
-      queue
-      |> AMQQueue.get_name()
-      |> Common.cache_name()
-      |> CacheManager.destroy()
+      :ok = delete_cache(queue)
     end
 
     passthrough1(state) do
@@ -468,28 +465,19 @@ defmodule RabbitMQMessageDeduplication.Queue do
 
   # Utility functions
 
-  # Set the state according to whether the queue is a deduplication one or not
-  def maybe_enable_dedup_queue(state = dqstate(queue: queue, queue_state: qs)) do
-    if dedup_queue?(state) do
-      :ok = init_cache(queue)
-      dqstate(queue: queue, queue_state: qs, dedup_enabled: true)
-    else
-      dqstate(queue: queue, queue_state: qs, dedup_enabled: :undefined)
+  # Enable/disable queue-level deduplication
+  def maybe_toggle_dedup_queue(state = dqstate(queue: queue, queue_state: qs)) do
+    cond do
+      enable_dedup_queue?(state) ->
+        :ok = init_cache(queue)
+        dqstate(queue: queue, queue_state: qs, dedup_enabled: true)
+      disable_dedup_queue?(state) ->
+        :ok = delete_cache(queue)
+        dqstate(queue: queue, queue_state: qs, dedup_enabled: false)
+      true -> dqstate(queue: queue, queue_state: qs, dedup_enabled: false)
     end
   end
 
-  # Check if it's a deduplication enabled queue
-  defp dedup_queue?(dqstate(dedup_enabled: val)) when is_boolean(val), do: val
-  defp dedup_queue?(dqstate(queue: queue)) do
-    args = AMQQueue.get_arguments(queue)
-    pols = case AMQQueue.get_policy(queue) do
-             :undefined -> []
-             policy -> policy[:definition]
-           end
-
-    Common.rabbit_argument(args ++ pols, "x-message-deduplication", default: false)
-  end
-
   # Initialize the deduplication cache
   defp init_cache(queue) do
     cache = queue |> AMQQueue.get_name() |> Common.cache_name()
@@ -509,6 +497,15 @@ defmodule RabbitMQMessageDeduplication.Queue do
     end
   end
 
+  # Remove the cache and all its content
+  defp delete_cache(queue) do
+    cache = queue |> AMQQueue.get_name() |> Common.cache_name()
+
+    RabbitLog.debug("Deleting queue deduplication cache ~s~n", [cache])
+
+    CacheManager.destroy(cache)
+  end
+
   # Returns true if the message is a duplicate.
   defp duplicate?(queue, message = basic_message()) do
     name = AMQQueue.get_name(queue)
@@ -559,4 +556,30 @@ defmodule RabbitMQMessageDeduplication.Queue do
       :exit, {:noproc, {GenServer, :call, [^cache | _]}} -> []
     end
   end
+
+  # True if `x-message-deduplication` is present within queue arguments or policy
+  defp dedup_arg?(queue) do
+    queue_policy(queue)
+    ++ AMQQueue.get_arguments(queue)
+    |> Common.rabbit_argument("x-message-deduplication", default: false)
+  end
+
+  # Return the list of policy arguments assigned to the queue
+  defp queue_policy(queue) do
+    case AMQQueue.get_policy(queue) do
+      :undefined -> []
+      policy -> policy[:definition]
+    end
+  end
+
+  # True if it's an active deduplication queue
+  defp dedup_queue?(dqstate(dedup_enabled: val)), do: val
+
+  # True if deduplication should be enabled for the queue
+  defp enable_dedup_queue?(dqstate(dedup_enabled: true)), do: false
+  defp enable_dedup_queue?(dqstate(queue: q, dedup_enabled: false)), do: dedup_arg?(q)
+
+  # True if deduplication should be disabled for the queue
+  defp disable_dedup_queue?(dqstate(dedup_enabled: false)), do: false
+  defp disable_dedup_queue?(dqstate(queue: q, dedup_enabled: true)), do: not dedup_arg?(q)
 end