%%%============================================================================= %%% @doc speed_trap_dynamic - Dynamic rate limiter based on rejection rates %%% %%% This module implements a dynamic rate limiter that automatically adjusts %%% bucket_size based on rejection rates. It monitors the rejection rate over a %%% configurable time interval and scales bucket_size up or down accordingly. %%% %%% Key features: %%% - Automatic upscaling when rejection rate exceeds threshold %%% - Automatic downscaling when rejection rate is below threshold AND %%% total requests are less than downscaled capacity during the scaling_interval %%% - Configurable min/max bucket_size, scaling time intervals, and adjustment increments %%% - Gradual bucket_size changes to allow downstream systems to scale %%% - Uses atomics for high-performance rejection/acceptance tracking %%% %%% @copyright 2025 Klarna Bank AB %%% @end %%%============================================================================= -module(speed_trap_dynamic). -export([start_link/0, register_dynamic_limiter/2, unregister_dynamic_limiter/1, record_rejection/1, record_acceptance/1]). -behaviour(gen_server). -export([init/1, handle_call/3, handle_cast/2, handle_info/2]). -define(SERVER, ?MODULE). -define(ETS_DYNAMIC_COUNTERS, speed_trap_dynamic_counters). -define(REJECTION_IDX, 1). -define(SUCCESS_IDX, 2). -record(dynamic_state, {id :: speed_trap:id(), counters :: atomics:atomics_ref(), timer_ref :: timer:tref() | undefined}). -type dynamic_state() :: #dynamic_state{}. -type state() :: #{speed_trap:id() => dynamic_state()}. %%----------------------------------------------------------------------------- %% API %%----------------------------------------------------------------------------- -spec start_link() -> {ok, pid()} | {error, term()}. start_link() -> gen_server:start_link({local, ?SERVER}, ?MODULE, [], []). -spec register_dynamic_limiter(speed_trap:id(), speed_trap:options()) -> ok | {error, already_registered}. register_dynamic_limiter(Id, Options) -> gen_server:call(?SERVER, {register, Id, Options}). -spec unregister_dynamic_limiter(speed_trap:id()) -> ok. unregister_dynamic_limiter(Id) -> gen_server:call(?SERVER, {unregister, Id}). -spec record_rejection(speed_trap:id()) -> ok. record_rejection(Id) -> case ets:lookup(?ETS_DYNAMIC_COUNTERS, Id) of [{Id, Counters}] -> atomics:add(Counters, ?REJECTION_IDX, 1); [] -> ok end. -spec record_acceptance(speed_trap:id()) -> ok. record_acceptance(Id) -> case ets:lookup(?ETS_DYNAMIC_COUNTERS, Id) of [{Id, Counters}] -> atomics:add(Counters, ?SUCCESS_IDX, 1); [] -> ok end. %%----------------------------------------------------------------------------- %% gen_server callbacks %%----------------------------------------------------------------------------- -spec init([]) -> {ok, state()}. init([]) -> ?ETS_DYNAMIC_COUNTERS = ets:new(?ETS_DYNAMIC_COUNTERS, [set, protected, named_table, {read_concurrency, true}]), {ok, #{}}. -spec handle_call(term(), {pid(), term()}, state()) -> {reply, term(), state()}. handle_call({register, Id, Options}, _From, State) -> case maps:is_key(Id, State) of true -> {reply, {error, already_registered}, State}; false -> #{scaling_time_interval := TimeInterval} = Options, % Create atomic counters for rejections and acceptances Counters = atomics:new(2, [{signed, false}]), atomics:put(Counters, ?REJECTION_IDX, 0), atomics:put(Counters, ?SUCCESS_IDX, 0), {ok, TimerRef} = timer:send_after(TimeInterval, self(), {check_and_adjust, Id}), DynState = #dynamic_state{id = Id, counters = Counters, timer_ref = TimerRef}, NewState = State#{Id => DynState}, ets:insert(?ETS_DYNAMIC_COUNTERS, {Id, Counters}), {reply, ok, NewState} end; handle_call({unregister, Id}, _From, State) -> case maps:take(Id, State) of {DynState, NewState} -> case DynState#dynamic_state.timer_ref of undefined -> ok; TimerRef -> timer:cancel(TimerRef) end, ets:delete(?ETS_DYNAMIC_COUNTERS, Id), {reply, ok, NewState}; error -> {reply, ok, State} end; handle_call(_Request, _From, State) -> {reply, {error, unknown_call}, State}. -spec handle_cast(term(), state()) -> {noreply, state()}. handle_cast(_Request, State) -> {noreply, State}. -spec handle_info(term(), state()) -> {noreply, state()}. handle_info({check_and_adjust, Id}, State) -> case maps:get(Id, State, undefined) of undefined -> {noreply, State}; DynState -> NewDynState = check_and_adjust_bucket_size(DynState), NewState = State#{Id => NewDynState}, {noreply, NewState} end; handle_info(_Info, State) -> {noreply, State}. %%----------------------------------------------------------------------------- %% Internal functions %%----------------------------------------------------------------------------- -spec check_and_adjust_bucket_size(dynamic_state()) -> dynamic_state(). check_and_adjust_bucket_size(DynState) -> #dynamic_state{id = Id, counters = Counters} = DynState, {ok, #{bucket_size := CurrentBucketSize, min_bucket_size := MinBucketSize, max_bucket_size := MaxBucketSize, refill_count := RefillCount, refill_interval := RefillInterval, scaling_time_interval := ScalingTimeInterval, rejection_rate_threshold := Threshold, scaling_bucket_size_adjust_count := AdjustCount}} = speed_trap_token_bucket:options(Id), % Read and reset counters atomically TotalRejections = atomics:exchange(Counters, ?REJECTION_IDX, 0), TotalSuccesses = atomics:exchange(Counters, ?SUCCESS_IDX, 0), TotalRequests = TotalRejections + TotalSuccesses, RejectionRate = case TotalRequests of 0 -> 0; _ -> round(TotalRejections / TotalRequests * 100) end, %% Calculate the target bucket size after potential downscaling PotentialDownscaleTarget = max(CurrentBucketSize - AdjustCount, MinBucketSize), %% Calculate the throughput capacity at the downscaled bucket size. %% Since refill_count scales proportionally with bucket_size, the capacity %% at the downscaled size is: (ScalingTimeInterval / RefillInterval) * DownscaledRefillCount %% where DownscaledRefillCount = RefillCount * (PotentialDownscaleTarget / CurrentBucketSize) PotentialDownscaledCapacity = trunc(ScalingTimeInterval / RefillInterval * RefillCount * (PotentialDownscaleTarget / CurrentBucketSize)), %% Only downscale if rejection rate is below threshold AND total requests %% are less than what the downscaled bucket could handle. This prevents %% downscaling when traffic would overwhelm the smaller bucket. ShouldDownscale = RejectionRate < Threshold andalso CurrentBucketSize > MinBucketSize andalso TotalRequests < PotentialDownscaledCapacity, if RejectionRate > Threshold andalso CurrentBucketSize < MaxBucketSize -> % Upscale: increase bucket_size NewQ = min(CurrentBucketSize + AdjustCount, MaxBucketSize), adjust_bucket_size(Id, CurrentBucketSize, NewQ, RefillCount); ShouldDownscale -> % Downscale: decrease bucket_size only when traffic is low enough adjust_bucket_size(Id, CurrentBucketSize, PotentialDownscaleTarget, RefillCount); true -> % No adjustment needed - either at optimal size or traffic justifies current size ok end, % Schedule next check {ok, TimerRef} = timer:send_after(ScalingTimeInterval, self(), {check_and_adjust, Id}), DynState#dynamic_state{timer_ref = TimerRef}. -spec adjust_bucket_size(speed_trap:id(), pos_integer(), pos_integer(), pos_integer()) -> ok. adjust_bucket_size(Id, CurrentBucketSize, NewBucketSize, RefillCount) -> %% Calculate scaling factor ScalingFactor = NewBucketSize / CurrentBucketSize, %% Scale refill_count proportionally NewRefillCount = round(RefillCount * ScalingFactor), %% Ensure minimum refill count AdjustedRefillCount = max(1, NewRefillCount), ModifyOpts = #{bucket_size => NewBucketSize, refill_count => AdjustedRefillCount}, ok = speed_trap:modify(Id, ModifyOpts).