Merge branch 'optimise-conn-count' of https://github.com/ferd/ranch

4 files changed, 78 insertions, 9 deletions

diff --git a/src/ranch_acceptor.erl b/src/ranch_acceptor.erl
index aef98d7..63d24c8 100644
--- a/src/ranch_acceptor.erl
+++ b/src/ranch_acceptor.erl
@@ -56,8 +56,13 @@ loop(LSocket, Transport, Protocol, MaxConns, Opts, ListenerPid, ConnsSup) ->
 				[ListenerPid, CSocket, Transport, Protocol, Opts]),
 			Transport:controlling_process(CSocket, ConnPid),
 			ConnPid ! {shoot, ListenerPid},
-			NbConns = ranch_listener:add_connection(ListenerPid, ConnPid),
-			{ok, MaxConns2} = maybe_wait(ListenerPid, MaxConns, NbConns),
+			{ok, MaxConns2} = case MaxConns of
+				infinity ->
+					{ok, infinity};
+				_ ->
+					NbConns = ranch_listener:add_connection(ListenerPid, ConnPid),
+					maybe_wait(ListenerPid, MaxConns, NbConns)
+			end,
 			?MODULE:init(LSocket, Transport, Protocol,
 				MaxConns2, Opts, ListenerPid, ConnsSup);
 		%% Upgrade the max number of connections allowed concurrently.
diff --git a/src/ranch_listener.erl b/src/ranch_listener.erl
index 81943d4..408fbcd 100644
--- a/src/ranch_listener.erl
+++ b/src/ranch_listener.erl
@@ -69,8 +69,14 @@ add_connection(ServerPid, ConnPid) ->
 %% connections.
 -spec remove_connection(pid()) -> non_neg_integer().
 remove_connection(ServerPid) ->
-	ok = gen_server:cast(ServerPid, remove_connection),
-	ranch_server:remove_connection(ServerPid).
+	try
+		Count = ranch_server:remove_connection(ServerPid),
+		ok = gen_server:cast(ServerPid, remove_connection),
+		Count
+	catch
+		error:badarg -> % Max conns = infinity
+			0
+	end.
 
 %% @doc Return the listener's port.
 -spec get_port(pid()) -> {ok, inet:port_number()}.
@@ -105,8 +111,12 @@ set_protocol_options(ServerPid, ProtoOpts) ->
 %% gen_server.
 
 %% @private
