Introduce timeout for stream server-side requests

The stream plugin can send frames to a client connection
and expect a response from it. This is used currently
for the consumer_update frame (single active consumer feature).
There was no timeout mechanism so far, so a slow or blocked
application could prevent a group of consumers to move on.

This commit introduces a timeout mechanism: if the expected
response takes too long to arrive, the server assumes the
connection is blocked and closes it.

The default timeout is 60 seconds but it can be changed by setting
the request_timeout parameter of the rabbitmq_stream application.

Note the mechanism does not enforce the exact duration of the timeout,
as a timer is set for the first request and re-used for other requests.
With bad timing, a request can time out after twice as long
as the set-up timeout.

References #7743

(cherry picked from commit 763acc2)

Author: acogoluegnes

deps/rabbitmq_stream/src/rabbit_stream_reader.erl

@@ -47,6 +47,9 @@
          send_limit :: non_neg_integer(),
          log :: undefined | osiris_log:state(),
          last_listener_offset = undefined :: undefined | osiris:offset()}).
+-record(request,
+        {start :: integer(),
+         content :: term()}).
 -record(stream_connection_state,
         {data :: rabbit_stream_core:state(), blocked :: boolean(),
          consumers :: #{subscription_id() => #consumer{}}}).
@@ -89,8 +92,10 @@
          transport :: tcp | ssl,
          proxy_socket :: undefined | ranch_transport:socket(),
          correlation_id_sequence :: integer(),
-         outstanding_requests :: #{integer() => term()},
-         deliver_version :: rabbit_stream_core:command_version()}).
+         outstanding_requests :: #{integer() => #request{}},
+         deliver_version :: rabbit_stream_core:command_version(),
+         request_timeout :: pos_integer(),
+         outstanding_requests_timer :: undefined | timer:tref()}).
 -record(configuration,
         {initial_credits :: integer(),
          credits_required_for_unblocking :: integer(),
@@ -223,6 +228,8 @@ init([KeepaliveSup,
                 socket_op(Sock,
                           fun(S) -> rabbit_net:socket_ends(S, inbound) end),
             DeliverVersion = ?VERSION_1,
+            RequestTimeout = application:get_env(rabbitmq_stream,
+                                                 request_timeout, 60_000),
             Connection =
                 #stream_connection{name =
                                        rabbit_data_coercion:to_binary(ConnStr),
@@ -249,6 +256,7 @@ init([KeepaliveSup,
                                        rabbit_net:maybe_get_proxy_socket(Sock),
                                    correlation_id_sequence = 0,
                                    outstanding_requests = #{},
+                                   request_timeout = RequestTimeout,
                                    deliver_version = DeliverVersion},
             State =
                 #stream_connection_state{consumers = #{},
@@ -962,6 +970,37 @@ open(info, emit_stats,
          StatemData) ->
     Connection1 = emit_stats(Connection, State),
     {keep_state, StatemData#statem_data{connection = Connection1}};
+open(info, check_outstanding_requests,
+     #statem_data{connection = #stream_connection{outstanding_requests = Requests,
+                                                  request_timeout = Timeout} = Connection0} =
+         StatemData) ->
+    Time = erlang:monotonic_time(millisecond),
+    rabbit_log:debug("Checking outstanding requests at ~tp: ~tp", [Time, Requests]),
+    HasTimedOut = maps:fold(fun(_, #request{}, true) ->
+                                    true;
+                               (K, #request{content = Ctnt, start = Start}, false) ->
+                                    case (Time - Start) > Timeout of
+                                        true ->
+                                            rabbit_log:debug("Request ~tp with content ~tp has timed out",
+                                                             [K, Ctnt]),
+
+                                            true;
+                                        false ->
+                                            false
+                                    end
+                            end, false, Requests),
+    case HasTimedOut of
+        true ->
+            rabbit_log_connection:info("Forcing stream connection ~tp closing: request to client timed out",
+                                       [self()]),
+            _ = demonitor_all_streams(Connection0),
+            {stop, {request_timeout, <<"Request timeout">>}};
+        false ->
+            Connection1 = ensure_outstanding_requests_timer(
+                            Connection0#stream_connection{outstanding_requests_timer = undefined}
+                           ),
+            {keep_state, StatemData#statem_data{connection = Connection1}}
+    end;
 open(info, {shutdown, Explanation} = Reason,
      #statem_data{connection = Connection}) ->
     %% rabbitmq_management or rabbitmq_stream_management plugin
@@ -2531,7 +2570,7 @@ handle_frame_post_auth(Transport,
                             [RC])
     end,
     case maps:take(CorrelationId, Requests0) of
-        {#{subscription_id := SubscriptionId} = Msg, Rs} ->
+        {#request{content = #{subscription_id := SubscriptionId} = Msg}, Rs} ->
             Stream = stream_from_consumers(SubscriptionId, Consumers),
             rabbit_log:debug("Received consumer update response for subscription "
                              "~tp on stream ~tp, correlation ID ~tp",
@@ -2904,31 +2943,58 @@ maybe_register_consumer(VirtualHost,
     Active.
 
 maybe_send_consumer_update(Transport,
-                           Connection =
-                           #stream_connection{socket = S,
-                                              correlation_id_sequence = CorrIdSeq,
-                                              outstanding_requests = OutstandingRequests0},
+                           Connection = #stream_connection{
+                                           socket = S,
+                                           correlation_id_sequence = CorrIdSeq},
                            Consumer,
                            Active,
                            Msg) ->
-    #consumer{configuration = #consumer_configuration{
-                                            stream = Stream,
-                                            subscription_id = SubscriptionId
-                                           }} = Consumer,
-    rabbit_log:debug("SAC subscription ~tp on ~tp, active = ~tp, " ++
-                     "sending consumer update with correlation ID ~tp",
-                     [SubscriptionId, Stream, Active, CorrIdSeq]),
-    Frame =
-        rabbit_stream_core:frame({request, CorrIdSeq,
-                                  {consumer_update, SubscriptionId, Active}}),
+    #consumer{configuration =
+              #consumer_configuration{subscription_id = SubscriptionId}} = Consumer,
+    Frame = rabbit_stream_core:frame({request, CorrIdSeq,
+                                      {consumer_update, SubscriptionId, Active}}),
+
+    Connection1 = register_request(Connection, Msg),
 
-    OutstandingRequests1 =
-        maps:put(CorrIdSeq,
-                 Msg,
-                 OutstandingRequests0),
     send(Transport, S, Frame),
-    Connection#stream_connection{correlation_id_sequence = CorrIdSeq + 1,
-                                 outstanding_requests = OutstandingRequests1}.
+    Connection1.
+
+register_request(#stream_connection{outstanding_requests = Requests0,
+                                    correlation_id_sequence = CorrIdSeq} = C,
+                 RequestContent) ->
+    rabbit_log:debug("Registering RPC request ~tp with correlation ID ~tp",
+                     [RequestContent, CorrIdSeq]),
+
+    Requests1 = maps:put(CorrIdSeq, request(RequestContent), Requests0),
+
+    ensure_outstanding_requests_timer(
+      C#stream_connection{correlation_id_sequence = CorrIdSeq + 1,
+                          outstanding_requests = Requests1}).
+
+request(Content) ->
+    #request{start = erlang:monotonic_time(millisecond),
+             content = Content}.
+
+ensure_outstanding_requests_timer(#stream_connection{
+                                     outstanding_requests = Requests,
+                                     outstanding_requests_timer = undefined
+                                    } = C) when map_size(Requests) =:= 0 ->
+    C;
+ensure_outstanding_requests_timer(#stream_connection{
+                                     outstanding_requests = Requests,
+                                     outstanding_requests_timer = TRef
+                                    } = C) when map_size(Requests) =:= 0 ->
+    erlang:cancel_timer(TRef, [{async, true}, {info, false}]),
+    C#stream_connection{outstanding_requests_timer = undefined};
+ensure_outstanding_requests_timer(#stream_connection{
+                                     outstanding_requests = Requests,
+                                     outstanding_requests_timer = undefined,
+                                     request_timeout = Timeout
+                                    } = C) when map_size(Requests) > 0 ->
+    TRef = erlang:send_after(Timeout, self(), check_outstanding_requests),
+    C#stream_connection{outstanding_requests_timer = TRef};
+ensure_outstanding_requests_timer(C) ->
+    C.
 
 maybe_unregister_consumer(_, _, false = _Sac, Requests) ->
     Requests;
@@ -2945,9 +3011,10 @@ maybe_unregister_consumer(VirtualHost,
     ConsumerName = consumer_name(Properties),
 
     Requests1 = maps:fold(
-                  fun(_, #{active := false,
-                           subscription_id := SubId,
-                           stepping_down := true}, Acc) when SubId =:= SubscriptionId ->
+                  fun(_, #request{content =
+                                  #{active := false,
+                                    subscription_id := SubId,
+                                    stepping_down := true}}, Acc) when SubId =:= SubscriptionId ->
                           _ = rabbit_stream_sac_coordinator:activate_consumer(VirtualHost,
                                                                               Stream,
                                                                               ConsumerName),

