diff --git a/apps/emqx/src/emqx_inflight.erl b/apps/emqx/src/emqx_inflight.erl index a3ff0ab10..96babd95a 100644 --- a/apps/emqx/src/emqx_inflight.erl +++ b/apps/emqx/src/emqx_inflight.erl @@ -28,6 +28,7 @@ update/3, resize/2, delete/2, + fold/3, values/1, to_list/1, to_list/2, @@ -77,6 +78,18 @@ delete(Key, ?INFLIGHT(MaxSize, Tree)) -> update(Key, Val, ?INFLIGHT(MaxSize, Tree)) -> ?INFLIGHT(MaxSize, gb_trees:update(Key, Val, Tree)). +-spec fold(fun((key(), Val :: term(), Acc) -> Acc), Acc, inflight()) -> Acc. +fold(FoldFun, AccIn, ?INFLIGHT(Tree)) -> + fold_iterator(FoldFun, AccIn, gb_trees:iterator(Tree)). + +fold_iterator(FoldFun, Acc, It) -> + case gb_trees:next(It) of + {Key, Val, ItNext} -> + fold_iterator(FoldFun, FoldFun(Key, Val, Acc), ItNext); + none -> + Acc + end. + -spec resize(integer(), inflight()) -> inflight(). resize(MaxSize, ?INFLIGHT(Tree)) -> ?INFLIGHT(MaxSize, Tree). diff --git a/apps/emqx/test/emqx_inflight_SUITE.erl b/apps/emqx/test/emqx_inflight_SUITE.erl index 2c0949b88..a56e62575 100644 --- a/apps/emqx/test/emqx_inflight_SUITE.erl +++ b/apps/emqx/test/emqx_inflight_SUITE.erl @@ -76,6 +76,17 @@ t_values(_) -> ?assertEqual([1, 2], emqx_inflight:values(Inflight)), ?assertEqual([{a, 1}, {b, 2}], emqx_inflight:to_list(Inflight)). +t_fold(_) -> + Inflight = maps:fold( + fun emqx_inflight:insert/3, + emqx_inflight:new(), + #{a => 1, b => 2, c => 42} + ), + ?assertEqual( + emqx_inflight:fold(fun(_, V, S) -> S + V end, 0, Inflight), + lists:foldl(fun({_, V}, S) -> S + V end, 0, emqx_inflight:to_list(Inflight)) + ). + t_is_full(_) -> Inflight = emqx_inflight:insert(k, v, emqx_inflight:new()), ?assertNot(emqx_inflight:is_full(Inflight)),