Merge fef6c2f into 21fa934

esl · Dec 2, 2019 · 64fa7e4 · 64fa7e4
2 parents 21fa934 + fef6c2f
commit 64fa7e4
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 22 deletions.
diff --git a/src/ejabberd_receiver.erl b/src/ejabberd_receiver.erl
@@ -356,7 +356,7 @@ process_data(Data, #state{parser = Parser,
                           stanza_chunk_size = ChunkSize,
                           c2s_pid = C2SPid} = State) ->
     ?DEBUG("Received XML on stream = \"~s\"", [Data]),
-    Size = size(Data),
+    Size = byte_size(Data),
     maybe_run_keep_alive_hook(Size, State),
     {C2SEvents, NewParser} =
         case exml_stream:parse(Parser, Data) of

diff --git a/src/shaper.erl b/src/shaper.erl
@@ -32,31 +32,34 @@
 -spec new(atom()) -> shaper().
 new(Name) ->
     case ejabberd_config:get_global_option({shaper, Name, global}) of
-        undefined -> none;
-        none -> none;
-        {maxrate, MaxRate} -> #shaper{max_rate = MaxRate, tokens = MaxRate}
+        {maxrate, MaxRatePerSecond} ->
+            #shaper{max_rate = MaxRatePerSecond,
+                    tokens = MaxRatePerSecond,
+                    last_update = erlang:monotonic_time(millisecond)};
+        _ -> none
     end.
 
 %% @doc Update shaper.
 %% `Delay' is how many milliseconds to wait.
 -spec update(shaper(), Size :: pos_integer()) -> {shaper(), Delay :: non_neg_integer()}.
 update(none, _Size) ->
     {none, 0};
-update(Shaper, Size) ->
-    Now = erlang:monotonic_time(),
-    Second = erlang:convert_time_unit(1, seconds, native),
-
-    SecondsSinceLastUpdate = (Now - Shaper#shaper.last_update) / Second,
-    TokenGrowth = round(Shaper#shaper.max_rate * SecondsSinceLastUpdate),
-    Tokens = min(Shaper#shaper.max_rate, Shaper#shaper.tokens + TokenGrowth),
-
-    TokensLeft = max(0, Tokens - Size),
-    AdditionalTokensNeeded = max(0, Size - Tokens),
-
-    TimeNeededForTokensToGrow = round(AdditionalTokensNeeded / Shaper#shaper.max_rate * Second),
-    Delay = erlang:convert_time_unit(TimeNeededForTokensToGrow, native, millisecond),
-    LastUpdate = Now + TimeNeededForTokensToGrow,
-
-    lager:debug("Tokens: ~p (+~p,-~p), delay: ~p ms", [TokensLeft, TokenGrowth, Size, Delay]),
-
-    {Shaper#shaper{last_update = LastUpdate, tokens = TokensLeft}, Delay}.
+update(#shaper{max_rate = MaxRatePerSecond,
+               tokens = LastAvailableTokens,
+               last_update = LastUpdate}, NowUsed) ->
+    Now = erlang:monotonic_time(millisecond),
+    % How much we might have recovered since last time, in milliseconds arithmetic
+    GrowthPerMillisecond = MaxRatePerSecond / 1000,
+    MilliSecondsSinceLastUpdate = (Now - LastUpdate),
+    PossibleTokenGrowth = round(GrowthPerMillisecond * MilliSecondsSinceLastUpdate),
+    % Available plus recovered cannot grow higher than the actual rate limit
+    ExactlyAvailableNow = min(MaxRatePerSecond, LastAvailableTokens + PossibleTokenGrowth),
+    TokensAvailable = max(0, ExactlyAvailableNow - NowUsed),
+    TokensOverused = max(0, NowUsed - ExactlyAvailableNow),
+    MaybeDelay = round(TokensOverused / GrowthPerMillisecond),
+    lager:debug("Tokens: ~p (+~p,-~p), delay: ~p ms",
+                [TokensOverused, TokensAvailable, NowUsed, MaybeDelay]),
+    {#shaper{max_rate = MaxRatePerSecond,
+             tokens = TokensAvailable,
+             last_update = Now + MaybeDelay},
+     MaybeDelay}.