12741272085

Committed 13 Jan 2025 05:33AM UTC coverage: 82.355%. First build

Build # 12741272085

Build Type

Pull #14475

github

Committed by

web-flow

Commit Message

Merge b56877530 into 0fc8025be

Pull Request Pull Request #14475: refactor(limiter): refactor and simplify the limiter

Run Details

199 of 249 new or added lines in 21 files covered. (79.92%)

57201 of 69457 relevant lines covered (82.35%)

15186.26 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

66.22

/apps/emqx/src/emqx_limiter/src/emqx_limiter_allocator.erl

%%--------------------------------------------------------------------
%% Copyright (c) 2021-2025 EMQ Technologies Co., Ltd. All Rights Reserved.
%%
%% Licensed under the Apache License, Version 2.0 (the "License");
%% you may not use this file except in compliance with the License.
%% You may obtain a copy of the License at
%%
%%     http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing, software
%% distributed under the License is distributed on an "AS IS" BASIS,
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%% See the License for the specific language governing permissions and
%% limitations under the License.
%%--------------------------------------------------------------------

%% Handle limiters for a zone.
%% Currently, each allocator corresponds to one zone,
%% if there are thousands of limiters, we can change to using allocator pools

-module(emqx_limiter_allocator).

-behaviour(gen_server).

-include_lib("emqx/include/logger.hrl").

%% gen_server callbacks
-export([
    init/1,
    handle_call/3,
    handle_cast/2,
    handle_info/2,
    terminate/2
]).

-export([
    add_bucket/2, add_bucket/3,
    delete_bucket/1, delete_bucket/2
]).

-export([start_link/1]).

-type bucket() :: #{
    name := bucket_name(),
    rate := rate(),
    burst => rate(),
    capacity := capacity(),
    counter := counters:counters_ref(),
    index := index(),
    correction := float()
}.

-type allocator_name() :: emqx_limiter:zone() | binary().

-type state() :: #{
    name := allocator_name(),
    buckets := buckets(),
    counter := counters:counters_ref(),
    index := index(),
    alloc_interval := millisecond(),
    lasttime := millisecond(),
    free_index := [index()]
}.

-type buckets() :: #{bucket_name() => bucket()}.
-type bucket_name() :: atom().
-type rate() :: number().
-type millisecond() :: non_neg_integer().
-type capacity() :: number().
-type index() :: pos_integer().

-define(COUNTER_SIZE, 8).
-define(NOW, erlang:system_time(millisecond)).

-export_type([allocator_name/0, state/0, index/0]).

%%--------------------------------------------------------------------
%% API
%%--------------------------------------------------------------------

-spec start_link(allocator_name()) -> _.
start_link(Name) ->
    gen_server:start_link({via, emqx_limiter_manager, Name}, ?MODULE, [Name], []).

add_bucket(Name, Cfg) ->
    add_bucket(emqx_limiter:internal_allocator(), Name, Cfg).

add_bucket(AllocatorName, Name, Cfg) ->
    gen_server:call({via, emqx_limiter_manager, AllocatorName}, {?FUNCTION_NAME, Name, Cfg}).

delete_bucket(Name) ->
    delete_bucket(emqx_limiter:internal_allocator(), Name).

delete_bucket(AllocatorName, Name) ->
    gen_server:call({via, emqx_limiter_manager, AllocatorName}, {?FUNCTION_NAME, Name}).

%%--------------------------------------------------------------------
%%% gen_server callbacks
%%--------------------------------------------------------------------

-spec init([allocator_name()]) -> {ok, State :: state()}.
init([Name]) ->
    State = init_state(Name),
    tick_alloc_event(State),
    {ok, State}.

handle_call(
    {add_bucket, Name, #{rate := Rate, burst := Burst}},
    _From,
    #{
        name := AllocatorName,
        counter := Counter,
        index := Index,
        alloc_interval := Interval,
        buckets := Buckets,
        free_index := Free
    } = State
) ->
    case Free of
        [ToUse | Free2] ->
            NewIndex = Index;
        [] ->
            ToUse = Index + 1,
            NewIndex = ToUse,
            Free2 = Free
    end,

    case ToUse > ?COUNTER_SIZE of
        false ->
            Bucket = do_create_bucket(
                Name, Rate, Burst, AllocatorName, Counter, Interval, ToUse
            ),

            {reply, ok, State#{
                buckets := Buckets#{Name => Bucket},
                index := NewIndex,
                free_index := Free2
            }};
        _ ->
            {reply, {error, <<"exceeded_max_index">>}, State}
    end;
