Add ranch:set_max_connections/2 and get_max_connections/1

5 files changed, 69 insertions, 9 deletions

diff --git a/guide/listeners.md b/guide/listeners.md
index 32fddde..1209706 100644
--- a/guide/listeners.md
+++ b/guide/listeners.md
@@ -172,6 +172,16 @@ ranch_listener:remove_connection(ListenerPid).
 As seen in the chapter covering protocols, this pid is received as the
 first argument of the protocol's `start_link/4` callback.
 
+You can modify the `max_connections` value on a running listener by
+using the `ranch:set_max_connections/2` function, with the name of the
+listener as first argument and the new value as the second.
+
+``` erlang
+ranch:set_max_connections(tcp_echo, MaxConns).
+```
+
+The change will occur immediately.
+
 Upgrading
 ---------
 
diff --git a/src/ranch.erl b/src/ranch.erl
index 59e1fa7..c582efa 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -20,6 +20,8 @@
 -export([child_spec/6]).
 -export([accept_ack/1]).
 -export([get_port/1]).
+-export([get_max_connections/1]).
+-export([set_max_connections/2]).
 -export([get_protocol_options/1]).
 -export([set_protocol_options/2]).
 -export([filter_options/3]).
@@ -121,6 +123,19 @@ get_port(Ref) ->
 	{ok, Port} = ranch_listener:get_port(ListenerPid),
 	Port.
 
+%% @doc Return the max number of connections allowed concurrently.
+-spec get_max_connections(any()) -> non_neg_integer().
+get_max_connections(Ref) ->
+	ListenerPid = ranch_server:lookup_listener(Ref),
+	{ok, MaxConnections} = ranch_listener:get_max_connections(ListenerPid),
+	MaxConnections.
+
+%% @doc Set the max number of connections allowed concurrently.
+-spec set_max_connections(any(), non_neg_integer()) -> ok.
+set_max_connections(Ref, MaxConnections) ->
+	ListenerPid = ranch_server:lookup_listener(Ref),
+	ok = ranch_listener:set_max_connections(ListenerPid, MaxConnections).
+
 %% @doc Return the current protocol options for the given listener.
 -spec get_protocol_options(any()) -> any().
 get_protocol_options(Ref) ->
diff --git a/src/ranch_acceptor.erl b/src/ranch_acceptor.erl
index d208a7f..1985826 100644
--- a/src/ranch_acceptor.erl
+++ b/src/ranch_acceptor.erl
@@ -57,22 +57,31 @@ loop(LSocket, Transport, Protocol, MaxConns, Opts, ListenerPid, ConnsSup) ->
 			Transport:controlling_process(CSocket, ConnPid),
 			ConnPid ! {shoot, ListenerPid},
 			NbConns = ranch_listener:add_connection(ListenerPid, ConnPid),
-			maybe_wait(ListenerPid, MaxConns, NbConns),
+			{ok, MaxConns2} = maybe_wait(ListenerPid, MaxConns, NbConns),
 			?MODULE:init(LSocket, Transport, Protocol,
-				MaxConns, Opts, ListenerPid, ConnsSup);
+				MaxConns2, Opts, ListenerPid, ConnsSup);
+		%% Upgrade the max number of connections allowed concurrently.
+		{set_max_conns, MaxConns2} ->
+			?MODULE:loop(LSocket, Transport, Protocol,
+				MaxConns2, Opts, ListenerPid, ConnsSup);
 		%% Upgrade the protocol options.
 		{set_opts, Opts2} ->
 			?MODULE:loop(LSocket, Transport, Protocol,
 				MaxConns, Opts2, ListenerPid, ConnsSup)
 	end.
 
--spec maybe_wait(pid(), non_neg_integer(), non_neg_integer()) -> ok.
+-spec maybe_wait(pid(), non_neg_integer(), non_neg_integer())
+	-> {ok, non_neg_integer()}.
 maybe_wait(_, MaxConns, NbConns) when MaxConns > NbConns ->
