bet365 · nsaadouni · Feb 3, 2020
diff --git a/include/riak_repl.hrl b/include/riak_repl.hrl
@@ -49,6 +49,9 @@
 -define(BT_META_BUCKET_NAME, bucket_name).
 -define(BT_META_BLACKLIST, realtime_blacklist).
 
+-define(RIAK_REPL2_RTQ_CONFIG_KEY, {riak_repl2_rtq, config}).
+-define(RTSOURCE_REALTIME_CONNECTIONS_KEY, {riak_repl2_rtsource, realtime_connections}).
+
 -type(ip_addr_str() :: string()).
 -type(ip_portnum() :: non_neg_integer()).
 -type(repl_addr() :: {ip_addr_str(), ip_portnum()}).

diff --git a/src/riak_core_cluster_mgr.erl b/src/riak_core_cluster_mgr.erl
@@ -807,7 +807,7 @@ shuffle(List) ->
 get_my_remote_ip_list(_Remote, [], _Return) ->
     {ok, []};
 get_my_remote_ip_list(Remote, RemoteUnsorted, Return) ->
-    case riak_repl2_rtsource_conn_data_mgr:read(active_nodes) of
+    case riak_repl_util:read_realtime_active_nodes() of
         no_leader ->
             {ok, []};
         %% SourceNodes = [SourceNode1, SourceNode3, SourceNode2 ...]
@@ -871,7 +871,7 @@ build_primary_dict([{Key, Value}| Rest], Dict) ->
 %% sink nodes to an index which would cause active connections to be dropped and made by other nodes.
 %% Output: dict -> Key-Value: {Index, SinkNode}
 link_addrs(AllPrimariesDict, SinkNodes, Remote) ->
-    case riak_repl2_rtsource_conn_data_mgr:read(realtime_connections, Remote) of
+    case riak_repl_util:read_realtime_endpoints(Remote) of
         no_leader ->
             no_leader;
         %% ActiveConnsDict: this is a dictionary of active realtime connections

diff --git a/src/riak_repl2_rtsink_conn.erl b/src/riak_repl2_rtsink_conn.erl
@@ -493,36 +493,13 @@ riak_repl2_rtsink_conn_test_() ->
       }]}.
 
 setup() ->
-    riak_repl_test_util:start_test_ring(),
-    riak_repl_test_util:abstract_gen_tcp(),
-    riak_repl_test_util:kill_and_wait(riak_repl2_rt),
-    {ok, _RT} = riak_repl2_rt:start_link(),
-    riak_repl_test_util:kill_and_wait(riak_repl2_rtq),
-    {ok, _} = riak_repl2_rtq:start_link(),
-    application:set_env(riak_repl, realtime_connection_rebalance_max_delay_secs, 1000),
-
     catch(meck:unload(riak_core_connection_mgr)),
     meck:new(riak_core_connection_mgr, [passthrough]),
     meck:expect(riak_core_connection_mgr, disconnect,
         fun(_Remote) ->
             ok
         end),
 
-    catch(meck:unload(riak_repl2_rtsource_conn_data_mgr)),
-    meck:new(riak_repl2_rtsource_conn_data_mgr, [passthrough]),
-    meck:expect(riak_repl2_rtsource_conn_data_mgr, read, fun(active_nodes) -> [node()]
-                                                         end),
-    meck:expect(riak_repl2_rtsource_conn_data_mgr, read,
-        fun(realtime_connections, _Remote) ->
-            dict:new()
-        end
-    ),
-    meck:expect(riak_repl2_rtsource_conn_data_mgr, write, 4,
-        fun(_, _, _, _) ->
-            ok
-        end
-    ),
-
     catch(meck:unload(riak_core_cluster_mgr)),
     meck:new(riak_core_cluster_mgr, [passthrough]),
     meck:expect(riak_core_cluster_mgr, get_unshuffled_ipaddrs_of_cluster, fun(_Remote) -> [] end ),
@@ -552,6 +529,15 @@ setup() ->
         {reply, ok,  #state{transport=Transport, socket=Socket}}
                                                end),
 
+
+
+    riak_repl_test_util:start_test_ring(),
+    riak_repl_test_util:abstract_gen_tcp(),
+    riak_repl_test_util:kill_and_wait(riak_repl2_rt),
+    {ok, _RT} = riak_repl2_rt:start_link(),
+    riak_repl_test_util:kill_and_wait(riak_repl2_rtq),
+    {ok, _} = riak_repl2_rtq:start_link(),
+    application:set_env(riak_repl, realtime_connection_rebalance_max_delay_secs, 1000),
     folsom:start(),
 
     ok.
