WIP; Rework Khepri clustering

In particular which node is picked while trying to make the Khepri
cluster match the Mnesia one.

4 files changed, 200 insertions, 55 deletions

diff --git a/deps/rabbit/src/rabbit_core_ff.erl b/deps/rabbit/src/rabbit_core_ff.erl
index f48b1e2462..d8208d384c 100644
--- a/deps/rabbit/src/rabbit_core_ff.erl
+++ b/deps/rabbit/src/rabbit_core_ff.erl
@@ -75,7 +75,9 @@
     #{desc          => "Use the new Raft-based metadata store [phase 1]",
       doc_url       => "", %% TODO
       stability     => experimental,
-      depends_on    => [virtual_host_metadata, user_limits],
+      depends_on    => [maintenance_mode_status,
+                        user_limits,
+                        virtual_host_metadata],
       migration_fun => {?MODULE, mds_phase1_migration}
      }}).
 
@@ -248,21 +250,27 @@ user_limits_migration(_FeatureName, _FeatureProps, post_enabled_locally) ->
                             rabbit_topic_permission]).
 
 mds_phase1_migration(_FeatureName, _FeatureProps, is_enabled) ->
-    %Tables = ?MDS_PHASE1_TABLES,
-    %is_mds_migration_done(Tables);
+    %% We don't check if the migration was done already because we also need
+    %% to make sure the cluster membership of Khepri matches the Mnesia
+    %% cluster.
     undefined;
 mds_phase1_migration(FeatureName, _FeatureProps, enable) ->
-    Tables = ?MDS_PHASE1_TABLES,
-    case is_mds_migration_done(Tables) of
-        false -> run_mds_phase1_migration(FeatureName, Tables);
-        true  -> ok
+    case ensure_khepri_cluster_matches_mnesia(FeatureName) of
+        ok ->
+            Tables = ?MDS_PHASE1_TABLES,
+            case is_mds_migration_done(Tables) of
+                false -> migrate_tables_to_khepri(FeatureName, Tables);
+                true  -> ok
+            end;
+        Error ->
+            Error
     end;
 mds_phase1_migration(FeatureName, _FeatureProps, post_enabled_locally) ->
     ?assert(rabbit_khepri:is_enabled(non_blocking)),
     Tables = ?MDS_PHASE1_TABLES,
     empty_unused_mnesia_tables(FeatureName, Tables).
 
-run_mds_phase1_migration(FeatureName, Tables) ->
+ensure_khepri_cluster_matches_mnesia(FeatureName) ->
     %% Initialize Khepri cluster based on Mnesia running nodes. Verify that
     %% all Mnesia nodes are running (all == running). It would be more
     %% difficult to add them later to the node when they start.
@@ -281,20 +289,13 @@ run_mds_phase1_migration(FeatureName, Tables) ->
             %% This is the first time Khepri will be used for real. Therefore
             %% we need to make sure the Khepri cluster matches the Mnesia
             %% cluster.
-            %%
-            %% Even if some nodes are already part of the Khepri cluster, we
-            %% reset every nodes (except the one running this function) and
-            %% add them to this node's Khepri cluster.
-            KhepriNodes = lists:sort(rabbit_khepri:nodes()),
-            OtherMnesiaNodes = AllMnesiaNodes -- KhepriNodes,
             ?LOG_DEBUG(
-               "Feature flag `~s`:   updating the Khepri cluster to match the "
-               "Mnesia cluster by adding the following nodes: ~p",
-               [FeatureName, OtherMnesiaNodes]),
-            case expand_khepri_cluster(OtherMnesiaNodes) of
+               "Feature flag `~s`:   updating the Khepri cluster to match "
+               "the Mnesia cluster",
+               [FeatureName]),
+            case expand_khepri_cluster(FeatureName) of
                 ok ->
