diff options
author | Simon MacMullen <simon@rabbitmq.com> | 2013-02-21 12:00:16 +0000 |
---|---|---|
committer | Simon MacMullen <simon@rabbitmq.com> | 2013-02-21 12:00:16 +0000 |
commit | ef744142e180f95adbeac63c7ef6628ec196793f (patch) | |
tree | b313e47ec5fe51763b173ee56122a878124109b0 | |
parent | 97757bbfc78bac5f02b5d7d9b1a4630cb41852f7 (diff) | |
download | rabbitmq-server-ef744142e180f95adbeac63c7ef6628ec196793f.tar.gz |
Use gb_trees rather than dict for performance.
-rw-r--r-- | src/rabbit_limiter.erl | 36 |
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 |