Support transient queues in Khepri

dumbbell · dumbbell · commit 8119283f5249 · 2023-09-28T17:57:44.000+02:00
diff --git a/deps/rabbit/src/amqqueue.erl b/deps/rabbit/src/amqqueue.erl
@@ -78,6 +78,8 @@
          pattern_match_all/0,
          pattern_match_on_name/1,
          pattern_match_on_type/1,
+         pattern_match_on_durable/1,
+         pattern_match_on_type_and_durable/2,
          reset_mirroring_and_decorators/1,
          set_immutable/1,
          qnode/1,
@@ -590,6 +592,17 @@ pattern_match_on_name(Name) ->
 pattern_match_on_type(Type) ->
     #amqqueue{type = Type, _ = '_'}.
 
+-spec pattern_match_on_durable(boolean()) -> amqqueue_pattern().
+
+pattern_match_on_durable(IsDurable) ->
+    #amqqueue{durable = IsDurable, _ = '_'}.
+
+-spec pattern_match_on_type_and_durable(atom(), boolean()) ->
+    amqqueue_pattern().
+
+pattern_match_on_type_and_durable(Type, IsDurable) ->
+    #amqqueue{type = Type, durable = IsDurable, _ = '_'}.
+
 -spec reset_mirroring_and_decorators(amqqueue()) -> amqqueue().
 
 reset_mirroring_and_decorators(#amqqueue{} = Queue) ->
diff --git a/deps/rabbit/src/rabbit_db_queue.erl b/deps/rabbit/src/rabbit_db_queue.erl
@@ -7,9 +7,12 @@
 
 -module(rabbit_db_queue).
 
+-include_lib("kernel/include/logger.hrl").
+-include_lib("stdlib/include/assert.hrl").
+-include_lib("stdlib/include/qlc.hrl").
+
 -include_lib("khepri/include/khepri.hrl").
 -include_lib("rabbit_common/include/rabbit.hrl").
--include_lib("stdlib/include/qlc.hrl").
 -include("amqqueue.hrl").
 
 -export([
@@ -169,7 +172,8 @@ get_all_durable_in_mnesia() ->
 get_all_durable_in_khepri() ->
     list_with_possible_retry_in_khepri(
       fun() ->
-              ets:tab2list(?KHEPRI_PROJECTION)
+              Pattern = amqqueue:pattern_match_on_durable(true),
+              ets:match_object(?KHEPRI_PROJECTION, Pattern)
       end).
 
 -spec get_all_durable_by_type(Type) -> [Queue] when
@@ -193,7 +197,7 @@ get_all_durable_by_type_in_mnesia(Type) ->
     rabbit_db:list_in_mnesia(?MNESIA_DURABLE_TABLE, Pattern).
 
 get_all_durable_by_type_in_khepri(Type) ->
-    Pattern = amqqueue:pattern_match_on_type(Type),
+    Pattern = amqqueue:pattern_match_on_type_and_durable(Type, true),
     ets:match_object(?KHEPRI_PROJECTION, Pattern).
 
 %% -------------------------------------------------------------------
@@ -227,7 +231,7 @@ filter_all_durable_in_mnesia(FilterFun) ->
 filter_all_durable_in_khepri(FilterFun) ->
     ets:foldl(
       fun(Q, Acc0) ->
-              case FilterFun(Q) of
+              case amqqueue:is_durable(Q) andalso FilterFun(Q) of
                   true -> [Q | Acc0];
                   false -> Acc0
               end
@@ -478,12 +482,23 @@ get_in_khepri(Name) ->
 get_durable(Name) ->
     rabbit_khepri:handle_fallback(
       #{mnesia => fun() -> get_durable_in_mnesia(Name) end,
-        khepri => fun() -> get_in_khepri(Name) end
+        khepri => fun() -> get_durable_in_khepri(Name) end
        }).
 
 get_durable_in_mnesia(Name) ->
     rabbit_mnesia:dirty_read({?MNESIA_DURABLE_TABLE, Name}).
 
+get_durable_in_khepri(Name) ->
+    case get_in_khepri(Name) of
+        {ok, Queue} = Ret ->
+            case amqqueue:is_durable(Queue) of
+                true  -> Ret;
+                false -> {error, not_found}
+            end;
+        Error ->
+            Error
+    end.
+
 %% -------------------------------------------------------------------
 %% get_many_durable().
 %% -------------------------------------------------------------------
@@ -494,10 +509,17 @@ get_durable_in_mnesia(Name) ->
 
 get_many_durable(Names) when is_list(Names) ->
     rabbit_khepri:handle_fallback(
-      #{mnesia => fun() -> get_many_in_ets(?MNESIA_DURABLE_TABLE, Names) end,
-        khepri => fun() -> get_many_in_ets(?KHEPRI_PROJECTION, Names) end
+      #{mnesia => fun() -> get_many_durable_in_mnesia(Names) end,
+        khepri => fun() -> get_many_durable_in_khepri(Names) end
        }).
 
+get_many_durable_in_mnesia(Names) ->
+    get_many_in_ets(?MNESIA_DURABLE_TABLE, Names).
+
+get_many_durable_in_khepri(Names) ->
+    Queues = get_many_in_ets(?KHEPRI_PROJECTION, Names),
+    [Q || Q <- Queues, amqqueue:is_durable(Q)].
+
 %% -------------------------------------------------------------------
 %% update().
 %% -------------------------------------------------------------------
@@ -632,7 +654,9 @@ update_durable_in_khepri(UpdateFun, FilterFun) ->
       fun() ->
               khepri_tx:foreach(Path,
                                 fun(Path0, #{data := Q}) ->
-                                        case FilterFun(Q) of
+                                        DoUpdate = amqqueue:is_durable(Q)
+                                            andalso FilterFun(Q),
+                                        case DoUpdate of
                                             true ->
                                                 khepri_tx:put(Path0, UpdateFun(Q));
                                             false ->
@@ -833,7 +857,7 @@ set_in_mnesia(Q) ->
       end).
 
 set_in_mnesia_tx(DurableQ, Q) ->
-    case ?amqqueue_is_durable(Q) of
+    case amqqueue:is_durable(Q) of
         true ->
             ok = mnesia:write(?MNESIA_DURABLE_TABLE, DurableQ, write);
         false ->
@@ -852,8 +876,9 @@ set_in_khepri(Q) ->
 -spec set_many([Queue]) -> ok when
       Queue :: amqqueue:amqqueue().
 %% @doc Writes a list of durable queue records.
-%% It is responsibility of the calling function to ensure all records are durable.
-%% Once transient entities are deprecated, this is a non-issue.
+%%
+%% It is responsibility of the calling function to ensure all records are
+%% durable.
 %%
 %% @private
 
@@ -868,7 +893,10 @@ set_many_in_mnesia(Qs) ->
         %% Just to be nested in forget_node_for_queue
         mnesia:transaction(
           fun() ->
-                  [ok = mnesia:write(?MNESIA_DURABLE_TABLE, Q, write) || Q <- Qs],
+                  [begin
+                       ?assert(amqqueue:is_durable(Q)),
+                       ok = mnesia:write(?MNESIA_DURABLE_TABLE, Q, write)
+                   end || Q <- Qs],
                   ok
           end),
     ok.
@@ -877,6 +905,7 @@ set_many_in_khepri(Qs) ->
     rabbit_khepri:transaction(
       fun() ->
               [begin
+                   ?assert(amqqueue:is_durable(Q)),
                    Path = khepri_queue_path(amqqueue:get_name(Q)),
                    case khepri_tx:put(Path, Q) of
                        ok      -> ok;
@@ -902,7 +931,7 @@ set_many_in_khepri(Qs) ->
 delete_transient(FilterFun) ->
     rabbit_khepri:handle_fallback(
       #{mnesia => fun() -> delete_transient_in_mnesia(FilterFun) end,
-        khepri => fun() -> ok end
+        khepri => fun() -> delete_transient_in_khepri(FilterFun) end
        }).
 
 delete_transient_in_mnesia(FilterFun) ->
@@ -944,6 +973,33 @@ partition_queues([Q0,Q1,Q2,Q3,Q4,Q5,Q6,Q7,Q8,Q9 | T]) ->
 partition_queues(T) ->
     [T].
 
+delete_transient_in_khepri(FilterFun) ->
+    PathPattern = khepri_queues_path() ++
+        [?KHEPRI_WILDCARD_STAR,
+         #if_data_matches{
+            pattern = amqqueue:pattern_match_on_durable(false)}],
+    Ret = rabbit_khepri:fold(
+            PathPattern,
+            fun(Path, #{data := Queue}, Acc) ->
+                    case FilterFun(Queue) of
+                        true ->
+                            QueueName = khepri_queue_path_to_name(Path),
+                            case delete_in_khepri(QueueName, false) of
+                                ok        -> Acc;
+                                Deletions -> [{QueueName, Deletions} | Acc]
+                            end;
+                        false ->
+                            Acc
+                    end
+            end, []),
+    case Ret of
+        {ok, Items} ->
+            {QueueNames, Deletions} = lists:unzip(Items),
+            {QueueNames, lists:flatten(Deletions)};
+        {error, _} = Error ->
+            Error
+    end.
+
 %% -------------------------------------------------------------------
 %% foreach_transient().
 %% -------------------------------------------------------------------
@@ -958,7 +1014,7 @@ partition_queues(T) ->
 foreach_transient(UpdateFun) ->
     rabbit_khepri:handle_fallback(
       #{mnesia => fun() -> foreach_transient_in_mnesia(UpdateFun) end,
-        khepri => fun() -> ok end
+        khepri => fun() -> throw(not_implemented) end
        }).
 
 foreach_transient_in_mnesia(UpdateFun) ->
@@ -1001,7 +1057,10 @@ foreach_durable_in_mnesia(UpdateFun, FilterFun) ->
     ok.
 
 foreach_durable_in_khepri(UpdateFun, FilterFun) ->
-    Path = khepri_queues_path() ++ [rabbit_khepri:if_has_data_wildcard()],
+    Path = khepri_queues_path() ++
+        [?KHEPRI_WILDCARD_STAR,
+         #if_data_matches{
+            pattern = amqqueue:pattern_match_on_durable(true)}],
     case rabbit_khepri:filter(Path, fun(_, #{data := Q}) ->
                                             FilterFun(Q)
                                     end) of
@@ -1190,3 +1249,6 @@ khepri_queues_path() ->
 
 khepri_queue_path(#resource{virtual_host = VHost, name = Name}) ->
     [?MODULE, queues, VHost, Name].
+
+khepri_queue_path_to_name([?MODULE, queues, VHost, Name]) ->
+    rabbit_misc:r(VHost, queue, Name).
diff --git a/deps/rabbit/src/rabbit_khepri.erl b/deps/rabbit/src/rabbit_khepri.erl
@@ -33,6 +33,7 @@
          update/2,
          cas/3,
          fold/3,
+         foreach/2,
          filter/2,
 
          get/1,
@@ -669,6 +670,8 @@ cas(Path, Pattern, Data) ->
 
 fold(Path, Pred, Acc) -> khepri:fold(?STORE_ID, Path, Pred, Acc).
 
+foreach(Path, Pred) -> khepri:foreach(?STORE_ID, Path, Pred).
+
 filter(Path, Pred) -> khepri:filter(?STORE_ID, Path, Pred).
 
 get(Path) ->
diff --git a/deps/rabbit/test/bindings_SUITE.erl b/deps/rabbit/test/bindings_SUITE.erl
@@ -910,7 +910,7 @@ transient_queue_on_node_down_khepri(Config) ->
 
     rabbit_ct_broker_helpers:stop_node(Config, Server),
 
-    Bindings1 = lists:sort([DirectBinding, DirectAltBinding]),
+    Bindings1 = lists:sort([DirectBinding]),
     ?awaitMatch(Bindings1,
                 lists:sort(
                   rabbit_ct_broker_helpers:rpc(Config, 1, rabbit_binding, list, [<<"/">>])),
@@ -921,9 +921,10 @@ transient_queue_on_node_down_khepri(Config) ->
 
     rabbit_ct_broker_helpers:start_node(Config, Server),
 
-    ?awaitMatch([_, _], rabbit_ct_broker_helpers:rpc(Config, 1, rabbit_amqqueue, list, [<<"/">>]),
+    Bindings2 = lists:sort([DefaultBinding, DirectBinding]),
+    ?awaitMatch([_], rabbit_ct_broker_helpers:rpc(Config, 1, rabbit_amqqueue, list, [<<"/">>]),
                 30000),
-    ?awaitMatch(Bindings,
+    ?awaitMatch(Bindings2,
                 lists:sort(
                   rabbit_ct_broker_helpers:rpc(Config, 1, rabbit_binding, list, [<<"/">>])),
                 30000),
