Merge pull request #10538 from rabbitmq/fix-offset-lag-calculation

michaelklishin · web-flow · commit e1a31e096739 · 2024-02-12T13:22:17.000-05:00
Fix consumer offset and offset lag calculation
diff --git a/deps/rabbitmq_stream/src/rabbit_stream_reader.erl b/deps/rabbitmq_stream/src/rabbit_stream_reader.erl
@@ -3908,16 +3908,21 @@ consumer_i(messages_consumed,
                          #consumer_configuration{counters = Counters}}) ->
     messages_consumed(Counters);
 consumer_i(offset,
-           #consumer{configuration =
-                         #consumer_configuration{counters = Counters}}) ->
-    consumer_offset(Counters);
+           #consumer{configuration = #consumer_configuration{counters = Counters},
+                     last_listener_offset = LLO}) ->
+    rabbit_stream_utils:consumer_offset(consumer_offset(Counters),
+                                        messages_consumed(Counters),
+                                        LLO);
 consumer_i(offset_lag, #consumer{log = undefined}) ->
     0;
 consumer_i(offset_lag,
-           #consumer{configuration =
-                         #consumer_configuration{counters = Counters},
+           #consumer{configuration = #consumer_configuration{counters = Counters},
+                     last_listener_offset = LLO,
                      log = Log}) ->
-    stream_stored_offset(Log) - consumer_offset(Counters);
+    rabbit_stream_utils:offset_lag(stream_stored_offset(Log),
+                                   consumer_offset(Counters),
+                                   messages_consumed(Counters),
+                                   LLO);
 consumer_i(connection_pid, _) ->
     self();
 consumer_i(node, _) ->
diff --git a/deps/rabbitmq_stream/src/rabbit_stream_utils.erl b/deps/rabbitmq_stream/src/rabbit_stream_utils.erl
@@ -35,7 +35,9 @@
          filter_spec/1,
          command_versions/0,
          filtering_supported/0,
-         check_super_stream_management_permitted/4]).
+         check_super_stream_management_permitted/4,
+         offset_lag/4,
+         consumer_offset/3]).
 
 -include_lib("rabbit_common/include/rabbit.hrl").
 -include_lib("rabbitmq_stream_common/include/rabbit_stream.hrl").
@@ -330,3 +332,26 @@ q(VirtualHost, Name) ->
 
 e(VirtualHost, Name) ->
     rabbit_misc:r(VirtualHost, exchange, Name).
+
+-spec consumer_offset(ConsumerOffsetFromCounter :: integer(),
+                      MessageConsumed :: non_neg_integer(),
+                      LastListenerOffset :: integer() | undefined) -> integer().
+consumer_offset(0, 0, undefined) ->
+    0;
+consumer_offset(0, 0, LastListenerOffset) when LastListenerOffset > 0 ->
+    %% consumer at "next" waiting for messages most likely
+    LastListenerOffset;
+consumer_offset(ConsumerOffsetFromCounter, _, _) ->
+    ConsumerOffsetFromCounter.
+
+-spec offset_lag(CommittedOffset :: integer(),
+                 ConsumerOffsetFromCounter :: integer(),
+                 MessageConsumed :: non_neg_integer(),
+                 LastListenerOffset :: integer() | undefined) -> integer().
+offset_lag(-1, _, _, _) ->
+    0;
+offset_lag(_, 0, 0, LastListenerOffset) when LastListenerOffset > 0 ->
+    %% consumer waiting for messages at the end of the stream, most likely
+    0;
+offset_lag(CommittedOffset, ConsumerOffset, _, _) ->
+    CommittedOffset - ConsumerOffset.
diff --git a/deps/rabbitmq_stream/test/rabbit_stream_SUITE.erl b/deps/rabbitmq_stream/test/rabbit_stream_SUITE.erl
@@ -60,7 +60,8 @@ groups() ->
        vhost_queue_limit,
        connection_should_be_closed_on_token_expiry,
        should_receive_metadata_update_after_update_secret,