deps/rabbitmq_stream/test/rabbit_stream_SUITE.erl

@@ -44,7 +44,8 @@ groups() ->
        unauthenticated_client_rejected_authenticating,
        timeout_authenticating,
        timeout_close_sent,
-       max_segment_size_bytes_validation]},
+       max_segment_size_bytes_validation,
+       close_connection_on_consumer_update_timeout]},
      %% Run `test_global_counters` on its own so the global metrics are
      %% initialised to 0 for each testcase
      {single_node_1, [], [test_global_counters]},
@@ -110,6 +111,26 @@ end_per_group(_, Config) ->
     rabbit_ct_helpers:run_steps(Config,
                                 rabbit_ct_broker_helpers:teardown_steps()).
 
+init_per_testcase(close_connection_on_consumer_update_timeout = TestCase, Config) ->
+    ok = rabbit_ct_broker_helpers:rpc(Config,
+                                      0,
+                                      application,
+                                      set_env,
+                                      [rabbitmq_stream, request_timeout, 2000]),
+    rabbit_ct_helpers:testcase_started(Config, TestCase);
+init_per_testcase(TestCase, Config) ->
+    rabbit_ct_helpers:testcase_started(Config, TestCase).
+
+end_per_testcase(close_connection_on_consumer_update_timeout = TestCase, Config) ->
+    ok = rabbit_ct_broker_helpers:rpc(Config,
+                                      0,
+                                      application,
+                                      set_env,
+                                      [rabbitmq_stream, request_timeout, 60000]),
+    rabbit_ct_helpers:testcase_finished(Config, TestCase);
+end_per_testcase(TestCase, Config) ->
+    rabbit_ct_helpers:testcase_finished(Config, TestCase).
+
 test_global_counters(Config) ->
     Stream = atom_to_binary(?FUNCTION_NAME, utf8),
     test_server(gen_tcp, Stream, Config),