+init([Ref, infinity, ProtoOpts]) ->
+	ok = ranch_server:insert_listener(Ref, self()),
+	{ok, #state{ref=Ref, max_conns=infinity, proto_opts=ProtoOpts}};
 init([Ref, MaxConns, ProtoOpts]) ->
 	ok = ranch_server:insert_listener(Ref, self()),
+	ranch_server:add_connections_counter(self()),
 	{ok, #state{ref=Ref, max_conns=MaxConns, proto_opts=ProtoOpts}}.
 
 %% @private
@@ -115,9 +125,19 @@ handle_call(get_port, _From, State=#state{port=Port}) ->
 handle_call(get_max_connections, _From, State=#state{max_conns=MaxConns}) ->
 	{reply, {ok, MaxConns}, State};
 handle_call({set_max_connections, MaxConnections}, _From,
-		State=#state{ref=Ref}) ->
+		State=#state{ref=Ref, max_conns=CurrMax, rm_diff=CurrDiff}) ->
+	RmDiff = case {MaxConnections, CurrMax} of
+		{infinity, _} -> % moving to infinity, delete connection key
+			ranch_server:remove_connections_counter(self()),
+			0;
+		{_, infinity} -> % moving away from infinity, create connection key
+			ranch_server:add_connections_counter(self()),
+			CurrDiff;
+		{_, _} -> % stay current
+			CurrDiff
+	end,
 	ranch_server:send_to_acceptors(Ref, {set_max_conns, MaxConnections}),
-	{reply, ok, State#state{max_conns=MaxConnections}};
+	{reply, ok, State#state{max_conns=MaxConnections, rm_diff=RmDiff}};
 handle_call(get_protocol_options, _From, State=#state{proto_opts=ProtoOpts}) ->
 	{reply, {ok, ProtoOpts}, State};
 handle_call({set_protocol_options, ProtoOpts}, _From, State=#state{ref=Ref}) ->
@@ -132,6 +152,8 @@ handle_call(_, _From, State) ->
 handle_cast({add_connection, ConnPid}, State) ->
 	_ = erlang:monitor(process, ConnPid),
 	{noreply, State};
+handle_cast(remove_connection, State=#state{max_conns=infinity}) ->
+	{noreply, State};
 handle_cast(remove_connection, State=#state{rm_diff=RmDiff}) ->
 	{noreply, State#state{rm_diff=RmDiff + 1}};
 handle_cast({set_port, Port}, State) ->
diff --git a/src/ranch_server.erl b/src/ranch_server.erl
index a17e103..c6d7c19 100644
--- a/src/ranch_server.erl
+++ b/src/ranch_server.erl
@@ -27,6 +27,8 @@
 -export([add_connection/1]).
 -export([count_connections/1]).
 -export([remove_connection/1]).
+-export([add_connections_counter/1]).
+-export([remove_connections_counter/1]).
 
 %% gen_server.
 -export([init/1]).
@@ -95,7 +97,12 @@ add_connection(ListenerPid) ->
 %% @doc Count the number of connections in the connection pool.
 -spec count_connections(pid()) -> non_neg_integer().
 count_connections(ListenerPid) ->
-	ets:update_counter(?TAB, {connections, ListenerPid}, 0).
+	try
+		ets:update_counter(?TAB, {connections, ListenerPid}, 0)
+	catch
+		error:badarg -> % Max conns = infinity
+			0
+	end.
 
 %% @doc Remove a connection from the connection pool.
 %%
@@ -104,6 +111,21 @@ count_connections(ListenerPid) ->
 remove_connection(ListenerPid) ->
 	ets:update_counter(?TAB, {connections, ListenerPid}, -1).
 
+
+%% @doc Add a connections counter to the connection pool
+%%
+%% Should only be used by ranch listeners when settings regarding the max
+%% number of connections change.
+add_connections_counter(Pid) ->
+	true = ets:insert_new(?TAB, {{connections, Pid}, 0}).
+
+%% @doc remove a connections counter from the connection pool
+%%
+%% Should only be used by ranch listeners when settings regarding the max
+%% number of connections change.
+remove_connections_counter(Pid) ->
+	true = ets:delete(?TAB, {connections, Pid}).
+
 %% gen_server.
 
 %% @private
@@ -117,7 +139,6 @@ handle_call(_Request, _From, State) ->
 %% @private
 handle_cast({insert_listener, Ref, Pid}, State=#state{monitors=Monitors}) ->
 	true = ets:insert_new(?TAB, {{acceptors, Ref}, []}),
-	true = ets:insert_new(?TAB, {{connections, Pid}, 0}),
 	MonitorRef = erlang:monitor(process, Pid),
 	{noreply, State#state{
 		monitors=[{{MonitorRef, Pid}, {listener, Ref}}|Monitors]}};
@@ -157,7 +178,7 @@ code_change(_OldVsn, State, _Extra) ->
 remove_process(Key = {listener, Ref}, MonitorRef, Pid, Monitors) ->
 	true = ets:delete(?TAB, Key),
 	true = ets:delete(?TAB, {acceptors, Ref}),
-	true = ets:delete(?TAB, {connections, Pid}),
+	remove_connections_counter(Pid),
 	lists:keydelete({MonitorRef, Pid}, 1, Monitors);
 remove_process(Key = {acceptors, _}, MonitorRef, Pid, Monitors) ->
 	try
diff --git a/test/acceptor_SUITE.erl b/test/acceptor_SUITE.erl
index f5db464..8827ec8 100644
--- a/test/acceptor_SUITE.erl
+++ b/test/acceptor_SUITE.erl
@@ -40,6 +40,7 @@
 -export([tcp_max_connections/1]).
 -export([tcp_max_connections_and_beyond/1]).
 -export([tcp_set_max_connections/1]).
+-export([tcp_infinity_max_connections/1]).
 -export([tcp_upgrade/1]).
 
 %% supervisor.
@@ -58,6 +59,7 @@ groups() ->
 		tcp_active_echo,
 		tcp_echo,
 		tcp_max_connections,
+		tcp_infinity_max_connections,
 		tcp_max_connections_and_beyond,
 		tcp_set_max_connections,
 		tcp_upgrade
@@ -275,6 +277,25 @@ tcp_set_max_connections(_) ->
 	10 = receive_loop(connected, 1000),
 	20 = ranch:get_max_connections(tcp_set_max_connections).
 
+tcp_infinity_max_connections(_) ->
+	{ok, _} = ranch:start_listener(tcp_infinity_max_connections, 1,
+		ranch_tcp, [{port, 0}, {max_connections, 10}],
+		notify_and_wait_protocol, [{msg, connected}, {pid, self()}]),
+	Port = ranch:get_port(tcp_infinity_max_connections),
+	%% @todo We'll probably want a more direct interface to count_connections.
+	ListenerPid = ranch_server:lookup_listener(tcp_infinity_max_connections),
+	ok = connect_loop(Port, 20, 0),
+	10 = ranch_server:count_connections(ListenerPid),
+	10 = receive_loop(connected, 1000),
+	10 = ranch:get_max_connections(tcp_infinity_max_connections),
+	ranch:set_max_connections(tcp_infinity_max_connections, infinity),
+	0 = ranch_server:count_connections(ListenerPid),
+	infinity = ranch:get_max_connections(tcp_infinity_max_connections),
+	ranch:set_max_connections(tcp_infinity_max_connections, 10),
+	0 = ranch_server:count_connections(ListenerPid),
+	10 = receive_loop(connected, 1000),
+	10 = ranch_server:count_connections(ListenerPid). % count could be off
+
 tcp_upgrade(_) ->
 	receive after 20000 -> ok end,
 	{ok, _} = ranch:start_listener(tcp_upgrade, 1,