-                    rabbit_table:wait(Tables, _Retry = true),
-                    migrate_tables_to_khepri(FeatureName, Tables);
+                    ok;
                 Error ->
                     ?LOG_ERROR(
                        "Feature flag `~s`:   failed to migrate from Mnesia "
@@ -311,15 +312,96 @@ run_mds_phase1_migration(FeatureName, Tables) ->
             {error, all_mnesia_nodes_must_run}
     end.
 
-expand_khepri_cluster([MnesiaNode | Rest]) ->
-    case rabbit_khepri:add_member(MnesiaNode) of
-        ok    -> expand_khepri_cluster(Rest);
-        Error -> Error
-    end;
-expand_khepri_cluster([]) ->
-    ok.
+expand_khepri_cluster(FeatureName) ->
+    %% All Mnesia nodes are running (this is a requirement to enable this
+    %% feature flag). We use this unique list of nodes to find the largest
+    %% Khepri clusters among all of them.
+    %%
+    %% The idea is that at the beginning, each Mnesia Node will also be an
+    %% unclustered Khepri node. Therefore, the first node in the sorted list
+    %% of Mnesia nodes will be picked (a "cluster" with 1 member, but the
+    %% largest at the beginning).
+    %%
+    %% After the first nodes join that single node, its cluster will grow and
+    %% will continue to be the largest.
+    %%
+    %% This function is executed on every nodes. Therefore it will add this
+    %% node to a cluster (except if it would join itself). Other nodes will
+    %% execute this function to add themselves.
+    %%
+    %% This should avoid the situation where a large established cluster is
+    %% reset and joins a single new/empty node.
+    %%
+    %% Also, we only consider Khepri clusters which are in use (the feature
+    %% flag is enabled). Here is an example:
+    %%     - Node2 is the only node in the Mnesia cluster at the time the
+    %%       feature flag is enabled. It joins no other node and runs its own
+    %%       one-node Khepri cluster.
+    %%     - Node1 joins the Mnesia cluster which is now Node1 + Node2. Given
+    %%       the sorting, Khepri cluster will be [[Node1], [Node2]] when
+    %%       sorted by name and size. With this order, Node1 should "join"
+    %%       itself. But the feature is not enabled yet on this node,
+    %%       therefore, we skip this cluster to consider the following one,
+    %%       [Node2].
+    ThisNode = node(),
+    KhepriCluster = find_largest_khepri_cluster(FeatureName),
+    case lists:member(ThisNode, KhepriCluster) of
+        true ->
+            ?LOG_DEBUG(
+               "Feature flag `~s`:   this node (~s) is already a member of "
+               "the largest cluster: ~p",
+               [FeatureName, ThisNode, KhepriCluster]),
+            ok;
+        false ->
+            [KhepriNode | _] = KhepriCluster,
+            ?LOG_DEBUG(
+               "Feature flag `~s`:   adding this node (~s) to the largest "
+               "Khepri cluster found among Mnesia nodes: ~p",
+               [FeatureName, ThisNode, KhepriCluster]),
+            rabbit_khepri:join_cluster(KhepriNode)
+    end.
+
+find_largest_khepri_cluster(FeatureName) ->
+    case list_all_khepri_clusters(FeatureName) of
+        [] ->
+            [node()];
+        KhepriClusters ->
+            KhepriClustersBySize = sort_khepri_clusters_by_size(
+                                     KhepriClusters),
+            ?LOG_DEBUG(
+               "Feature flag `~s`:   existing Khepri clusters (sorted by "
+               "size): ~p",
+               [FeatureName, KhepriClustersBySize]),
+            LargestKhepriCluster = hd(KhepriClustersBySize),
+            LargestKhepriCluster
+    end.
+
+list_all_khepri_clusters(FeatureName) ->
+    MnesiaNodes = lists:sort(rabbit_mnesia:cluster_nodes(all)),
+    ?LOG_DEBUG(
+       "Feature flag `~s`:   querying the following Mnesia nodes to learn "
+       "their Khepri cluster membership: ~p",
+       [FeatureName, MnesiaNodes]),
+    KhepriClusters = lists:foldl(
+                       fun(MnesiaNode, Acc) ->
+                               case khepri_cluster_on_node(MnesiaNode) of
+                                   []        -> Acc;
+                                   Cluster   -> Acc#{Cluster => true}
+                               end
+                       end, #{}, MnesiaNodes),
+    lists:sort(maps:keys(KhepriClusters)).
+
+sort_khepri_clusters_by_size(KhepriCluster) ->
+    lists:sort(
+      fun(A, B) -> length(A) >= length(B) end,
+      KhepriCluster).
+
+khepri_cluster_on_node(Node) ->
+    lists:sort(
+      rabbit_misc:rpc_call(Node, rabbit_khepri, nodes_if_khepri_enabled, [])).
 
 migrate_tables_to_khepri(FeatureName, Tables) ->
+    rabbit_table:wait(Tables, _Retry = true),
     ?LOG_NOTICE(
        "Feature flag `~s`:   starting migration from Mnesia "
        "to Khepri; expect decrease in performance and "
diff --git a/deps/rabbit/src/rabbit_khepri.erl b/deps/rabbit/src/rabbit_khepri.erl
index df3793185f..4511eaf274 100644
--- a/deps/rabbit/src/rabbit_khepri.erl
+++ b/deps/rabbit/src/rabbit_khepri.erl
@@ -15,6 +15,7 @@
 
 -export([setup/0,
          setup/1,
+         join_cluster/1,
          add_member/1,
          remove_member/1,
          members/0,
@@ -53,6 +54,7 @@
          info/0,
          is_enabled/0,
          is_enabled/1,
+         nodes_if_khepri_enabled/0,
          try_mnesia_or_khepri/2]).
 -export([priv_reset/0]).
 
@@ -61,7 +63,7 @@
          clear_forced_metadata_store/0]).
 -endif.
 
--compile({no_auto_import, [get/1, get/2]}).
+-compile({no_auto_import, [get/1, get/2, nodes/0]}).
 
 -define(RA_SYSTEM, coordination).
 -define(RA_CLUSTER_NAME, metadata_store).
@@ -124,9 +126,31 @@ wait_for_leader(Fun, Timeout) ->
             Error
     end.
 
-%% FIXME: For both add_member() and remove_member(), we don't need to manage
-%% the `rabbit` application on the node running this code. However, the node
-%% to add or move must not run RabbitMQ.
+join_cluster(RemoteNode) when RemoteNode =/= node() ->
+    ThisNode = node(),
+    ?LOG_INFO(
+      "Khepri clustering: Attempt to add node ~p to Khepri cluster through "
+      "node ~p",
+      [ThisNode, RemoteNode],
+      #{domain => ?RMQLOG_DOMAIN_GLOBAL}),
+    Ret = rabbit_misc:rpc_call(
+            RemoteNode, rabbit_khepri, add_member, [ThisNode]),
+    case Ret of
+        ok ->
+            ?LOG_INFO(
+              "Khepri clustering: Node ~p successfully added to Khepri "
+              "cluster through node ~p",
+              [ThisNode, RemoteNode],
+              #{domain => ?RMQLOG_DOMAIN_GLOBAL}),
+            ok;
+        Error ->
+            ?LOG_INFO(
+              "Khepri clustering: Failed to add node ~p to Khepri cluster "
+              "through ~p: ~p",
+              [ThisNode, RemoteNode, Error],
+              #{domain => ?RMQLOG_DOMAIN_GLOBAL}),
+            Error
+    end.
 
 add_member(NewNode) when NewNode =/= node() ->
     ?LOG_DEBUG(
@@ -146,7 +170,23 @@ add_member(NewNode) when NewNode =/= node() ->
                "Resetting Khepri on remote node ~s",
                [NewNode],
                #{domain => ?RMQLOG_DOMAIN_GLOBAL}),
-            Ret1 = rpc:call(NewNode, rabbit_khepri, priv_reset, []),
+            %% If the remote node to add is running RabbitMQ, we need to put
+            %% it in maintenance mode at least. We remember that state to
+            %% revive the node only if it was fully running before this code.
+            RemoteIsRunning = rabbit:is_running(NewNode),
+            RemoteAlreadyBeingDrained =
+            rabbit_maintenance:is_being_drained_consistent_read(NewNode),
+            NeedToReviveRemote =
+            RemoteIsRunning andalso RemoteAlreadyBeingDrained,
+            case RemoteIsRunning of
+                true ->
+                    ok = rabbit_misc:rpc_call(
+                           NewNode, rabbit_maintenance, drain, []);
+                false ->
+                    ok
+            end,
+            Ret1 = rabbit_misc:rpc_call(
+                     NewNode, rabbit_khepri, priv_reset, []),
             case Ret1 of
                 ok ->
                     ?LOG_DEBUG(
@@ -157,6 +197,15 @@ add_member(NewNode) when NewNode =/= node() ->
                     Ret2 = khepri:add_member(
                              ?RA_SYSTEM, ?RA_CLUSTER_NAME, ?RA_FRIENDLY_NAME,
                              NewNode),
+                    %% Revive the remote node if it was running and not under
+                    %% maintenance before we changed the cluster membership.
+                    case NeedToReviveRemote of
+                        true ->
+                            ok = rabbit_misc:rpc_call(
+                                   NewNode, rabbit_maintenance, revive, []);
+                        false ->
+                            ok
+                    end,
                     case Ret2 of
                         ok ->
                             ?LOG_DEBUG(
@@ -236,7 +285,17 @@ remove_member(NodeToRemove) when NodeToRemove =/= node() ->
     end.
 
 priv_reset() ->
-    ok = rabbit:stop(),
+    %% To reset the content of Khepri, we need RabbitMQ to be under
+    %% maintenance or stopped. The reason is we can't let RabbitMQ do any
+    %% reads from or writes to Khepri because the content will be garbage.
+    %%
+    %% Stopping RabbitMQ would be done by the `stop_app' CLI command.
+    %%
+    %% If RabbitMQ is running, the node is put under maintenance by the
+    %% `add_member/1' function above.
+    ?assert(
+       not rabbit:is_running() orelse
+       rabbit_maintenance:is_being_drained_consistent_read(node())),
     ok = ensure_ra_system_started(),
     ok = khepri:reset(?RA_SYSTEM, ?RA_CLUSTER_NAME).
 
@@ -386,6 +445,12 @@ is_enabled(Blocking) ->
     rabbit_feature_flags:is_enabled(
       raft_based_metadata_store_phase1, Blocking) =:= true.
 
+nodes_if_khepri_enabled() ->
+    case is_enabled(non_blocking) of
+        true  -> nodes();
+        false -> []
+    end.
+
 -ifdef(TEST).
 -define(FORCED_MDS_KEY, {?MODULE, forced_metadata_store}).
 
diff --git a/deps/rabbit/src/rabbit_mnesia.erl b/deps/rabbit/src/rabbit_mnesia.erl
index 14eb90d41c..7a0029b9fd 100644
--- a/deps/rabbit/src/rabbit_mnesia.erl
+++ b/deps/rabbit/src/rabbit_mnesia.erl
@@ -212,9 +212,15 @@ join_discovered_peers_with_retries(TryNodes, NodeType, RetriesLeft, DelayInterva
 %% we cluster to its cluster.
 
 -spec join_cluster(node(), node_type())
-                        -> ok | {ok, already_member} | {error, {inconsistent_cluster, string()}}.
+                        -> ok | {ok, already_member} | {error, {inconsistent_cluster, string()}} | {badrpc, term()}.
 
 join_cluster(DiscoveryNode, NodeType) ->
+    case join_mnesia_cluster(DiscoveryNode, NodeType) of
+        ok    -> join_khepri_cluster(DiscoveryNode);
+        Other -> Other
+    end.
+
+join_mnesia_cluster(DiscoveryNode, NodeType) ->
     ensure_mnesia_not_running(),
     ensure_mnesia_dir(),
     case is_only_clustered_disc_node() of
@@ -258,6 +264,18 @@ join_cluster(DiscoveryNode, NodeType) ->
             end
     end.
 
+join_khepri_cluster(DiscoveryNode) ->
+    case rabbit_khepri:join_cluster(DiscoveryNode) of
+        {badrpc, {'EXIT', {undef, [{rabbit_khepri, add_member, _, _}]}}} ->
+            rabbit_log:debug(
+              "Skip Khepri cluster expansion; Khepri is unavailable on "
+              "remote node ~p",
+              [DiscoveryNode]),
+            ok;
+        Other ->
+            Other
+    end.
+
 %% return node to its virgin state, where it is not member of any
 %% cluster, has no cluster configuration, no local database, and no
 %% persisted messages
diff --git a/deps/rabbitmq_cli/lib/rabbitmq/cli/ctl/commands/join_cluster_command.ex b/deps/rabbitmq_cli/lib/rabbitmq/cli/ctl/commands/join_cluster_command.ex
index 9620b52d71..c92965c8c9 100644
--- a/deps/rabbitmq_cli/lib/rabbitmq/cli/ctl/commands/join_cluster_command.ex
+++ b/deps/rabbitmq_cli/lib/rabbitmq/cli/ctl/commands/join_cluster_command.ex
@@ -42,32 +42,12 @@ defmodule RabbitMQ.CLI.Ctl.Commands.JoinClusterCommand do
     long_or_short_names = Config.get_option(:longnames, opts)
     target_node_normalised = Helpers.normalise_node(target_node, long_or_short_names)
 
-    ret1 = :rabbit_misc.rpc_call(
+    :rabbit_misc.rpc_call(
       node_name,
       :rabbit_mnesia,
       :join_cluster,
       [target_node_normalised, node_type]
     )
-
-    case ret1 do
-      :ok ->
-        ret2 = :rabbit_misc.rpc_call(
-          target_node_normalised,
-          :rabbit_khepri,
-          :add_member,
-          [node_name]
-        )
-
-        # TODO: Handle errors.
-        case ret2 do
-          {:undef, [{:rabbit_khepri, :add_member, _, []}]} ->
-            ret1
-          _ ->
-            ret2
-        end
-      _ ->
-        ret1
-    end
   end
 
   def output({:ok, :already_member}, _) do