@@ -565,7 +551,7 @@ cleanup(_Ctx) ->
     riak_repl_test_util:stop_test_ring(),
     riak_repl_test_util:maybe_unload_mecks(
       [riak_core_service_mgr,
-        riak_repl2_rtsource_conn_data_mgr,
+        riak_repl_util,
         riak_core_connection_mgr,
        riak_repl_util,
        riak_core_tcp_mon,
@@ -632,6 +618,14 @@ start_source(NegotiatedVer) ->
     end),
 
 
+    meck:expect(riak_repl_util, read_realtime_active_nodes, fun() -> [node()] end),
+    meck:expect(riak_repl_util, read_realtime_active_nodes, fun(_) -> [node()] end),
+    meck:expect(riak_repl_util, read_realtime_endpoints, fun(_) -> dict:new() end),
+    meck:expect(riak_repl_util, read_realtime_endpoints, fun(_, _) -> dict:new() end),
+    meck:expect(riak_repl_util, write_realtime_endpoints, fun(_, _) -> ok end),
+    meck:expect(riak_repl_util, write_realtime_endpoints, fun(_, _, _) -> ok end),
+    meck:expect(riak_repl_util, write_realtime_endpoints, fun(_, _, _, _) -> ok end),
+    meck:expect(riak_repl_util, delete_realtime_endpoints, fun(_) -> ok end),
     {ok, SourcePid} = riak_repl2_rtsource_conn_mgr:start_link("sink_cluster"),
     receive
         {sink_started, SinkPid} ->
@@ -653,7 +647,7 @@ unload_mecks() ->
     riak_repl_test_util:maybe_unload_mecks([
         stateful, riak_core_ring_manager, riak_core_ring,
         riak_repl2_rtsink_helper, gen_tcp, fake_source, riak_repl2_rtq,
-        riak_core_capability, riak_repl2_rtsource_conn_data_mgr]).
+        riak_core_capability, riak_repl2_rtsource_conn_data_mgr, riak_repl_util]).
 
 
 -endif.
diff --git a/src/riak_repl2_rtsource_conn_data_mgr.erl b/src/riak_repl2_rtsource_conn_data_mgr.erl
@@ -30,8 +30,7 @@
     is_leader,
     connections,
     active_nodes,
-    restoration,
-    core_capability_polling_interval
+    restoration
 
 }).
 
@@ -78,23 +77,27 @@ node_watcher_update(_Services) ->
 %%%===================================================================
 
 init([]) ->
+    riak_core_node_watcher_events:add_sup_callback(fun ?MODULE:node_watcher_update/1),
+
+    CoreCapabilityPollingIntervalSecs = app_helper:get_env(riak_repl, realtime_core_capability_polling_interval, 60),
+    CoreCapabilityPollingInterval = CoreCapabilityPollingIntervalSecs * 1000,
+    erlang:send_after(CoreCapabilityPollingInterval, self(), poll_core_capability),
+
     Version = riak_core_capability:get({riak_repl, realtime_connections}, legacy),
     State = case Version of
                 legacy ->
                     legacy_init(#state{});
                 v1 ->
-                    v1_init(#state{})
+                    v1_init(#state{});
+                v2 ->
+                    #state{version = v2}
             end,
     {ok, State}.
 
 
 legacy_init(State) ->
-    riak_core_node_watcher_events:add_sup_callback(fun ?MODULE:node_watcher_update/1),
-    CoreCapabilityPollingIntervalSecs = app_helper:get_env(riak_repl, realtime_core_capability_polling_interval, 60),
-    CoreCapabilityPollingInterval = CoreCapabilityPollingIntervalSecs * 1000,
-    erlang:send_after(CoreCapabilityPollingInterval, self(), poll_core_capability),
     State#state{version = legacy, leader_node = undefined, is_leader = false, connections = dict:new(), active_nodes = [],
-        restoration = false, core_capability_polling_interval = CoreCapabilityPollingInterval}.
+        restoration = false}.
 
 v1_init(State) ->
     {Leader, IsLeader} =
@@ -124,11 +127,7 @@ v1_init(State) ->
 
     C = dict:new(),
     AN = [],
-    riak_core_node_watcher_events:add_sup_callback(fun ?MODULE:node_watcher_update/1),
-    CoreCapabilityPollingIntervalSecs = app_helper:get_env(riak_repl, realtime_core_capability_polling_interval, 60),
-    CoreCapabilityPollingInterval = CoreCapabilityPollingIntervalSecs * 1000,
-    State#state{version = v1, leader_node = Leader, is_leader = IsLeader, connections = C, active_nodes = AN, restoration = false,
-        core_capability_polling_interval = CoreCapabilityPollingInterval}.
+    State#state{version = v1, leader_node = Leader, is_leader = IsLeader, connections = C, active_nodes = AN, restoration = false}.
 
 
 %% -------------------------------------------------- Read ---------------------------------------------------------- %%