handle_call(
    {delete_bucket, Name},
    _From,
    #{name := AllocatorName, buckets := Buckets, free_index := Free} = State
) ->
    case maps:take(Name, Buckets) of
        error ->
            {reply, ok, State};
        {#{index := Index}, Buckets2} ->
            emqx_limiter_manager:delete_bucket(AllocatorName, Name),
            {reply, ok, State#{buckets := Buckets2, free_index := [Index | Free]}}
    end;
handle_call(Req, _From, State) ->
    ?SLOG(error, #{msg => "emqx_limiter_allocator_unexpected_call", call => Req}),
    {reply, ignored, State}.

handle_cast(Req, State) ->
    ?SLOG(error, #{msg => "emqx_limiter_allocator_unexpected_cast", cast => Req}),
    {noreply, State}.

handle_info(tick_alloc_event, State) ->
    {noreply, do_alloc(State)};
handle_info(Info, State) ->
    ?SLOG(error, #{msg => "emqx_limiter_allocator_unexpected_info", info => Info}),
    {noreply, State}.

terminate(_Reason, #{name := Name, buckets := Buckets} = _State) ->
    maps:foreach(
        fun(LimiterName, _) ->
            emqx_limiter_manager:delete_bucket(Name, LimiterName)
        end,
        Buckets
    ),
    ok.

%%--------------------------------------------------------------------
%%% Internal functions
%%--------------------------------------------------------------------
tick_alloc_event(#{alloc_interval := Interval}) ->
    erlang:send_after(Interval, self(), ?FUNCTION_NAME).

%% an allocator server with a zone name is used for this zone, or it is a dynamic server
init_state(Name) when is_binary(Name) ->
    Counter = counters:new(?COUNTER_SIZE, [write_concurrency]),
    #{
        name => Name,
        counter => Counter,
        buckets => #{},
        index => 0,
        alloc_interval => emqx_limiter:default_alloc_interval(),
        lasttime => ?NOW,
        free_index => []
    };
init_state(Zone) when is_atom(Zone) ->
    Cfg = emqx_config:get_zone_conf(Zone, [mqtt, limiter]),
    init_state(Zone, Cfg).

init_state(Zone, #{alloc_interval := Interval} = Cfg) ->
    Counter = counters:new(?COUNTER_SIZE, [write_concurrency]),
    Names = emqx_limiter_schema:mqtt_limiter_names(),
    Buckets = init_buckets(Names, Zone, Counter, Cfg, #{}),
    #{
        name => Zone,
        counter => Counter,
        buckets => Buckets,
        index => maps:size(Buckets),
        alloc_interval => Interval,
        lasttime => ?NOW,
        free_index => []
    }.

init_buckets([Name | Names], Zone, Counter, #{alloc_interval := Interval} = Cfg, Buckets) ->
    {ok, RateKey} = emqx_limiter:to_rate_key(Name),
    {ok, BurstKey} = emqx_limiter:to_burst_key(Name),
    case maps:get(RateKey, Cfg, infinity) of
        infinity ->
            init_buckets(Names, Zone, Counter, Cfg, Buckets);
        Rate ->
            Burst = maps:get(BurstKey, Cfg, 0),
            Bucket = do_create_bucket(
                Name, Rate, Burst, Zone, Counter, Interval, maps:size(Buckets) + 1
            ),
            init_buckets(Names, Zone, Counter, Cfg, Buckets#{Name => Bucket})
    end;
init_buckets([], _Zone, _Counter, _, Buckets) ->
    Buckets.

%% @doc generate tokens, and then spread to leaf nodes
-spec do_alloc(state()) -> state().
do_alloc(
    #{
        lasttime := LastTime,
        buckets := Buckets
    } = State
) ->
    tick_alloc_event(State),
    Now = ?NOW,
    Elapsed = Now - LastTime,
    Buckets2 = do_buckets_alloc(Buckets, Elapsed),
    State#{
        lasttime := Now,
        buckets := Buckets2
    }.

do_buckets_alloc(Buckets, Elapsed) ->
    maps:map(
        fun(_, Bucket) ->
            do_bucket_alloc(Bucket, Elapsed)
        end,
        Buckets
    ).

do_bucket_alloc(
    #{
        rate := Rate,
        correction := Correction,
        counter := Counter,
        index := Index,
        capacity := Capacity
    } = Bucket,
    Elapsed
) ->
    Val = counters:get(Counter, Index),
    case Val >= Capacity of
        true ->
            Bucket;
        _ ->
            Inc = Rate * Elapsed + Correction,
            Inc2 = erlang:floor(Inc),
            Correction2 = Inc - Inc2,
            add_tokens(Bucket, Inc2),
            Bucket#{correction := Correction2}
    end.

set_tokens(Counter, Ix, Tokens) ->
    counters:put(Counter, Ix, erlang:floor(Tokens)).

add_tokens(_, 0) ->
    ok;
add_tokens(#{counter := Counter, index := Index, capacity := Capacity}, Tokens) ->
    Val = counters:get(Counter, Index),
    case erlang:min(Capacity, Val + Tokens) - Val of
        Inc when Inc > 0 ->
            counters:put(Counter, Index, Inc + Val);
        _ ->
            ok
    end.

do_create_bucket(
    LimiterName,
    Rate,
    Burst,
    AllocatorName,
    Counter,
    Interval,
    Index
) ->
    Capacity = emqx_limiter:calc_capacity(Rate, Interval),

    set_tokens(Counter, Index, Capacity),
    Ref = emqx_limiter_bucket_ref:new(Counter, Index),
    emqx_limiter_manager:insert_bucket(AllocatorName, LimiterName, Ref),
    #{
        name => LimiterName,
        rate => Rate,
        burst => Burst,
        capacity => Capacity,
        counter => Counter,
        index => Index,
        correction => 0
    }.

1	%%--------------------------------------------------------------------
2	%% Copyright (c) 2021-2025 EMQ Technologies Co., Ltd. All Rights Reserved.
3	%%
4	%% Licensed under the Apache License, Version 2.0 (the "License");
5	%% you may not use this file except in compliance with the License.
6	%% You may obtain a copy of the License at
7	%%
8	%% http://www.apache.org/licenses/LICENSE-2.0
9	%%
10	%% Unless required by applicable law or agreed to in writing, software
11	%% distributed under the License is distributed on an "AS IS" BASIS,
12	%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	%% See the License for the specific language governing permissions and
14	%% limitations under the License.
15	%%--------------------------------------------------------------------
16
17	%% Handle limiters for a zone.
18	%% Currently, each allocator corresponds to one zone,
19	%% if there are thousands of limiters, we can change to using allocator pools
20
21	-module(emqx_limiter_allocator).
22
23	-behaviour(gen_server).
24
25	-include_lib("emqx/include/logger.hrl").
26
27	%% gen_server callbacks
28	-export([
29	init/1,
30	handle_call/3,
31	handle_cast/2,
32	handle_info/2,
33	terminate/2
34	]).
35
36	-export([
37	add_bucket/2, add_bucket/3,
38	delete_bucket/1, delete_bucket/2
39	]).
40
41	-export([start_link/1]).
42
43	-type bucket() :: #{
44	name := bucket_name(),
45	rate := rate(),
46	burst => rate(),
47	capacity := capacity(),
48	counter := counters:counters_ref(),
49	index := index(),
50	correction := float()
51	}.
52
53	-type allocator_name() :: emqx_limiter:zone() \| binary().
54
55	-type state() :: #{
56	name := allocator_name(),
57	buckets := buckets(),
58	counter := counters:counters_ref(),
59	index := index(),
60	alloc_interval := millisecond(),
61	lasttime := millisecond(),
62	free_index := [index()]
63	}.
64
65	-type buckets() :: #{bucket_name() => bucket()}.
66	-type bucket_name() :: atom().
67	-type rate() :: number().
68	-type millisecond() :: non_neg_integer().
69	-type capacity() :: number().
70	-type index() :: pos_integer().
71
72	-define(COUNTER_SIZE, 8).
73	-define(NOW, erlang:system_time(millisecond)).
74
75	-export_type([allocator_name/0, state/0, index/0]).
76
77	%%--------------------------------------------------------------------
78	%% API
79	%%--------------------------------------------------------------------
80
81	-spec start_link(allocator_name()) -> _.
82	start_link(Name) ->
83	gen_server:start_link({via, emqx_limiter_manager, Name}, ?MODULE, [Name], []).	1,179✔
84
85	add_bucket(Name, Cfg) ->
86	add_bucket(emqx_limiter:internal_allocator(), Name, Cfg).	191✔
87
88	add_bucket(AllocatorName, Name, Cfg) ->
89	gen_server:call({via, emqx_limiter_manager, AllocatorName}, {?FUNCTION_NAME, Name, Cfg}).	191✔
90
91	delete_bucket(Name) ->
92	delete_bucket(emqx_limiter:internal_allocator(), Name).	143✔
93
94	delete_bucket(AllocatorName, Name) ->
95	gen_server:call({via, emqx_limiter_manager, AllocatorName}, {?FUNCTION_NAME, Name}).	143✔
96
97	%%--------------------------------------------------------------------
98	%%% gen_server callbacks
99	%%--------------------------------------------------------------------
100
101	-spec init([allocator_name()]) -> {ok, State :: state()}.
102	init([Name]) ->
103	State = init_state(Name),	1,179✔
104	tick_alloc_event(State),	1,179✔
105	{ok, State}.	1,179✔
106
107	handle_call(
108	{add_bucket, Name, #{rate := Rate, burst := Burst}},
109	_From,
110	#{
111	name := AllocatorName,
112	counter := Counter,
113	index := Index,
114	alloc_interval := Interval,
115	buckets := Buckets,
116	free_index := Free
117	} = State
118	) ->
119	case Free of	191✔
120	[ToUse \| Free2] ->
121	NewIndex = Index;	105✔
122	[] ->
123	ToUse = Index + 1,	86✔
124	NewIndex = ToUse,	86✔
125	Free2 = Free	86✔
126	end,
127
128	case ToUse > ?COUNTER_SIZE of	191✔
129	false ->
130	Bucket = do_create_bucket(	191✔
131	Name, Rate, Burst, AllocatorName, Counter, Interval, ToUse
132	),
133
134	{reply, ok, State#{	191✔
135	buckets := Buckets#{Name => Bucket},
136	index := NewIndex,
137	free_index := Free2
138	}};
139	_ ->
NEW 140	{reply, {error, <<"exceeded_max_index">>}, State}	×
141	end;
142	handle_call(
143	{delete_bucket, Name},
144	_From,
145	#{name := AllocatorName, buckets := Buckets, free_index := Free} = State
146	) ->
147	case maps:take(Name, Buckets) of	143✔
148	error ->
149	{reply, ok, State};	1✔
150	{#{index := Index}, Buckets2} ->
151	emqx_limiter_manager:delete_bucket(AllocatorName, Name),	142✔
152	{reply, ok, State#{buckets := Buckets2, free_index := [Index \| Free]}}	142✔
153	end;
154	handle_call(Req, _From, State) ->
NEW 155	?SLOG(error, #{msg => "emqx_limiter_allocator_unexpected_call", call => Req}),	×
NEW 156	{reply, ignored, State}.	×
157
158	handle_cast(Req, State) ->
NEW 159	?SLOG(error, #{msg => "emqx_limiter_allocator_unexpected_cast", cast => Req}),	×
NEW 160	{noreply, State}.	×
161
162	handle_info(tick_alloc_event, State) ->
163	{noreply, do_alloc(State)};	114,153✔
164	handle_info(Info, State) ->
NEW 165	?SLOG(error, #{msg => "emqx_limiter_allocator_unexpected_info", info => Info}),	×
NEW 166	{noreply, State}.	×
167
168	terminate(_Reason, #{name := Name, buckets := Buckets} = _State) ->
NEW 169	maps:foreach(	×
170	fun(LimiterName, _) ->
NEW 171	emqx_limiter_manager:delete_bucket(Name, LimiterName)	×
172	end,
173	Buckets
174	),
NEW 175	ok.	×
176
177	%%--------------------------------------------------------------------
178	%%% Internal functions
179	%%--------------------------------------------------------------------
180	tick_alloc_event(#{alloc_interval := Interval}) ->
181	erlang:send_after(Interval, self(), ?FUNCTION_NAME).	115,332✔
182
183	%% an allocator server with a zone name is used for this zone, or it is a dynamic server
184	init_state(Name) when is_binary(Name) ->
185	Counter = counters:new(?COUNTER_SIZE, [write_concurrency]),	1,179✔
186	#{	1,179✔
187	name => Name,
188	counter => Counter,
189	buckets => #{},
190	index => 0,
191	alloc_interval => emqx_limiter:default_alloc_interval(),
192	lasttime => ?NOW,
193	free_index => []
194	};
195	init_state(Zone) when is_atom(Zone) ->
NEW 196	Cfg = emqx_config:get_zone_conf(Zone, [mqtt, limiter]),	×
NEW 197	init_state(Zone, Cfg).	×
198
199	init_state(Zone, #{alloc_interval := Interval} = Cfg) ->
NEW 200	Counter = counters:new(?COUNTER_SIZE, [write_concurrency]),	×
NEW 201	Names = emqx_limiter_schema:mqtt_limiter_names(),	×
NEW 202	Buckets = init_buckets(Names, Zone, Counter, Cfg, #{}),	×
NEW 203	#{	×
204	name => Zone,
205	counter => Counter,
206	buckets => Buckets,
207	index => maps:size(Buckets),
208	alloc_interval => Interval,
209	lasttime => ?NOW,
210	free_index => []
211	}.
212
213	init_buckets([Name \| Names], Zone, Counter, #{alloc_interval := Interval} = Cfg, Buckets) ->
NEW 214	{ok, RateKey} = emqx_limiter:to_rate_key(Name),	×
NEW 215	{ok, BurstKey} = emqx_limiter:to_burst_key(Name),	×
NEW 216	case maps:get(RateKey, Cfg, infinity) of	×
217	infinity ->
NEW 218	init_buckets(Names, Zone, Counter, Cfg, Buckets);	×
219	Rate ->
NEW 220	Burst = maps:get(BurstKey, Cfg, 0),	×
NEW 221	Bucket = do_create_bucket(	×
222	Name, Rate, Burst, Zone, Counter, Interval, maps:size(Buckets) + 1
223	),
NEW 224	init_buckets(Names, Zone, Counter, Cfg, Buckets#{Name => Bucket})	×
225	end;
226	init_buckets([], _Zone, _Counter, _, Buckets) ->
NEW 227	Buckets.	×
228
229	%% @doc generate tokens, and then spread to leaf nodes
230	-spec do_alloc(state()) -> state().
231	do_alloc(
232	#{
233	lasttime := LastTime,
234	buckets := Buckets
235	} = State
236	) ->
237	tick_alloc_event(State),	114,153✔
238	Now = ?NOW,	114,153✔
239	Elapsed = Now - LastTime,	114,153✔
240	Buckets2 = do_buckets_alloc(Buckets, Elapsed),	114,153✔
241	State#{	114,153✔
242	lasttime := Now,
243	buckets := Buckets2
244	}.
245
246	do_buckets_alloc(Buckets, Elapsed) ->
247	maps:map(	114,153✔
248	fun(_, Bucket) ->
249	do_bucket_alloc(Bucket, Elapsed)	6,973✔
250	end,
251	Buckets
252	).
253
254	do_bucket_alloc(
255	#{
256	rate := Rate,
257	correction := Correction,
258	counter := Counter,
259	index := Index,
260	capacity := Capacity
261	} = Bucket,
262	Elapsed
263	) ->
264	Val = counters:get(Counter, Index),	6,973✔
265	case Val >= Capacity of	6,973✔
266	true ->
267	Bucket;	6,874✔
268	_ ->
269	Inc = Rate * Elapsed + Correction,	99✔
270	Inc2 = erlang:floor(Inc),	99✔
271	Correction2 = Inc - Inc2,	99✔
272	add_tokens(Bucket, Inc2),	99✔
273	Bucket#{correction := Correction2}	99✔
274	end.
275
276	set_tokens(Counter, Ix, Tokens) ->
277	counters:put(Counter, Ix, erlang:floor(Tokens)).	191✔
278
279	add_tokens(_, 0) ->
280	ok;	60✔
281	add_tokens(#{counter := Counter, index := Index, capacity := Capacity}, Tokens) ->
282	Val = counters:get(Counter, Index),	39✔
283	case erlang:min(Capacity, Val + Tokens) - Val of	39✔
284	Inc when Inc > 0 ->
285	counters:put(Counter, Index, Inc + Val);	39✔
286	_ ->
NEW 287	ok	×
288	end.
289
290	do_create_bucket(
291	LimiterName,
292	Rate,
293	Burst,
294	AllocatorName,
295	Counter,
296	Interval,
297	Index
298	) ->
299	Capacity = emqx_limiter:calc_capacity(Rate, Interval),	191✔
300
301	set_tokens(Counter, Index, Capacity),	191✔
302	Ref = emqx_limiter_bucket_ref:new(Counter, Index),	191✔
303	emqx_limiter_manager:insert_bucket(AllocatorName, LimiterName, Ref),	191✔
304	#{	191✔
305	name => LimiterName,
306	rate => Rate,
307	burst => Burst,
308	capacity => Capacity,
309	counter => Counter,
310	index => Index,
311	correction => 0
312	}.

emqx / emqx / 12741272085

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous