Merge branch 'fix/listener_sup_failures' of git://github.com/keynslug/syncranch

6 files changed, 55 insertions, 32 deletions

diff --git a/src/ranch.erl b/src/ranch.erl
index 0924ee7..59e1fa7 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -55,6 +55,7 @@
 start_listener(Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts)
 		when is_integer(NbAcceptors) andalso is_atom(Transport)
 		andalso is_atom(Protocol) ->
+	_ = code:ensure_loaded(Transport),
 	case erlang:function_exported(Transport, name, 0) of
 		false ->
 			{error, badarg};
diff --git a/src/ranch_acceptors_sup.erl b/src/ranch_acceptors_sup.erl
index 1d9503c..f1908a6 100644
--- a/src/ranch_acceptors_sup.erl
+++ b/src/ranch_acceptors_sup.erl
@@ -17,7 +17,7 @@
 -behaviour(supervisor).
 
 %% API.
--export([start_link/7]).
+-export([start_link/5]).
 
 %% supervisor.
 -export([init/1]).
@@ -25,16 +25,16 @@
 %% API.
 
 -spec start_link(any(), non_neg_integer(), module(), any(),
-	module(), pid(), pid()) -> {ok, pid()}.
-start_link(Ref, NbAcceptors, Transport, TransOpts,
-		Protocol, ListenerPid, ConnsPid) ->
+	module()) -> {ok, pid()}.
+start_link(Ref, NbAcceptors, Transport, TransOpts, Protocol) ->
 	supervisor:start_link(?MODULE, [Ref, NbAcceptors, Transport, TransOpts,
-		Protocol, ListenerPid, ConnsPid]).
+		Protocol]).
 
 %% supervisor.
 
-init([Ref, NbAcceptors, Transport, TransOpts,
-		Protocol, ListenerPid, ConnsPid]) ->
+init([Ref, NbAcceptors, Transport, TransOpts, Protocol]) ->
+	ListenerPid = ranch_server:lookup_listener(Ref),
+	ConnsPid = ranch_server:lookup_connections_sup(Ref),
 	LSocket = case proplists:get_value(socket, TransOpts) of
 		undefined ->
 			{ok, Socket} = Transport:listen(TransOpts),
diff --git a/src/ranch_conns_sup.erl b/src/ranch_conns_sup.erl
index 99853c5..3cc09be 100644
--- a/src/ranch_conns_sup.erl
+++ b/src/ranch_conns_sup.erl
@@ -17,7 +17,7 @@
 -behaviour(supervisor).
 
 %% API.
--export([start_link/0]).
+-export([start_link/1]).
 -export([start_protocol/5]).
 
 %% supervisor.
@@ -25,9 +25,9 @@
 
 %% API.
 
--spec start_link() -> {ok, pid()}.
-start_link() ->
-	supervisor:start_link(?MODULE, []).
+-spec start_link(any()) -> {ok, pid()}.
+start_link(Ref) ->
+	supervisor:start_link(?MODULE, Ref).
 
 -spec start_protocol(pid(), inet:socket(), module(), module(), any())
 	-> {ok, pid()}.
@@ -36,6 +36,7 @@ start_protocol(ListenerPid, Socket, Transport, Protocol, Opts) ->
 
 %% supervisor.
 
-init([]) ->
+init(Ref) ->
+	ok = ranch_server:set_connections_sup(Ref, self()),
 	{ok, {{simple_one_for_one, 0, 1}, [{?MODULE, {?MODULE, start_protocol, []},
 		temporary, brutal_kill, worker, [?MODULE]}]}}.
diff --git a/src/ranch_listener.erl b/src/ranch_listener.erl
index 83ee658..52e4c32 100644
--- a/src/ranch_listener.erl
+++ b/src/ranch_listener.erl
@@ -101,6 +101,7 @@ set_protocol_options(ServerPid, ProtoOpts) ->
 
 %% @private
 init([Ref, MaxConns, ProtoOpts]) ->