@@ -182,6 +181,8 @@ handle_call(_Request, _From, State) ->
     {reply, ok, State}.
 
 
+handle_cast({set_leader_node, _LeaderNode}, State=#state{version = v2}) ->
+    {noreply, State};
 handle_cast({set_leader_node, LeaderNode}, State=#state{version = V}) ->
     lager:info("setting leader node as: ~p", [LeaderNode]),
     case {V, node()} of
@@ -190,13 +191,15 @@ handle_cast({set_leader_node, LeaderNode}, State=#state{version = V}) ->
             {noreply, State#state{leader_node = LeaderNode, is_leader = true}};
         {legacy, _} ->
             {noreply, State#state{leader_node = LeaderNode, is_leader = false}};
-        {v1, LeaderNode} ->
+        {_, LeaderNode} ->
             gen_server:cast(?SERVER, node_watcher_update),
             {noreply, become_leader(State, LeaderNode)};
-        {v1, _} ->
+        {_, _} ->
             {noreply, become_proxy(State, LeaderNode)}
     end;
 
+handle_cast(node_watcher_update, State=#state{version = v2}) ->
+    {noreply, State};
 handle_cast(node_watcher_update, State=#state{active_nodes = OldActiveNodes, connections = C}) when State#state.is_leader == true ->
     NewActiveNodes = riak_core_node_watcher:nodes(riak_kv),
     DownNodes = OldActiveNodes -- NewActiveNodes,
@@ -302,8 +305,7 @@ handle_cast(Msg = {write_realtime_connections, Remote, Node, ConnectionList}, St
             {noreply, State#state{connections = NewConnections}};
 
         false ->
-            lager:info("data_mgr is proxy sending to leader -> ~p
-      node = ~p", [Msg, node()]),
+            lager:info("data_mgr is proxy sending to leader -> ~p node = ~p", [Msg, node()]),
             proxy_cast(Msg, State),
             {noreply, State}
     end;
@@ -342,18 +344,26 @@ handle_info(restore_leader_data, State) ->
 handle_info(reset_restoration_flag, State) ->
     {noreply, State#state{restoration = false}};
 
-handle_info(poll_core_capability, State=#state{version = OldVersion, core_capability_polling_interval = PI}) ->
-    NewVersion = riak_core_capability:get({riak_repl, realtime_connections}, legacy),
-    case {OldVersion, NewVersion} of
-        {legacy, legacy} ->
+handle_info(poll_core_capability, State=#state{version = OldVersion}) ->
+    CoreCapabilityPollingIntervalSecs = app_helper:get_env(riak_repl, realtime_core_capability_polling_interval, 60),
+    PI = CoreCapabilityPollingIntervalSecs * 1000,
+    case {OldVersion, riak_core_capability:get({riak_repl, realtime_connections}, legacy)} of
+        {_, v2} ->
+            case State#state.leader_node == node() of
+                true ->
+                    riak_core_ring_manager:ring_trans(fun riak_repl_ring:cleanup_realtime_connection_data/2, []);
+                false ->
+                    ok
+            end,
+            {noreply, State#state{version = v2}};
+        {v1, v1} ->
             erlang:send_after(PI, self(), poll_core_capability),
             {noreply, State};
         {legacy, v1} ->
-            %% send upgrade message to conn_mgr after 30 seconds to give the data manager time to get into the correct state
-            %% after changing versions
-            [erlang:send_after(30000, Pid, upgrade_connection_version) || {_Remote, Pid} <- riak_repl2_rtsource_conn_sup:enabled()],
-            {noreply, v1_init(State)};
-        {v1, _} ->
+            erlang:send_after(PI, self(), poll_core_capability),
+            {noreply, v1_init(State#state{version = v1})};
+        _ ->
+            erlang:send_after(PI, self(), poll_core_capability),
             {noreply, State}
     end;
 
@@ -365,8 +375,7 @@ handle_info(_Info, State) ->
 
 
 terminate(Reason, State) ->
-    lager:info("riak_repl2_rtsource_conn_data_mgr termianting due to: ~p
-            State: ~p", [Reason, State]),
+    lager:info("riak_repl2_rtsource_conn_data_mgr termianting due to: ~p State: ~p", [Reason, State]),
     ok.
 
 code_change(_OldVsn, State, _Extra) ->