Use gb_trees rather than dict for performance.

1 files changed, 21 insertions, 15 deletions

diff --git a/src/rabbit_limiter.erl b/src/rabbit_limiter.erl
index ece3d1a4..801b748e 100644
--- a/src/rabbit_limiter.erl
+++ b/src/rabbit_limiter.erl
@@ -86,7 +86,7 @@ start_link() -> gen_server2:start_link(?MODULE, [], []).
 make_token() -> make_token(undefined).
 make_token(Pid) -> #token{pid     = Pid,
                           enabled = false,
-                          credits = dict:new()}.
+                          credits = gb_trees:empty()}.
 
 is_enabled(#token{enabled = Enabled}) -> Enabled.
 
@@ -141,10 +141,10 @@ unblock(Limiter) ->
     maybe_call(Limiter, {unblock, Limiter}, ok).
 
 is_consumer_blocked(#token{credits = Credits}, CTag) ->
-    case dict:find(CTag, Credits) of
-        {ok, #credit{credit = C}} when C > 0 -> false;
-        {ok, #credit{}}                      -> true;
-        error                                -> false
+    case gb_trees:lookup(CTag, Credits) of
+        {value, #credit{credit = C}} when C > 0 -> false;
+        {value, #credit{}}                      -> true;
+        none                                    -> false
     end.
 
 is_blocked(Limiter) ->
@@ -155,16 +155,16 @@ credit(Limiter = #token{credits = Credits}, CTag, Credit, Drain) ->
 
 drained(Limiter = #token{credits = Credits}) ->
     {CTagCredits, Credits2} =
-        dict:fold(
-          fun (CTag,  #credit{credit = C,  drain = true},  {Acc, Creds0}) ->
+        lists:foldl(
+          fun ({CTag,  #credit{credit = C,  drain = true}},  {Acc, Creds0}) ->
                   {[{CTag, C} | Acc], write_credit(CTag, 0, false, Creds0)};
-              (_CTag, #credit{credit = _C, drain = false}, {Acc, Creds0}) ->
+              ({_CTag, #credit{credit = _C, drain = false}}, {Acc, Creds0}) ->
                   {Acc, Creds0}
-          end, {[], Credits}, Credits),
+          end, {[], Credits}, gb_trees:to_list(Credits)),
     {CTagCredits, Limiter#token{credits = Credits2}}.
 
 forget_consumer(Limiter = #token{credits = Credits}, CTag) ->
-    Limiter#token{credits = dict:erase(CTag, Credits)}.
+    Limiter#token{credits = gb_trees:delete_any(CTag, Credits)}.
 
 copy_queue_state(#token{credits = Credits}, Token) ->
     Token#token{credits = Credits}.
@@ -180,10 +180,10 @@ copy_queue_state(#token{credits = Credits}, Token) ->
 %% maintain a fiction that the limiter is making the decisions...
 
 record_send_q(CTag, Credits) ->
-    case dict:find(CTag, Credits) of
-        {ok, #credit{credit = Credit, drain = Drain}} ->
+    case gb_trees:lookup(CTag, Credits) of
+        {value, #credit{credit = Credit, drain = Drain}} ->
             write_credit(CTag, Credit, Drain, Credits);
-        error ->
+        none ->
             Credits
     end.
 
@@ -195,10 +195,16 @@ update_credit(CTag, Credit, Drain, Credits) ->
     end.
 
 write_credit(CTag, Credit, Drain, Credits) when Credit > 0 ->
-    dict:store(CTag, #credit{credit = Credit, drain = Drain}, Credits);
+    write_credit0(CTag, #credit{credit = Credit, drain = Drain}, Credits);
 %% Using up all credit means we do not need to send a drained event
 write_credit(CTag, Credit, _Drain, Credits) ->
-    dict:store(CTag, #credit{credit = Credit, drain = false}, Credits).
+    write_credit0(CTag, #credit{credit = Credit, drain = false}, Credits).
+
+write_credit0(CTag, Credit, Credits) ->
+    case gb_trees:is_defined(CTag, Credits) of
+        true  -> gb_trees:update(CTag, Credit, Credits);
+        false -> gb_trees:insert(CTag, Credit, Credits)
+    end.
 
 %%----------------------------------------------------------------------------
 %% gen_server callbacks