Extend aggregator requests and fix stall-after-error bug

src/async_pools/mongoose_aggregator_worker.erl

@@ -21,7 +21,7 @@
                     mongoose_async_pools:pool_extra()) ->
     {ok, mongoose_async_pools:task()} | {error, term()}.
 -callback request(mongoose_async_pools:task(), mongoose_async_pools:pool_extra()) ->
-    gen_server:request_id().
+    gen_server:request_id() | drop.
 -callback verify(term(), mongoose_async_pools:task(), mongoose_async_pools:pool_extra()) ->
     term().
 -optional_callbacks([verify/3]).
@@ -55,16 +55,16 @@
 -spec init(map()) -> {ok, state()}.
 init(#{host_type := HostType,
        pool_id := PoolId,
-       request_callback := Requester,
+       request_callback := Requestor,
        aggregate_callback := Aggregator,
        flush_extra := FlushExtra} = Opts)
-  when is_function(Requester, 2),
+  when is_function(Requestor, 2),
        is_function(Aggregator, 3),
        is_map(FlushExtra) ->
     ?LOG_DEBUG(#{what => aggregator_worker_start, host_type => HostType, pool_id => PoolId}),
     {ok, #state{host_type = HostType,
                 pool_id = PoolId,
-                request_callback = Requester,
+                request_callback = Requestor,
                 aggregate_callback = Aggregator,
                 verify_callback = maps:get(verify_callback, Opts, undefined),
                 flush_extra = FlushExtra}}.
