fix: do not consume limiter token if failed to accept

emqx · Oct 10, 2024 · 42b71da · 42b71da
1 parent 903fb7d
commit 42b71da
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/src/esockd_acceptor.erl b/src/esockd_acceptor.erl
@@ -116,10 +116,11 @@ accepting(info, {inet_async, LSock, Ref, {ok, Sock}},
     %% Inc accepted stats.
     esockd_server:inc_stats({Proto, ListenOn}, accepted, 1),
 
-    case eval_tune_socket_fun(TuneFun, Sock) of
-        {ok, Sock} ->
+    Result = case eval_tune_socket_fun(TuneFun, Sock) of
+        {ok, _} ->
             case esockd_connection_sup:start_connection(ConnSup, Sock, UpgradeFuns) of
-                {ok, _Pid} -> ok;
+                {ok, _Pid} ->
+                    consume_limiter;
                 {error, enotconn} ->
                     close(Sock); %% quiet...issue #10
                 {error, einval} ->
@@ -140,8 +141,7 @@ accepting(info, {inet_async, LSock, Ref, {ok, Sock}},
                                    [esockd:format(Sockname), Reason]),
             close(Sock)
     end,
-    rate_limit(State);
-
+    rate_limit(State, Result);
 accepting(info, {inet_async, LSock, Ref, {error, closed}},
           State = #state{lsock = LSock, accept_ref = Ref}) ->
     {stop, normal, State};
@@ -184,13 +184,16 @@ code_change(_OldVsn, StateName, State, _Extra) ->
 
 close(Sock) -> catch port_close(Sock).
 
-rate_limit(State = #state{conn_limiter = Limiter}) ->
+rate_limit(State = #state{conn_limiter = Limiter}, consume_limiter) ->
     case esockd_limiter:consume(Limiter, 1) of
         {I, Pause} when I =< 0 ->
             {next_state, suspending, State, Pause};
         _ ->
             {keep_state, State, {next_event, internal, accept}}
-    end.
+    end;
+rate_limit(State, _NotAccepted) ->
+    %% Socket closed or error by the time when accepting it
+    {keep_state, State, {next_event, internal, accept}}.
 
 eval_tune_socket_fun({Fun, Args1}, Sock) ->
     apply(Fun, [Sock|Args1]).
diff --git a/src/esockd_dtls_acceptor.erl b/src/esockd_dtls_acceptor.erl
@@ -93,10 +93,11 @@ accepting(internal, accept,
         {ok, Sock} ->
             %% Inc accepted stats.
             esockd_server:inc_stats({Proto, ListenOn}, accepted, 1),
-            _ = case eval_tune_socket_fun(TuneFun, Sock) of
+            Result = case eval_tune_socket_fun(TuneFun, Sock) of
                 {ok, Sock} ->
                     case esockd_connection_sup:start_connection(ConnSup, Sock, UpgradeFuns) of
-                        {ok, _Pid} -> ok;
+                        {ok, _Pid} ->
+                            consume_limiter;
                         {error, enotconn} ->
                             close(Sock); %% quiet...issue #10
                         {error, einval} ->
@@ -117,7 +118,7 @@ accepting(internal, accept,
                                            [esockd:format(Sockname), Reason]),
                     close(Sock)
             end,
-            rate_limit(State);
+            rate_limit(State, Result);
         {error, Reason} when Reason =:= emfile;
                              Reason =:= enfile ->
             {next_state, suspending, State, 1000};
@@ -148,13 +149,16 @@ code_change(_OldVsn, StateName, State, _Extra) ->
 
 close(Sock) -> ssl:close(Sock).
 
-rate_limit(State = #state{conn_limiter = Limiter}) ->
+rate_limit(State = #state{conn_limiter = Limiter}, consume_limiter) ->
     case esockd_limiter:consume(Limiter, 1) of
         {I, Pause} when I =< 0 ->
             {next_state, suspending, State, Pause};
         _ ->
             {keep_state, State, {next_event, internal, accept}}
-    end.
+    end;
+rate_limit(State, _NotAccepted) ->
+    %% Socket closed or error by the time when accepting it
+    {keep_state, State, {next_event, internal, accept}}.
 
 eval_tune_socket_fun({Fun, Args1}, Sock) ->
     apply(Fun, [Sock|Args1]).