-	ok;
-maybe_wait(ListenerPid, MaxConns, _) ->
-	erlang:yield(),
-	NbConns2 = ranch_server:count_connections(ListenerPid),
-	maybe_wait(ListenerPid, MaxConns, NbConns2).
+	{ok, MaxConns};
+maybe_wait(ListenerPid, MaxConns, NbConns) ->
+	receive
+		{set_max_conns, MaxConns2} ->
+			maybe_wait(ListenerPid, MaxConns2, NbConns)
+	after 0 ->
+		NbConns2 = ranch_server:count_connections(ListenerPid),
+		maybe_wait(ListenerPid, MaxConns, NbConns2)
+	end.
 
 -spec async_accept(inet:socket(), module()) -> ok.
 async_accept(LSocket, Transport) ->
diff --git a/src/ranch_listener.erl b/src/ranch_listener.erl
index 52e4c32..873d860 100644
--- a/src/ranch_listener.erl
+++ b/src/ranch_listener.erl
@@ -24,6 +24,7 @@
 -export([get_port/1]).
 -export([set_port/2]).
 -export([get_max_connections/1]).
+-export([set_max_connections/2]).
 -export([get_protocol_options/1]).
 -export([set_protocol_options/2]).
 
@@ -82,11 +83,15 @@ set_port(ServerPid, Port) ->
 	gen_server:cast(ServerPid, {set_port, Port}).
 
 %% @doc Return the max number of connections allowed concurrently.
-%% @todo Add set_max_connections.
 -spec get_max_connections(pid()) -> {ok, non_neg_integer()}.
 get_max_connections(ServerPid) ->
 	gen_server:call(ServerPid, get_max_connections).
 
+%% @doc Set the max number of connections allowed concurrently.
+-spec set_max_connections(pid(), non_neg_integer()) -> ok.
+set_max_connections(ServerPid, MaxConnections) ->
+	gen_server:call(ServerPid, {set_max_connections, MaxConnections}).
+
 %% @doc Return the current protocol options.
 -spec get_protocol_options(pid()) -> {ok, any()}.
 get_protocol_options(ServerPid) ->
@@ -109,6 +114,10 @@ handle_call(get_port, _From, State=#state{port=Port}) ->
 	{reply, {ok, Port}, State};
 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}) ->
+	ranch_server:send_to_acceptors(Ref, {set_max_conns, MaxConnections}),
+	{reply, ok, State#state{max_conns=MaxConnections}};
 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}) ->
diff --git a/test/acceptor_SUITE.erl b/test/acceptor_SUITE.erl
index 582e564..f5db464 100644
--- a/test/acceptor_SUITE.erl
+++ b/test/acceptor_SUITE.erl
@@ -39,6 +39,7 @@
 -export([tcp_echo/1]).
 -export([tcp_max_connections/1]).
 -export([tcp_max_connections_and_beyond/1]).
+-export([tcp_set_max_connections/1]).
 -export([tcp_upgrade/1]).
 
 %% supervisor.
@@ -58,6 +59,7 @@ groups() ->
 		tcp_echo,
 		tcp_max_connections,
 		tcp_max_connections_and_beyond,
+		tcp_set_max_connections,
 		tcp_upgrade
 	]}, {ssl, [
 		ssl_accept_error,
@@ -258,6 +260,21 @@ tcp_max_connections_and_beyond(_) ->
 	receive after 500 -> ok end,
 	10 = ranch_server:count_connections(ListenerPid).
 
+tcp_set_max_connections(_) ->
+	{ok, _} = ranch:start_listener(tcp_set_max_connections, 1,
+		ranch_tcp, [{port, 0}, {max_connections, 10}],
+		notify_and_wait_protocol, [{msg, connected}, {pid, self()}]),
+	Port = ranch:get_port(tcp_set_max_connections),
+	%% @todo We'll probably want a more direct interface to count_connections.
+	ListenerPid = ranch_server:lookup_listener(tcp_set_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_set_max_connections),
+	ranch:set_max_connections(tcp_set_max_connections, 20),
+	10 = receive_loop(connected, 1000),
+	20 = ranch:get_max_connections(tcp_set_max_connections).
+
 tcp_upgrade(_) ->
 	receive after 20000 -> ok end,
 	{ok, _} = ranch:start_listener(tcp_upgrade, 1,