@@ -96,10 +96,10 @@ handle_info(Msg, #state{async_request = {AsyncRequest, ReqTask}} = State) ->
     case gen_server:check_response(Msg, AsyncRequest) of
         {error, {Reason, _Ref}} ->
             ?LOG_ERROR(log_fields(State, #{what => asynchronous_request_failed, reason => Reason})),
-            {noreply, State};
+            {noreply, State#state{async_request = no_request_pending}};
         {reply, {error, Reason}} ->
             ?LOG_ERROR(log_fields(State, #{what => asynchronous_request_failed, reason => Reason})),
-            {noreply, State};
+            {noreply, State#state{async_request = no_request_pending}};
         {reply, Reply} ->
             maybe_verify_reply(Reply, ReqTask, State),
             {noreply, maybe_request_next(State)};
@@ -129,7 +129,7 @@ format_status(_Opt, [_PDict, State | _]) ->
 % If we don't have any request pending, it means that it is the first task submitted,
 % so aggregation is not needed.
 handle_task(_, Value, #state{async_request = no_request_pending} = State) ->
-    State#state{async_request = make_async_request(Value, State)};
+    make_async_request(Value, State);
 handle_task(Key, NewValue, #state{aggregate_callback = Aggregator,
                                   flush_elems = Acc,
                                   flush_queue = Queue,
@@ -153,7 +153,7 @@ handle_task(Key, NewValue, #state{aggregate_callback = Aggregator,
 % If we don't have any request pending, it means that it is the first task submitted,
 % so aggregation is not needed.
 handle_broadcast(Task, #state{async_request = no_request_pending} = State) ->
-    State#state{async_request = make_async_request(Task, State)};
+    make_async_request(Task, State);
 handle_broadcast(Task, #state{aggregate_callback = Aggregator,
                               flush_elems = Acc,
                               flush_extra = Extra} = State) ->
@@ -172,16 +172,26 @@ maybe_request_next(#state{flush_elems = Acc, flush_queue = Queue} = State) ->
     case queue:out(Queue) of
         {{value, Key}, NewQueue} ->
             {Value, NewAcc} = maps:take(Key, Acc),
-            State#state{async_request = make_async_request(Value, State),
-                        flush_elems = NewAcc, flush_queue = NewQueue};
+            NewState1 = State#state{flush_elems = NewAcc, flush_queue = NewQueue},
+            case make_async_request(Value, NewState1) of
+                NewState2 = #state{async_request = no_request_pending} ->
+                    maybe_request_next(NewState2);
+                NewState2 ->
+                    NewState2
+            end;
         {empty, _} ->
             State#state{async_request = no_request_pending}
     end.
 
-make_async_request(Value, #state{host_type = HostType, pool_id = PoolId,
-                                 request_callback = Requestor, flush_extra = Extra}) ->
-    mongoose_metrics:update(HostType, [mongoose_async_pools, PoolId, async_request], 1),
-    {Requestor(Value, Extra), Value}.
+make_async_request(Request, #state{host_type = HostType, pool_id = PoolId,
+                                   request_callback = Requestor, flush_extra = Extra} = State) ->
+    case Requestor(Request, Extra) of
+        drop ->
+            State;
+        ReqId ->
+            mongoose_metrics:update(HostType, [mongoose_async_pools, PoolId, async_request], 1),
+            State#state{async_request = {ReqId, Request}}
+    end.
 
 maybe_verify_reply(_, _, #state{verify_callback = undefined}) ->
     ok;

src/rdbms/mongoose_rdbms.erl

@@ -288,7 +288,7 @@ sql_transaction(HostType, F) when is_function(F) ->
     sql_call(HostType, {sql_transaction, F}).
 
 %% @doc SQL transaction based on a list of queries
--spec sql_transaction_request(server(), fun() | maybe_improper_list()) -> transaction_result().
+-spec sql_transaction_request(server(), fun() | maybe_improper_list()) -> gen_server:request_id().
 sql_transaction_request(HostType, Queries) when is_list(Queries) ->
     F = fun() -> lists:map(fun sql_query_t/1, Queries) end,
     sql_transaction_request(HostType, F);

test/batches_SUITE.erl

@@ -31,7 +31,8 @@ groups() ->
        prepare_task_works,
        sync_flushes_down_everything,
        sync_aggregates_down_everything,
-       aggregating_error_is_handled,
+       aggregating_error_is_handled_and_can_continue,
+       aggregation_might_produce_noop_requests,
        async_request
       ]}
     ].
@@ -109,13 +110,21 @@ shared_cache_inserts_in_shared_table(_) ->
     mongoose_user_cache:merge_entry(host_type(), ?mod(2), some_jid(), #{}),
     ?assert(mongoose_user_cache:is_member(host_type(), ?mod(1), some_jid())).
 
+aggregation_might_produce_noop_requests(_) ->
+    {ok, Server} = gen_server:start_link(?MODULE, [], []),
+    Requestor = fun(1, _) -> timer:sleep(1), gen_server:send_request(Server, 1);
+                   (_, _) -> drop end,
+    Opts = (default_aggregator_opts(Server))#{pool_id => ?FUNCTION_NAME,
+                                              request_callback => Requestor},
+    {ok, Pid} = gen_server:start_link(mongoose_aggregator_worker, Opts, []),
+    [ gen_server:cast(Pid, {task, key, N}) || N <- lists:seq(1, 1000) ],
+    async_helper:wait_until(
+      fun() -> gen_server:call(Server, get_acc) end, 1).
+
 broadcast_reaches_all_workers(_) ->
     {ok, Server} = gen_server:start_link(?MODULE, [], []),
-    WPoolOpts = #{pool_type => aggregate,
-                  pool_size => 10,
-                  request_callback => fun(Task, _) -> timer:sleep(1), gen_server:send_request(Server, Task) end,
-                  aggregate_callback => fun(T1, T2, _) -> {ok, T1 + T2} end,
-                  verify_callback => fun(ok, _T, _) -> ok end},
+    WPoolOpts = (default_aggregator_opts(Server))#{pool_type => aggregate,
+                                                   pool_size => 10},
     {ok, _} = mongoose_async_pools:start_pool(host_type(), ?FUNCTION_NAME, WPoolOpts),
     mongoose_async_pools:broadcast_task(host_type(), ?FUNCTION_NAME, key, 1),
     async_helper:wait_until(
@@ -134,11 +143,9 @@ broadcast_reaches_all_keys(_) ->
                           gen_server:send_request(Server, 0)
                   end
           end,
-    WPoolOpts = #{pool_type => aggregate,
-                  pool_size => 3,
-                  request_callback => Req,
-                  aggregate_callback => fun(T1, T2, _) -> {ok, T1 + T2} end,
-                  verify_callback => fun(ok, _T, _) -> ok end},
+    WPoolOpts = (default_aggregator_opts(Server))#{pool_type => aggregate,
+                                                   pool_size => 3,
+                                                   request_callback => Req},
     {ok, _} = mongoose_async_pools:start_pool(HostType, ?FUNCTION_NAME, WPoolOpts),
     [ mongoose_async_pools:put_task(HostType, ?FUNCTION_NAME, N, 1) || N <- lists:seq(0, 1000) ],
     mongoose_async_pools:broadcast(HostType, ?FUNCTION_NAME, -1),
@@ -227,39 +234,31 @@ sync_flushes_down_everything(_) ->
 
 sync_aggregates_down_everything(_) ->
     {ok, Server} = gen_server:start_link(?MODULE, [], []),
-    Opts = #{host_type => host_type(),
-             pool_id => ?FUNCTION_NAME,
-             request_callback => fun(Task, _) -> timer:sleep(1), gen_server:send_request(Server, Task) end,
-             aggregate_callback => fun(T1, T2, _) -> {ok, T1 + T2} end,
-             verify_callback => fun(ok, _T, _) -> ok end,
-             flush_extra => #{host_type => host_type()}},
+    Opts = (default_aggregator_opts(Server))#{pool_id => ?FUNCTION_NAME},
     {ok, Pid} = gen_server:start_link(mongoose_aggregator_worker, Opts, []),
     ?assertEqual(skipped, gen_server:call(Pid, sync)),
     [ gen_server:cast(Pid, {task, key, N}) || N <- lists:seq(1, 1000) ],
     ?assertEqual(ok, gen_server:call(Pid, sync)),
     ?assertEqual(500500, gen_server:call(Server, get_acc)).
 
-aggregating_error_is_handled(_) ->
+aggregating_error_is_handled_and_can_continue(_) ->
     {ok, Server} = gen_server:start_link(?MODULE, [], []),
-    Opts = #{host_type => host_type(),
-             pool_id => ?FUNCTION_NAME,
-             request_callback => fun(_, _) -> gen_server:send_request(Server, return_error) end,
-             aggregate_callback => fun(T1, T2, _) -> {ok, T1 + T2} end,
-             verify_callback => fun(ok, _T, _) -> ok end,
-             flush_extra => #{host_type => host_type()}},
+    Requestor = fun(Task, _) -> timer:sleep(1), gen_server:send_request(Server, Task) end,
+    Opts = (default_aggregator_opts(Server))#{pool_id => ?FUNCTION_NAME,
+                                              request_callback => Requestor},
     {ok, Pid} = gen_server:start_link(mongoose_aggregator_worker, Opts, []),
-    gen_server:cast(Pid, {task, key, 0}),
-    async_helper:wait_until(
-      fun() -> gen_server:call(Server, get_acc) end, 0).
+    [ gen_server:cast(Pid, {task, key, N}) || N <- lists:seq(1, 10) ],
+    gen_server:cast(Pid, {task, return_error, return_error}),
+    ct:sleep(100),
+    [ gen_server:cast(Pid, {task, key, N}) || N <- lists:seq(11, 100) ],
+    %% We don't call sync here because sync is force flushing,
+    %% we want to test that it flushes alone
+    ct:sleep(100),
+    ?assert(55 < gen_server:call(Server, get_acc)).
 
 async_request(_) ->
     {ok, Server} = gen_server:start_link(?MODULE, [], []),
-    Opts = #{host_type => host_type(),
-             pool_id => ?FUNCTION_NAME,
-             request_callback => fun(Task, _) -> timer:sleep(1), gen_server:send_request(Server, Task) end,
-             aggregate_callback => fun(T1, T2, _) -> {ok, T1 + T2} end,
-             verify_callback => fun(ok, _T, _) -> ok end,
-             flush_extra => #{host_type => host_type()}},
+    Opts = (default_aggregator_opts(Server))#{pool_id => ?FUNCTION_NAME},
     {ok, Pid} = gen_server:start_link(mongoose_aggregator_worker, Opts, []),
     [ gen_server:cast(Pid, {task, key, N}) || N <- lists:seq(1, 1000) ],
     async_helper:wait_until(
@@ -272,6 +271,26 @@ host_type() ->
 some_jid() ->
     jid:make_noprep(<<"alice">>, <<"localhost">>, <<>>).
 
+default_aggregator_opts(Server) ->
+    #{host_type => host_type(),
+      request_callback => requester(Server),
+      aggregate_callback => fun aggregate_sum/3,
+      verify_callback => fun validate_all_ok/3,
+      flush_extra => #{host_type => host_type()}}.
+
+validate_all_ok(ok, _, _) ->
+    ok.
+
+aggregate_sum(T1, T2, _) ->
+    {ok, T1 + T2}.
+
+requester(Server) ->
+    fun(return_error, _) ->
+            gen_server:send_request(Server, return_error);
+       (Task, _) ->
+            timer:sleep(1), gen_server:send_request(Server, Task)
+    end.
+
 init([]) ->
     {ok, 0}.