Gracefully shutdown when stop_listener/1 is called

Implements the `shutdown` option as documented previously.

Makefile

@@ -10,7 +10,7 @@ dep_ct_helper = https://github.com/extend/ct_helper.git master
 # Options.
 
 COMPILE_FIRST = ranch_transport
-CT_SUITES = acceptor sendfile
+CT_SUITES = acceptor sendfile shutdown
 PLT_APPS = crypto public_key ssl
 
 # Standard targets.

src/ranch.erl

@@ -120,7 +120,7 @@ child_spec(Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts)
 		andalso is_atom(Protocol) ->
 	{{ranch_listener_sup, Ref}, {ranch_listener_sup, start_link, [
 		Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts
-	]}, permanent, 5000, supervisor, [ranch_listener_sup]}.
+	]}, permanent, infinity, supervisor, [ranch_listener_sup]}.
 
 %% @doc Acknowledge the accepted connection.
 %%

src/ranch_conns_sup.erl

@@ -20,22 +20,24 @@
 -module(ranch_conns_sup).
 
 %% API.
--export([start_link/5]).
+-export([start_link/6]).
 -export([start_protocol/2]).
 -export([active_connections/1]).
 
 %% Supervisor internals.
--export([init/6]).
+-export([init/7]).
 -export([system_continue/3]).
 -export([system_terminate/4]).
 -export([system_code_change/4]).
 
 -type conn_type() :: worker | supervisor.