+	ok = ranch_server:insert_listener(Ref, self()),
 	{ok, #state{ref=Ref, max_conns=MaxConns, proto_opts=ProtoOpts}}.
 
 %% @private
diff --git a/src/ranch_listener_sup.erl b/src/ranch_listener_sup.erl
index c8ba12d..0147cf2 100644
--- a/src/ranch_listener_sup.erl
+++ b/src/ranch_listener_sup.erl
@@ -28,23 +28,25 @@
 	-> {ok, pid()}.
 start_link(Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts) ->
 	MaxConns = proplists:get_value(max_connections, TransOpts, 1024),
-	{ok, SupPid} = supervisor:start_link(?MODULE, []),
-	{ok, ListenerPid} = supervisor:start_child(SupPid,
-		{ranch_listener, {ranch_listener, start_link,
-			[Ref, MaxConns, ProtoOpts]},
-		 permanent, 5000, worker, [ranch_listener]}),
-	ok = ranch_server:insert_listener(Ref, ListenerPid),
-	{ok, ConnsPid} = supervisor:start_child(SupPid,
-		{ranch_conns_sup, {ranch_conns_sup, start_link, []},
-		 permanent, 5000, supervisor, [ranch_conns_sup]}),
-	{ok, _PoolPid} = supervisor:start_child(SupPid,
-		{ranch_acceptors_sup, {ranch_acceptors_sup, start_link, [
-			Ref, NbAcceptors, Transport, TransOpts,
-			Protocol, ListenerPid, ConnsPid
-		]}, permanent, 5000, supervisor, [ranch_acceptors_sup]}),
-	{ok, SupPid}.
+	supervisor:start_link(?MODULE, {
+		Ref, NbAcceptors, MaxConns, Transport, TransOpts, Protocol, ProtoOpts
+		}).
 
 %% supervisor.
 
-init([]) ->
-	{ok, {{one_for_all, 10, 10}, []}}.
+init({Ref, NbAcceptors, MaxConns, Transport, TransOpts, Protocol, ProtoOpts}) ->
+	ChildSpecs = [
+		%% listener
+		{ranch_listener, {ranch_listener, start_link,
+			[Ref, MaxConns, ProtoOpts]},
+		 permanent, 5000, worker, [ranch_listener]},
+		%% conns_sup
+		{ranch_conns_sup, {ranch_conns_sup, start_link, [Ref]},
+		 permanent, infinity, supervisor, [ranch_conns_sup]},
+		%% acceptors_sup
+		{ranch_acceptors_sup, {ranch_acceptors_sup, start_link,
+			[Ref, NbAcceptors, Transport, TransOpts, Protocol]
+		 }, permanent, infinity, supervisor, [ranch_acceptors_sup]}
+	],
+	{ok, {{rest_for_one, 10, 10}, ChildSpecs}}.
+
diff --git a/src/ranch_server.erl b/src/ranch_server.erl
index b0ae612..a17e103 100644
--- a/src/ranch_server.erl
+++ b/src/ranch_server.erl
@@ -20,6 +20,8 @@
 -export([start_link/0]).
 -export([insert_listener/2]).
 -export([lookup_listener/1]).
+-export([set_connections_sup/2]).
+-export([lookup_connections_sup/1]).
 -export([add_acceptor/2]).
 -export([send_to_acceptors/2]).
 -export([add_connection/1]).
@@ -52,7 +54,7 @@ start_link() ->
 %% @doc Insert a listener into the database.
 -spec insert_listener(any(), pid()) -> ok.
 insert_listener(Ref, Pid) ->
-	true = ets:insert_new(?TAB, {{listener, Ref}, Pid}),
+	true = ets:insert_new(?TAB, {{listener, Ref}, Pid, undefined}),
 	gen_server:cast(?MODULE, {insert_listener, Ref, Pid}).
 
 %% @doc Lookup a listener in the database.
@@ -60,6 +62,17 @@ insert_listener(Ref, Pid) ->
 lookup_listener(Ref) ->
 	ets:lookup_element(?TAB, {listener, Ref}, 2).
 
+%% @doc Set a connection supervisor associated with specific listener.
+-spec set_connections_sup(any(), pid()) -> ok.
+set_connections_sup(Ref, Pid) ->
+	true = ets:update_element(?TAB, {listener, Ref}, {3, Pid}),
+	ok.
+
+%% @doc Lookup a connection supervisor used by specific listener.
+-spec lookup_connections_sup(any()) -> pid() | undefined.
+lookup_connections_sup(Ref) ->
+	ets:lookup_element(?TAB, {listener, Ref}, 3).
+
 %% @doc Add an acceptor for the given listener.
 -spec add_acceptor(any(), pid()) -> ok.
 add_acceptor(Ref, Pid) ->
@@ -147,6 +160,11 @@ remove_process(Key = {listener, Ref}, MonitorRef, Pid, Monitors) ->
 	true = ets:delete(?TAB, {connections, Pid}),
 	lists:keydelete({MonitorRef, Pid}, 1, Monitors);
 remove_process(Key = {acceptors, _}, MonitorRef, Pid, Monitors) ->
-	Acceptors = ets:lookup_element(?TAB, Key, 2),
-	true = ets:insert(?TAB, {Key, lists:delete(Pid, Acceptors)}),
+	try
+		Acceptors = ets:lookup_element(?TAB, Key, 2),
+		true = ets:update_element(?TAB, Key, {2, lists:delete(Pid, Acceptors)})
+	catch
+		error:_ ->
+			ok
+	end,
 	lists:keydelete({MonitorRef, Pid}, 1, Monitors).