Use coordinator to get notified on stream deletion

This commit is contained in:
Arnaud Cogoluègnes 2020-06-18 08:58:59 +02:00
parent 625eaec004
commit a1314a13c9
2 changed files with 121 additions and 45 deletions

View File

@ -21,7 +21,8 @@
%% API %% API
-export([init/1, handle_call/3, handle_cast/2, handle_info/2]). -export([init/1, handle_call/3, handle_cast/2, handle_info/2]).
-export([start_link/1, create/4, register/0, delete/3, lookup_leader/2, lookup_local_member/2, topology/2, unregister/0]). -export([start_link/1, create/4, delete/3, lookup_leader/2, lookup_local_member/2, topology/2]).
-export([subscribe/2, unsubscribe/2]).
-record(state, { -record(state, {
configuration, listeners, monitors configuration, listeners, monitors
@ -43,14 +44,6 @@ create(VirtualHost, Reference, Arguments, Username) ->
delete(VirtualHost, Reference, Username) -> delete(VirtualHost, Reference, Username) ->
gen_server:call(?MODULE, {delete, VirtualHost, Reference, Username}). gen_server:call(?MODULE, {delete, VirtualHost, Reference, Username}).
-spec register() -> ok.
register() ->
gen_server:call(?MODULE, {register, self()}).
-spec unregister() -> ok.
unregister() ->
gen_server:call(?MODULE, {unregister, self()}).
-spec lookup_leader(binary(), binary()) -> pid() | cluster_not_found. -spec lookup_leader(binary(), binary()) -> pid() | cluster_not_found.
lookup_leader(VirtualHost, Stream) -> lookup_leader(VirtualHost, Stream) ->
gen_server:call(?MODULE, {lookup_leader, VirtualHost, Stream}). gen_server:call(?MODULE, {lookup_leader, VirtualHost, Stream}).
@ -64,6 +57,14 @@ lookup_local_member(VirtualHost, Stream) ->
topology(VirtualHost, Stream) -> topology(VirtualHost, Stream) ->
gen_server:call(?MODULE, {topology, VirtualHost, Stream}). gen_server:call(?MODULE, {topology, VirtualHost, Stream}).
-spec subscribe(binary(), binary()) -> ok | {error, stream_not_found} | {error, internal_error}.
subscribe(VirtualHost, Stream) ->
gen_server:call(?MODULE, {subscribe, VirtualHost, Stream, self()}).
-spec unsubscribe(binary(), binary()) -> ok | {error, stream_not_found}.
unsubscribe(VirtualHost, Stream) ->
gen_server:call(?MODULE, {unsubscribe, VirtualHost, Stream, self()}).
stream_queue_arguments(Arguments) -> stream_queue_arguments(Arguments) ->
stream_queue_arguments([{<<"x-queue-type">>, longstr, <<"stream">>}], Arguments). stream_queue_arguments([{<<"x-queue-type">>, longstr, <<"stream">>}], Arguments).
@ -189,6 +190,44 @@ handle_call({topology, VirtualHost, Stream}, _From, State) ->
_ -> _ ->
{error, stream_not_found} {error, stream_not_found}
end, end,
{reply, Res, State};
handle_call({subscribe, VirtualHost, Stream, SubscriberPid}, _From, State) ->
Name = #resource{virtual_host = VirtualHost, kind = queue, name = Stream},
Res = case rabbit_amqqueue:lookup(Name) of
{ok, Q} ->
case is_stream_queue(Q) of
true ->
#{name := StreamInternalName} = amqqueue:get_type_state(Q),
case rabbit_stream_coordinator:subscribe(StreamInternalName, SubscriberPid) of
{ok, ok, _} ->
ok;
{ok, {error, not_found}, _} ->
{error, stream_not_found};
_ ->
{error, internal_error}
end;
_ ->
{error, stream_not_found}
end;
_ ->
{error, stream_not_found}
end,
{reply, Res, State};
handle_call({unsubscribe, VirtualHost, Stream, SubscriberPid}, _From, State) ->
Name = #resource{virtual_host = VirtualHost, kind = queue, name = Stream},
Res = case rabbit_amqqueue:lookup(Name) of
{ok, Q} ->
case is_stream_queue(Q) of
true ->
#{name := StreamInternalName} = amqqueue:get_type_state(Q),
rabbit_stream_coordinator:unsubscribe(StreamInternalName, SubscriberPid),
ok;
_ ->
{error, stream_not_found}
end;
_ ->
{error, stream_not_found}
end,
{reply, Res, State}. {reply, Res, State}.
handle_cast(_, State) -> handle_cast(_, State) ->

View File

@ -39,7 +39,7 @@
name :: string(), name :: string(),
helper_sup :: pid(), helper_sup :: pid(),
socket :: rabbit_net:socket(), socket :: rabbit_net:socket(),
clusters :: #{binary() => pid()}, stream_leaders :: #{binary() => pid()},
stream_subscriptions :: #{binary() => [integer()]}, stream_subscriptions :: #{binary() => [integer()]},
credits :: atomics:atomics_ref(), credits :: atomics:atomics_ref(),
authentication_state :: atom(), authentication_state :: atom(),
@ -48,7 +48,8 @@
connection_step :: atom(), % tcp_connected, peer_properties_exchanged, authenticating, authenticated, tuning, tuned, opened, failure, closing, closing_done connection_step :: atom(), % tcp_connected, peer_properties_exchanged, authenticating, authenticated, tuning, tuned, opened, failure, closing, closing_done
frame_max :: integer(), frame_max :: integer(),
heartbeater :: any(), heartbeater :: any(),
client_properties = #{} :: #{binary() => binary()} client_properties = #{} :: #{binary() => binary()},
subscribed_streams = sets:new() :: sets:set(binary())
}). }).
-record(configuration, { -record(configuration, {
@ -79,14 +80,13 @@ init([KeepaliveSup, Transport, Ref, #{initial_credits := InitialCredits,
RealSocket = rabbit_net:unwrap_socket(Sock), RealSocket = rabbit_net:unwrap_socket(Sock),
case rabbit_net:connection_string(Sock, inbound) of case rabbit_net:connection_string(Sock, inbound) of
{ok, ConnStr} -> {ok, ConnStr} ->
rabbit_stream_manager:register(),
Credits = atomics:new(1, [{signed, true}]), Credits = atomics:new(1, [{signed, true}]),
init_credit(Credits, InitialCredits), init_credit(Credits, InitialCredits),
Connection = #stream_connection{ Connection = #stream_connection{
name = ConnStr, name = ConnStr,
helper_sup = KeepaliveSup, helper_sup = KeepaliveSup,
socket = RealSocket, socket = RealSocket,
clusters = #{}, stream_leaders = #{},
stream_subscriptions = #{}, stream_subscriptions = #{},
credits = Credits, credits = Credits,
authentication_state = none, authentication_state = none,
@ -161,7 +161,8 @@ close(Transport, S) ->
Transport:close(S). Transport:close(S).
listen_loop_post_auth(Transport, #stream_connection{socket = S, listen_loop_post_auth(Transport, #stream_connection{socket = S,
stream_subscriptions = StreamSubscriptions, credits = Credits, heartbeater = Heartbeater} = Connection, stream_subscriptions = StreamSubscriptions, credits = Credits,
heartbeater = Heartbeater} = Connection,
#stream_connection_state{consumers = Consumers, blocked = Blocked} = State, #stream_connection_state{consumers = Consumers, blocked = Blocked} = State,
#configuration{credits_required_for_unblocking = CreditsRequiredForUnblocking} = Configuration) -> #configuration{credits_required_for_unblocking = CreditsRequiredForUnblocking} = Configuration) ->
{OK, Closed, Error} = Transport:messages(), {OK, Closed, Error} = Transport:messages(),
@ -197,13 +198,14 @@ listen_loop_post_auth(Transport, #stream_connection{socket = S,
end, end,
listen_loop_post_auth(Transport, Connection1, State2, Configuration) listen_loop_post_auth(Transport, Connection1, State2, Configuration)
end; end;
{stream_manager, cluster_deleted, ClusterReference} -> {queue_deleted, #resource{name = Stream}} ->
{Connection1, State1} = case clean_state_after_stream_deletion(ClusterReference, Connection, State) of C = unsubscribe_from_stream(Stream, Connection),
{Connection1, State1} = case clean_state_after_stream_deletion(Stream, C, State) of
{cleaned, NewConnection, NewState} -> {cleaned, NewConnection, NewState} ->
StreamSize = byte_size(ClusterReference), StreamSize = byte_size(Stream),
FrameSize = 2 + 2 + 2 + 2 + StreamSize, FrameSize = 2 + 2 + 2 + 2 + StreamSize,
Transport:send(S, [<<FrameSize:32, ?COMMAND_METADATA_UPDATE:16, ?VERSION_0:16, Transport:send(S, [<<FrameSize:32, ?COMMAND_METADATA_UPDATE:16, ?VERSION_0:16,
?RESPONSE_CODE_STREAM_DELETED:16, StreamSize:16, ClusterReference/binary>>]), ?RESPONSE_CODE_STREAM_DELETED:16, StreamSize:16, Stream/binary>>]),
{NewConnection, NewState}; {NewConnection, NewState};
{not_cleaned, SameConnection, SameState} -> {not_cleaned, SameConnection, SameState} ->
{SameConnection, SameState} {SameConnection, SameState}
@ -265,18 +267,18 @@ listen_loop_post_auth(Transport, #stream_connection{socket = S,
listen_loop_post_auth(Transport, Connection1, State1, Configuration); listen_loop_post_auth(Transport, Connection1, State1, Configuration);
{heartbeat_send_error, Reason} -> {heartbeat_send_error, Reason} ->
rabbit_log:info("Heartbeat send error ~p, closing connection~n", [Reason]), rabbit_log:info("Heartbeat send error ~p, closing connection~n", [Reason]),
rabbit_stream_manager:unregister(), S1 = unsubscribe_from_all_streams(S),
close(Transport, S); close(Transport, S1);
heartbeat_timeout -> heartbeat_timeout ->
rabbit_log:info("Heartbeat timeout, closing connection~n"), rabbit_log:info("Heartbeat timeout, closing connection~n"),
rabbit_stream_manager:unregister(), S1 = unsubscribe_from_all_streams(S),
close(Transport, S); close(Transport, S1);
{Closed, S} -> {Closed, S} ->
rabbit_stream_manager:unregister(), unsubscribe_from_all_streams(S),
rabbit_log:info("Socket ~w closed [~w]~n", [S, self()]), rabbit_log:info("Socket ~w closed [~w]~n", [S, self()]),
ok; ok;
{Error, S, Reason} -> {Error, S, Reason} ->
rabbit_stream_manager:unregister(), unsubscribe_from_all_streams(S),
rabbit_log:info("Socket error ~p [~w]~n", [Reason, S, self()]); rabbit_log:info("Socket error ~p [~w]~n", [Reason, S, self()]);
M -> M ->
rabbit_log:warning("Unknown message ~p~n", [M]), rabbit_log:warning("Unknown message ~p~n", [M]),
@ -487,7 +489,6 @@ handle_frame_pre_auth(Transport, #stream_connection{helper_sup = SupPid, socket
SupPid, Sock, ConnectionName, SupPid, Sock, ConnectionName,
Heartbeat, SendFun, Heartbeat, ReceiveFun), Heartbeat, SendFun, Heartbeat, ReceiveFun),
{Connection#stream_connection{connection_step = tuned, frame_max = FrameMax, heartbeater = Heartbeater}, State, Rest}; {Connection#stream_connection{connection_step = tuned, frame_max = FrameMax, heartbeater = Heartbeater}, State, Rest};
handle_frame_pre_auth(Transport, #stream_connection{user = User, socket = S} = Connection, State, handle_frame_pre_auth(Transport, #stream_connection{user = User, socket = S} = Connection, State,
<<?COMMAND_OPEN:16, ?VERSION_0:16, CorrelationId:32, <<?COMMAND_OPEN:16, ?VERSION_0:16, CorrelationId:32,
@ -571,7 +572,8 @@ handle_frame_post_auth(Transport, #stream_connection{socket = Socket,
credit = Credit, credit = Credit,
stream = Stream stream = Stream
}, },
rabbit_log:info("registering consumer ~p in ~p~n", [ConsumerState, self()]),
Connection1 = subscribe_to_stream(Stream, Connection),
response_ok(Transport, Connection, ?COMMAND_SUBSCRIBE, CorrelationId), response_ok(Transport, Connection, ?COMMAND_SUBSCRIBE, CorrelationId),
@ -588,10 +590,11 @@ handle_frame_post_auth(Transport, #stream_connection{socket = Socket,
_ -> _ ->
StreamSubscriptions#{Stream => [SubscriptionId]} StreamSubscriptions#{Stream => [SubscriptionId]}
end, end,
{Connection#stream_connection{stream_subscriptions = StreamSubscriptions1}, State#stream_connection_state{consumers = Consumers1}, Rest} {Connection1#stream_connection{stream_subscriptions = StreamSubscriptions1}, State#stream_connection_state{consumers = Consumers1}, Rest}
end end
end; end;
handle_frame_post_auth(Transport, #stream_connection{stream_subscriptions = StreamSubscriptions, clusters = Clusters} = Connection, handle_frame_post_auth(Transport, #stream_connection{stream_subscriptions = StreamSubscriptions,
stream_leaders = StreamLeaders} = Connection,
#stream_connection_state{consumers = Consumers} = State, #stream_connection_state{consumers = Consumers} = State,
<<?COMMAND_UNSUBSCRIBE:16, ?VERSION_0:16, CorrelationId:32, SubscriptionId:32>>, Rest) -> <<?COMMAND_UNSUBSCRIBE:16, ?VERSION_0:16, CorrelationId:32, SubscriptionId:32>>, Rest) ->
case subscription_exists(StreamSubscriptions, SubscriptionId) of case subscription_exists(StreamSubscriptions, SubscriptionId) of
@ -603,21 +606,27 @@ handle_frame_post_auth(Transport, #stream_connection{stream_subscriptions = Stre
Stream = Consumer#consumer.stream, Stream = Consumer#consumer.stream,
#{Stream := SubscriptionsForThisStream} = StreamSubscriptions, #{Stream := SubscriptionsForThisStream} = StreamSubscriptions,
SubscriptionsForThisStream1 = lists:delete(SubscriptionId, SubscriptionsForThisStream), SubscriptionsForThisStream1 = lists:delete(SubscriptionId, SubscriptionsForThisStream),
{StreamSubscriptions1, Clusters1} = {Connection1, StreamSubscriptions1, StreamLeaders1} =
case length(SubscriptionsForThisStream1) of case length(SubscriptionsForThisStream1) of
0 -> 0 ->
%% no more subscriptions for this stream %% no more subscriptions for this stream
{maps:remove(Stream, StreamSubscriptions), %% we unregister even though it could affect publishing if the stream is published to
maps:remove(Stream, Clusters) %% from this connection and is deleted.
%% to mitigate this, we remove the stream from the leaders cache
%% this way the stream leader will be looked up in the next publish command
%% and registered to.
C = unsubscribe_from_stream(Stream, Connection),
{C, maps:remove(Stream, StreamSubscriptions),
maps:remove(Stream, StreamLeaders)
}; };
_ -> _ ->
{StreamSubscriptions#{Stream => SubscriptionsForThisStream1}, Clusters} {Connection, StreamSubscriptions#{Stream => SubscriptionsForThisStream1}, StreamLeaders}
end, end,
Consumers1 = maps:remove(SubscriptionId, Consumers), Consumers1 = maps:remove(SubscriptionId, Consumers),
response_ok(Transport, Connection, ?COMMAND_SUBSCRIBE, CorrelationId), response_ok(Transport, Connection, ?COMMAND_SUBSCRIBE, CorrelationId),
{Connection#stream_connection{ {Connection1#stream_connection{
stream_subscriptions = StreamSubscriptions1, stream_subscriptions = StreamSubscriptions1,
clusters = Clusters1 stream_leaders = StreamLeaders1
}, State#stream_connection_state{consumers = Consumers1}, Rest} }, State#stream_connection_state{consumers = Consumers1}, Rest}
end; end;
handle_frame_post_auth(Transport, Connection, #stream_connection_state{consumers = Consumers} = State, handle_frame_post_auth(Transport, Connection, #stream_connection_state{consumers = Consumers} = State,
@ -663,8 +672,8 @@ handle_frame_post_auth(Transport, #stream_connection{socket = S, virtual_host =
case rabbit_stream_manager:delete(VirtualHost, Stream, Username) of case rabbit_stream_manager:delete(VirtualHost, Stream, Username) of
{ok, deleted} -> {ok, deleted} ->
response_ok(Transport, Connection, ?COMMAND_DELETE_STREAM, CorrelationId), response_ok(Transport, Connection, ?COMMAND_DELETE_STREAM, CorrelationId),
%% FIXME add connection and state parameters (and change return type) C = unsubscribe_from_stream(Stream, Connection),
{Connection1, State1} = case clean_state_after_stream_deletion(Stream, Connection, State) of {Connection1, State1} = case clean_state_after_stream_deletion(Stream, C, State) of
{cleaned, NewConnection, NewState} -> {cleaned, NewConnection, NewState} ->
StreamSize = byte_size(Stream), StreamSize = byte_size(Stream),
FrameSize = 2 + 2 + 2 + 2 + StreamSize, FrameSize = 2 + 2 + 2 + 2 + StreamSize,
@ -803,35 +812,63 @@ extract_stream_list(<<>>, Streams) ->
extract_stream_list(<<Length:16, Stream:Length/binary, Rest/binary>>, Streams) -> extract_stream_list(<<Length:16, Stream:Length/binary, Rest/binary>>, Streams) ->
extract_stream_list(Rest, [Stream | Streams]). extract_stream_list(Rest, [Stream | Streams]).
clean_state_after_stream_deletion(Stream, #stream_connection{clusters = Clusters, stream_subscriptions = StreamSubscriptions} = Connection, clean_state_after_stream_deletion(Stream, #stream_connection{stream_leaders = StreamLeaders, stream_subscriptions = StreamSubscriptions} = Connection,
#stream_connection_state{consumers = Consumers} = State) -> #stream_connection_state{consumers = Consumers} = State) ->
case maps:is_key(Stream, StreamSubscriptions) of case {maps:is_key(Stream, StreamSubscriptions), maps:is_key(Stream, StreamLeaders)} of
true -> {true, _} ->
#{Stream := SubscriptionIds} = StreamSubscriptions, #{Stream := SubscriptionIds} = StreamSubscriptions,
{cleaned, Connection#stream_connection{ {cleaned, Connection#stream_connection{
clusters = maps:remove(Stream, Clusters), stream_leaders = maps:remove(Stream, StreamLeaders),
stream_subscriptions = maps:remove(Stream, StreamSubscriptions) stream_subscriptions = maps:remove(Stream, StreamSubscriptions)
}, State#stream_connection_state{consumers = maps:without(SubscriptionIds, Consumers)}}; }, State#stream_connection_state{consumers = maps:without(SubscriptionIds, Consumers)}};
false -> {false, true} ->
{cleaned, Connection#stream_connection{
stream_leaders = maps:remove(Stream, StreamLeaders)
}, State};
{false, false} ->
{not_cleaned, Connection, State} {not_cleaned, Connection, State}
end. end.
lookup_leader(Stream, #stream_connection{clusters = Clusters, virtual_host = VirtualHost} = State) -> lookup_leader(Stream, #stream_connection{stream_leaders = StreamLeaders, virtual_host = VirtualHost} = Connection) ->
case maps:get(Stream, Clusters, undefined) of case maps:get(Stream, StreamLeaders, undefined) of
undefined -> undefined ->
case lookup_leader_from_manager(VirtualHost, Stream) of case lookup_leader_from_manager(VirtualHost, Stream) of
cluster_not_found -> cluster_not_found ->
cluster_not_found; cluster_not_found;
ClusterPid -> ClusterPid ->
{ClusterPid, State#stream_connection{clusters = Clusters#{Stream => ClusterPid}}} Connection1 = subscribe_to_stream(Stream, Connection),
{ClusterPid, Connection1#stream_connection{stream_leaders = StreamLeaders#{Stream => ClusterPid}}}
end; end;
ClusterPid -> ClusterPid ->
{ClusterPid, State} {ClusterPid, Connection}
end. end.
lookup_leader_from_manager(VirtualHost, Stream) -> lookup_leader_from_manager(VirtualHost, Stream) ->
rabbit_stream_manager:lookup_leader(VirtualHost, Stream). rabbit_stream_manager:lookup_leader(VirtualHost, Stream).
subscribe_to_stream(Stream, #stream_connection{virtual_host = VirtualHost, subscribed_streams = SubscribedStreams} = Connection) ->
case rabbit_stream_manager:subscribe(VirtualHost, Stream) of
ok ->
Connection#stream_connection{subscribed_streams = sets:add_element(Stream, SubscribedStreams)};
M ->
rabbit_log:warning("Could not subscribe to ~p stream: ~p~n", [Stream, M]),
Connection
end.
unsubscribe_from_stream(Stream, #stream_connection{virtual_host = VirtualHost, subscribed_streams = SubscribedStreams} = Connection) ->
case rabbit_stream_manager:unsubscribe(VirtualHost, Stream) of
ok ->
ok;
M ->
rabbit_log:warning("Could not unsubscribe from ~p stream: ~p~n", [Stream, M]),
Connection
end,
Connection#stream_connection{subscribed_streams = sets:del_element(Stream, SubscribedStreams)}.
unsubscribe_from_all_streams(#stream_connection{subscribed_streams = SubscribedStreams} = Connection) ->
[unsubscribe_from_stream(Stream, Connection) || Stream <- sets:to_list(SubscribedStreams)],
Connection#stream_connection{subscribed_streams = sets:new()}.
frame(Transport, #stream_connection{socket = S}, Frame) -> frame(Transport, #stream_connection{socket = S}, Frame) ->
FrameSize = byte_size(Frame), FrameSize = byte_size(Frame),
Transport:send(S, [<<FrameSize:32>>, Frame]). Transport:send(S, [<<FrameSize:32>>, Frame]).