+-type shutdown() :: brutal_kill | timeout().
 
 -record(state, {
 	parent = undefined :: pid(),
 	ref :: ranch:ref(),
 	conn_type :: conn_type(),
+	shutdown :: shutdown(),
 	transport = undefined :: module(),
 	protocol = undefined :: module(),
 	opts :: any(),
@@ -45,11 +47,11 @@
 
 %% API.
 
--spec start_link(ranch:ref(), conn_type(), module(), timeout(), module())
-	-> {ok, pid()}.
-start_link(Ref, ConnType, Transport, AckTimeout, Protocol) ->
+-spec start_link(ranch:ref(), conn_type(), shutdown(), module(),
+	timeout(), module()) -> {ok, pid()}.
+start_link(Ref, ConnType, Shutdown, Transport, AckTimeout, Protocol) ->
 	proc_lib:start_link(?MODULE, init,
-		[self(), Ref, ConnType, Transport, AckTimeout, Protocol]).
+		[self(), Ref, ConnType, Shutdown, Transport, AckTimeout, Protocol]).
 
 %% We can safely assume we are on the same node as the supervisor.
 %%
@@ -94,17 +96,17 @@ active_connections(SupPid) ->
 
 %% Supervisor internals.
 
--spec init(pid(), ranch:ref(), conn_type(), module(), timeout(), module())
-	-> no_return().
-init(Parent, Ref, ConnType, Transport, AckTimeout, Protocol) ->
+-spec init(pid(), ranch:ref(), conn_type(), shutdown(),
+	module(), timeout(), module()) -> no_return().
+init(Parent, Ref, ConnType, Shutdown, Transport, AckTimeout, Protocol) ->
 	process_flag(trap_exit, true),
 	ok = ranch_server:set_connections_sup(Ref, self()),
 	MaxConns = ranch_server:get_max_connections(Ref),
 	Opts = ranch_server:get_protocol_options(Ref),
 	ok = proc_lib:init_ack(Parent, {ok, self()}),
 	loop(#state{parent=Parent, ref=Ref, conn_type=ConnType,
-		transport=Transport, protocol=Protocol, opts=Opts,
-		ack_timeout=AckTimeout, max_conns=MaxConns}, 0, 0, []).
+		shutdown=Shutdown, transport=Transport, protocol=Protocol,
+		opts=Opts, ack_timeout=AckTimeout, max_conns=MaxConns}, 0, 0, []).
 
 loop(State=#state{parent=Parent, ref=Ref, conn_type=ConnType,
 		transport=Transport, protocol=Protocol, opts=Opts,
@@ -151,7 +153,7 @@ loop(State=#state{parent=Parent, ref=Ref, conn_type=ConnType,
 			loop(State#state{opts=Opts2},
 				CurConns, NbChildren, Sleepers);
 		{'EXIT', Parent, Reason} ->
-			exit(Reason);
+			terminate(State, Reason, NbChildren);
 		{'EXIT', Pid, Reason} when Sleepers =:= [] ->
 			report_error(Ref, Protocol, Pid, Reason),
 			erase(Pid),
@@ -190,12 +192,59 @@ loop(State=#state{parent=Parent, ref=Ref, conn_type=ConnType,
 				[Ref, Msg])
 	end.
 
+-spec terminate(#state{}, any(), non_neg_integer()) -> no_return().
+%% Kill all children and then exit. We unlink first to avoid
+%% getting a message for each child getting killed.
+terminate(#state{shutdown=brutal_kill}, Reason, _) ->
+	Pids = get_keys(true),
+	_ = [begin
+		unlink(P),
+		exit(P, kill)
+	end || P <- Pids],
+	exit(Reason);
+%% Attempt to gracefully shutdown all children.
+terminate(#state{shutdown=Shutdown}, Reason, NbChildren) ->
+	shutdown_children(),
+	_ = if
+		Shutdown =:= infinity ->
+			ok;
+		true ->
+			erlang:send_after(Shutdown, self(), kill)
+	end,
+	wait_children(NbChildren),
+	exit(Reason).
+
+%% Monitor processes so we can know which ones have shutdown
+%% before the timeout. Unlink so we avoid receiving an extra
+%% message. Then send a shutdown exit signal.
+shutdown_children() ->
+	Pids = get_keys(true),
+	_ = [begin
+		monitor(process, P),
+		unlink(P),
+		exit(P, shutdown)
+	end || P <- Pids],
+	ok.
+
+wait_children(0) ->
+	ok;
+wait_children(NbChildren) ->
+	receive
+        {'DOWN', _, process, Pid, _} ->
+			_ = erase(Pid),
+			wait_children(NbChildren - 1);
+		kill ->
+			Pids = get_keys(true),
+			_ = [exit(P, kill) || P <- Pids],
+			ok
+	end.
+
 system_continue(_, _, {State, CurConns, NbChildren, Sleepers}) ->
 	loop(State, CurConns, NbChildren, Sleepers).
 
 -spec system_terminate(any(), _, _, _) -> no_return().
-system_terminate(Reason, _, _, _) ->
-	exit(Reason).
+system_terminate(Reason, _, _, {State, _, NbChildren, _}) ->
+	terminate(State, Reason, NbChildren).
 
 system_code_change(Misc, _, _, _) ->
 	{ok, Misc}.

src/ranch_listener_sup.erl

@@ -38,9 +38,10 @@ start_link(Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts) ->
 init({Ref, NbAcceptors, Transport, TransOpts, Protocol}) ->
 	AckTimeout = proplists:get_value(ack_timeout, TransOpts, 5000),
 	ConnType = proplists:get_value(connection_type, TransOpts, worker),
+	Shutdown = proplists:get_value(shutdown, TransOpts, 5000),
 	ChildSpecs = [
 		{ranch_conns_sup, {ranch_conns_sup, start_link,
-				[Ref, ConnType, Transport, AckTimeout, Protocol]},
+				[Ref, ConnType, Shutdown, Transport, AckTimeout, Protocol]},
 			permanent, infinity, supervisor, [ranch_conns_sup]},
 		{ranch_acceptors_sup, {ranch_acceptors_sup, start_link,
 				[Ref, NbAcceptors, Transport, TransOpts]},

test/shutdown_SUITE.erl

@@ -0,0 +1,164 @@
+%% Copyright (c) 2013, Loïc Hoguin <essen@ninenines.eu>
+%%
+%% Permission to use, copy, modify, and/or distribute this software for any
+%% purpose with or without fee is hereby granted, provided that the above
+%% copyright notice and this permission notice appear in all copies.
+%%
+%% THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+%% WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+%% MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+%% ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+%% WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+%% ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+%% OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+-module(shutdown_SUITE).
+
+-include_lib("common_test/include/ct.hrl").
+
+%% ct.
+-export([all/0]).
+-export([init_per_suite/1]).
+-export([end_per_suite/1]).
+
+%% Tests.
+
+-export([brutal_kill/1]).
+-export([infinity/1]).
+-export([infinity_trap_exit/1]).
+-export([timeout/1]).
+-export([timeout_trap_exit/1]).
+
+%% ct.
+
+all() ->
+	[brutal_kill, infinity, infinity_trap_exit, timeout, timeout_trap_exit].
+
+init_per_suite(Config) ->
+	ok = application:start(ranch),
+	Config.
+
+end_per_suite(_) ->
+	application:stop(ranch),
+	ok.
+
+%% Tests.
+
+brutal_kill(_) ->
+	Name = brutal_kill,
+	{ok, ListenerSup} = ranch:start_listener(Name, 1,
+		ranch_tcp, [{port, 0}, {shutdown, brutal_kill}],
+		echo_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, _} = gen_tcp:connect("localhost", Port, []),
+	receive after 100 -> ok end,
+	ListenerSupChildren = supervisor:which_children(ListenerSup),
+	{_, ConnsSup, _, _}
+		= lists:keyfind(ranch_conns_sup, 1, ListenerSupChildren),
+	[{_, Pid, _, _}] = supervisor:which_children(ConnsSup),
+	true = is_process_alive(Pid),
+	ranch:stop_listener(Name),
+	receive after 100 -> ok end,
+	false = is_process_alive(Pid),
+	false = is_process_alive(ListenerSup),
+	{error, _} = gen_tcp:connect("localhost", Port, []),
+	ok.
+
+infinity(_) ->
+	Name = infinity,
+	{ok, ListenerSup} = ranch:start_listener(Name, 1,
+		ranch_tcp, [{port, 0}, {shutdown, infinity}],
+		echo_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, _} = gen_tcp:connect("localhost", Port, []),
+	receive after 100 -> ok end,
+	ListenerSupChildren = supervisor:which_children(ListenerSup),
+	{_, ConnsSup, _, _}
+		= lists:keyfind(ranch_conns_sup, 1, ListenerSupChildren),
+	[{_, Pid, _, _}] = supervisor:which_children(ConnsSup),
+	true = is_process_alive(Pid),
+	ranch:stop_listener(Name),
+	receive after 100 -> ok end,
+	false = is_process_alive(Pid),
+	false = is_process_alive(ListenerSup),
+	{error, _} = gen_tcp:connect("localhost", Port, []),
+	ok.
+
+infinity_trap_exit(_) ->
+	Name = infinity_trap_exit,
+	{ok, ListenerSup} = ranch:start_listener(Name, 1,
+		ranch_tcp, [{port, 0}, {shutdown, infinity}],
+		trap_exit_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, _} = gen_tcp:connect("localhost", Port, []),
+	receive after 100 -> ok end,
+	ListenerSupChildren = supervisor:which_children(ListenerSup),
+	{_, ConnsSup, _, _}
+		= lists:keyfind(ranch_conns_sup, 1, ListenerSupChildren),
+	[{_, Pid, _, _}] = supervisor:which_children(ConnsSup),
+	true = is_process_alive(Pid),
+	%% This call will block infinitely.
+	SpawnPid = spawn(fun() -> ranch:stop_listener(Name) end),
+	receive after 100 -> ok end,
+	%% The protocol traps exit signals, and ignore them, so it won't die.
+	true = is_process_alive(Pid),
+	%% The listener will stay up forever too.
+	true = is_process_alive(ListenerSup),
+	%% We can't connect, though.
+	{error, _} = gen_tcp:connect("localhost", Port, []),
+	%% Killing the process unblocks everything.
+	exit(Pid, kill),
+	receive after 100 -> ok end,
+	false = is_process_alive(ListenerSup),
+	false = is_process_alive(SpawnPid),
+	ok.
+
+%% Same as infinity because the protocol doesn't trap exits.
+timeout(_) ->
+	Name = timeout,
+	{ok, ListenerSup} = ranch:start_listener(Name, 1,
+		ranch_tcp, [{port, 0}, {shutdown, 500}],
+		echo_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, _} = gen_tcp:connect("localhost", Port, []),
+	receive after 100 -> ok end,
+	ListenerSupChildren = supervisor:which_children(ListenerSup),
+	{_, ConnsSup, _, _}
+		= lists:keyfind(ranch_conns_sup, 1, ListenerSupChildren),
+	[{_, Pid, _, _}] = supervisor:which_children(ConnsSup),
+	true = is_process_alive(Pid),
+	ranch:stop_listener(Name),
+	receive after 100 -> ok end,
+	false = is_process_alive(Pid),
+	false = is_process_alive(ListenerSup),
+	{error, _} = gen_tcp:connect("localhost", Port, []),
+	ok.
+
+timeout_trap_exit(_) ->
+	Name = timeout_trap_exit,
+	{ok, ListenerSup} = ranch:start_listener(Name, 1,
+		ranch_tcp, [{port, 0}, {shutdown, 500}],
+		trap_exit_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, _} = gen_tcp:connect("localhost", Port, []),
+	receive after 100 -> ok end,
+	ListenerSupChildren = supervisor:which_children(ListenerSup),
+	{_, ConnsSup, _, _}
+		= lists:keyfind(ranch_conns_sup, 1, ListenerSupChildren),
+	[{_, Pid, _, _}] = supervisor:which_children(ConnsSup),
+	true = is_process_alive(Pid),
+	%% This call will block for the duration of the shutdown.
+	SpawnPid = spawn(fun() -> ranch:stop_listener(Name) end),
+	receive after 100 -> ok end,
+	%% The protocol traps exit signals, and ignore them, so it won't die.
+	true = is_process_alive(Pid),
+	%% The listener will stay up for now too.
+	true = is_process_alive(ListenerSup),
+	%% We can't connect, though.
+	{error, _} = gen_tcp:connect("localhost", Port, []),
+	%% Wait for the timeout to finish and see that everything is killed.
+	receive after 500 -> ok end,
+	false = is_process_alive(Pid),
+	false = is_process_alive(ListenerSup),
+	false = is_process_alive(SpawnPid),
+	ok.

test/trap_exit_protocol.erl

@@ -0,0 +1,23 @@
+-module(trap_exit_protocol).
+-behaviour(ranch_protocol).
+
+-export([start_link/4]).
+-export([init/4]).
+
+start_link(Ref, Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+	{ok, Pid}.
+
+init(Ref, Socket, Transport, _Opts = []) ->
+	process_flag(trap_exit, true),
+	ok = ranch:accept_ack(Ref),
+	loop(Socket, Transport).
+
+loop(Socket, Transport) ->
+	case Transport:recv(Socket, 0, infinity) of
+		{ok, Data} ->
+			Transport:send(Socket, Data),
+			loop(Socket, Transport);
+		_ ->
+			ok = Transport:close(Socket)
+	end.