@@ -350,6 +371,37 @@ max_segment_size_bytes_validation(Config) ->
     test_close(Transport, S, C3),
     ok.
 
+close_connection_on_consumer_update_timeout(Config) ->
+    Transport = gen_tcp,
+    Port = get_stream_port(Config),
+    {ok, S} =
+    Transport:connect("localhost", Port,
+                      [{active, false}, {mode, binary}]),
+    C0 = rabbit_stream_core:init(0),
+    C1 = test_peer_properties(Transport, S, C0),
+    C2 = test_authenticate(Transport, S, C1),
+    Stream = atom_to_binary(?FUNCTION_NAME, utf8),
+    C3 = test_create_stream(Transport, S, Stream, C2),
+
+    SubId = 42,
+    C4 = test_subscribe(Transport, S, SubId, Stream,
+                        #{<<"single-active-consumer">> => <<"true">>,
+                          <<"name">> => <<"foo">>},
+                        C3),
+    {Cmd, _C5} = receive_commands(Transport, S, C4),
+    ?assertMatch({request, _, {consumer_update, SubId, true}}, Cmd),
+    closed = wait_for_socket_close(Transport, S, 10),
+    {ok, Sb} =
+    Transport:connect("localhost", Port,
+                      [{active, false}, {mode, binary}]),
+    Cb0 = rabbit_stream_core:init(0),
+    Cb1 = test_peer_properties(Transport, Sb, Cb0),
+    Cb2 = test_authenticate(Transport, Sb, Cb1),
+    Cb3 = test_delete_stream(Transport, Sb, Stream, Cb2, false),
+    _Cb4 = test_close(Transport, Sb, Cb3),
+    closed = wait_for_socket_close(Transport, Sb, 10),
+    ok.
+
 consumer_count(Config) ->
     ets_count(Config, ?TABLE_CONSUMER).
 
@@ -525,12 +577,21 @@ test_create_stream(Transport, S, Stream, C0) ->
     C.
 
 test_delete_stream(Transport, S, Stream, C0) ->
+    test_delete_stream(Transport, S, Stream, C0, true).
+
+test_delete_stream(Transport, S, Stream, C0, false) ->
+    do_test_delete_stream(Transport, S, Stream, C0);
+test_delete_stream(Transport, S, Stream, C0, true) ->
+    C1 = do_test_delete_stream(Transport, S, Stream, C0),
+    test_metadata_update_stream_deleted(Transport, S, Stream, C1).
+
+do_test_delete_stream(Transport, S, Stream, C0) ->
     DeleteStreamFrame =
         rabbit_stream_core:frame({request, 1, {delete_stream, Stream}}),
     ok = Transport:send(S, DeleteStreamFrame),
     {Cmd, C1} = receive_commands(Transport, S, C0),
     ?assertMatch({response, 1, {delete_stream, ?RESPONSE_CODE_OK}}, Cmd),
-    test_metadata_update_stream_deleted(Transport, S, Stream, C1).
+    C1.
 
 test_metadata_update_stream_deleted(Transport, S, Stream, C0) ->
     {Meta, C1} = receive_commands(Transport, S, C0),