Jelajahi Sumber

feat(emqx): add `emqx_inflight:fold/3` generic function

Andrew Mayorov 2 tahun lalu
induk
melakukan
f022c9b1a4
2 mengubah file dengan 24 tambahan dan 0 penghapusan
  1. 13 0
      apps/emqx/src/emqx_inflight.erl
  2. 11 0
      apps/emqx/test/emqx_inflight_SUITE.erl

+ 13 - 0
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).

+ 11 - 0
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)),