-       store_offset_requires_read_access
+       store_offset_requires_read_access,
+       offset_lag_calculation
       ]},
      %% Run `test_global_counters` on its own so the global metrics are
      %% initialised to 0 for each testcase
@@ -814,6 +815,116 @@ store_offset_requires_read_access(Config) ->
     closed = wait_for_socket_close(T, S, 10),
     ok.
 
+offset_lag_calculation(Config) ->
+    FunctionName = atom_to_binary(?FUNCTION_NAME, utf8),
+    T = gen_tcp,
+    Port = get_port(T, Config),
+    Opts = get_opts(T),
+    {ok, S} = T:connect("localhost", Port, Opts),
+    C = rabbit_stream_core:init(0),
+    ConnectionName = FunctionName,
+    test_peer_properties(T, S, #{<<"connection_name">> => ConnectionName}, C),
+    test_authenticate(T, S, C),
+
+    Stream = FunctionName,
+    test_create_stream(T, S, Stream, C),
+
+    SubId = 1,
+    TheFuture = os:system_time(millisecond) + 60 * 60 * 1_000,
+    lists:foreach(fun(OffsetSpec) ->
+                          test_subscribe(T, S, SubId, Stream,
+                                         OffsetSpec, 10, #{},
+                                         ?RESPONSE_CODE_OK, C),
+                          ConsumerInfo = consumer_offset_info(Config, ConnectionName),
+                          ?assertEqual({0, 0}, ConsumerInfo),
+                          test_unsubscribe(T, S, SubId, C)
+                  end, [first, last, next, 0, 1_000, {timestamp, TheFuture}]),
+
+
+    PublisherId = 1,
+    test_declare_publisher(T, S, PublisherId, Stream, C),
+    MessageCount = 10,
+    Body = <<"hello">>,
+    lists:foreach(fun(_) ->
+                          test_publish_confirm(T, S, PublisherId, Body, C)
+                  end, lists:seq(1, MessageCount - 1)),
+    %% to make sure to have 2 chunks
+    timer:sleep(200),
+    test_publish_confirm(T, S, PublisherId, Body, C),
+    test_delete_publisher(T, S, PublisherId, C),
+
+    NextOffset = MessageCount,
+    lists:foreach(fun({OffsetSpec, ReceiveDeliver, CheckFun}) ->
+                          test_subscribe(T, S, SubId, Stream,
+                                         OffsetSpec, 1, #{},
+                                         ?RESPONSE_CODE_OK, C),
+                          case ReceiveDeliver of
+                              true ->
+                                  {{deliver, SubId, _}, _} = receive_commands(T, S, C);
+                              _ ->
+                                  ok
+                          end,
+                          {Offset, Lag} = consumer_offset_info(Config, ConnectionName),
+                          CheckFun(Offset, Lag),
+                          test_unsubscribe(T, S, SubId, C)
+                  end, [{first, true,
+                         fun(Offset, Lag) ->
+                                 ?assert(Offset >= 0, "first, at least one chunk consumed"),
+                                 ?assert(Lag > 0, "first, not all messages consumed")
+                         end},
+                        {last, true,
+                         fun(Offset, _Lag) ->
+                                 ?assert(Offset > 0, "offset expected for last")
+                         end},
+                        {next, false,
+                         fun(Offset, Lag) ->
+                                 ?assertEqual(NextOffset, Offset, "next, offset should be at the end of the stream"),
+                                 ?assert(Lag =:= 0, "next, offset lag should be 0")
+                         end},
+                        {0, true,
+                         fun(Offset, Lag) ->
+                                 ?assert(Offset >= 0, "offset spec = 0, at least one chunk consumed"),
+                                 ?assert(Lag > 0, "offset spec = 0, not all messages consumed")
+                         end},
+                        {1_000, false,
+                         fun(Offset, Lag) ->
+                                 ?assertEqual(NextOffset, Offset, "offset spec = 1000, offset should be at the end of the stream"),
+                                 ?assert(Lag =:= 0, "offset spec = 1000, offset lag should be 0")
+                         end},
+                        {{timestamp, TheFuture}, false,
+                         fun(Offset, Lag) ->
+                                 ?assertEqual(NextOffset, Offset, "offset spec in future, offset should be at the end of the stream"),
+                                 ?assert(Lag =:= 0, "offset spec in future , offset lag should be 0")
+                         end}]),
+
+    test_delete_stream(T, S, Stream, C, false),
+    test_close(T, S, C),
+
+    ok.
+
+consumer_offset_info(Config, ConnectionName) ->
+    [[{offset, Offset},
+      {offset_lag, Lag}]] = rpc(Config, 0, ?MODULE,
+                                list_consumer_info, [ConnectionName, [offset, offset_lag]]),
+    {Offset, Lag}.
+
+list_consumer_info(ConnectionName, Infos) ->
+    Pids = rabbit_stream:list(<<"/">>),
+    [ConnPid] = lists:filter(fun(ConnectionPid) ->
+                                     ConnectionPid ! {infos, self()},
+                                     receive
+                                         {ConnectionPid,
+                                          #{<<"connection_name">> := ConnectionName}} ->
+                                             true;
+                                         {ConnectionPid, _ClientProperties} ->
+                                             false
+                                     after 1000 ->
+                                               false
+                                     end
+                             end,
+                             Pids),
+    rabbit_stream_reader:consumers_info(ConnPid, Infos).
+
 store_offset(Transport, S, Reference, Stream, Offset) ->
     StoreFrame = rabbit_stream_core:frame({store_offset, Reference, Stream, Offset}),
     ok = Transport:send(S, StoreFrame).
@@ -983,7 +1094,10 @@ test_server(Transport, Stream, Config) ->
     ok.
 
 test_peer_properties(Transport, S, C0) ->
-    PeerPropertiesFrame = request({peer_properties, #{}}),
+    test_peer_properties(Transport, S, #{}, C0).
+
+test_peer_properties(Transport, S, Properties, C0) ->
+    PeerPropertiesFrame = request({peer_properties, Properties}),
     ok = Transport:send(S, PeerPropertiesFrame),
     {Cmd, C} = receive_commands(Transport, S, C0),
     ?assertMatch({response, 1, {peer_properties, ?RESPONSE_CODE_OK, _}},
@@ -1133,6 +1247,13 @@ test_publish_confirm(Transport, S, publish_v2 = PublishCmd, PublisherId,
     end,
     C.
 
+test_delete_publisher(Transport, Socket, PublisherId, C0) ->
+    Frame = request({delete_publisher, PublisherId}),
+    ok = Transport:send(Socket, Frame),
+    {Cmd, C} = receive_commands(Transport, Socket, C0),
+    ?assertMatch({response, 1, {delete_publisher, ?RESPONSE_CODE_OK}}, Cmd),
+    C.
+
 test_subscribe(Transport, S, SubscriptionId, Stream, C0) ->
     test_subscribe(Transport,
                    S,
@@ -1149,7 +1270,21 @@ test_subscribe(Transport,
                SubscriptionProperties,
                ExpectedResponseCode,
                C0) ->
-    SubscribeFrame = request({subscribe, SubscriptionId, Stream, 0, 10, SubscriptionProperties}),
+    test_subscribe(Transport, S, SubscriptionId, Stream, 0, 10,
+                   SubscriptionProperties,
+                   ExpectedResponseCode, C0).
+
+test_subscribe(Transport,
+               S,
+               SubscriptionId,
+               Stream,
+               OffsetSpec,
+               Credit,
+               SubscriptionProperties,
+               ExpectedResponseCode,
+               C0) ->
+    SubscribeFrame = request({subscribe, SubscriptionId, Stream,
+                              OffsetSpec, Credit, SubscriptionProperties}),
     ok = Transport:send(S, SubscribeFrame),
     {Cmd, C} = receive_commands(Transport, S, C0),
     ?assertMatch({response, 1, {subscribe, ExpectedResponseCode}}, Cmd),
