Introduce the ranch_server registry, make it handle listeners

6 files changed, 110 insertions, 21 deletions

diff --git a/src/ranch.app.src b/src/ranch.app.src
index dd2840c..8d0a208 100644
--- a/src/ranch.app.src
+++ b/src/ranch.app.src
@@ -16,7 +16,7 @@
 	{description, "Socket acceptor pool for TCP protocols."},
 	{vsn, git},
 	{modules, []},
-	{registered, [ranch_sup]},
+	{registered, [ranch_sup, ranch_server]},
 	{applications, [
 		kernel,
 		stdlib
diff --git a/src/ranch.erl b/src/ranch.erl
index 3f07df7..3360154 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -78,7 +78,7 @@ child_spec(Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts)
 		when is_integer(NbAcceptors) andalso is_atom(Transport)
 		andalso is_atom(Protocol) ->
 	{{ranch_listener_sup, Ref}, {ranch_listener_sup, start_link, [
-		NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts
+		Ref, NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts
 	]}, permanent, 5000, supervisor, [ranch_listener_sup]}.
 
 %% @doc Acknowledge the accepted connection.
@@ -92,14 +92,14 @@ accept_ack(ListenerPid) ->
 %% @doc Return the listener's port.
 -spec get_port(any()) -> inet:port_number().
 get_port(Ref) ->
-	ListenerPid = ref_to_listener_pid(Ref),
+	ListenerPid = ranch_server:lookup_listener(Ref),
 	{ok, Port} = ranch_listener:get_port(ListenerPid),
 	Port.
 
 %% @doc Return the current protocol options for the given listener.
 -spec get_protocol_options(any()) -> any().
 get_protocol_options(Ref) ->
-	ListenerPid = ref_to_listener_pid(Ref),
+	ListenerPid = ranch_server:lookup_listener(Ref),
 	{ok, ProtoOpts} = ranch_listener:get_protocol_options(ListenerPid),
 	ProtoOpts.
 
@@ -110,17 +110,5 @@ get_protocol_options(Ref) ->
 %% no effect on the currently opened connections.
 -spec set_protocol_options(any(), any()) -> ok.
 set_protocol_options(Ref, ProtoOpts) ->
-	ListenerPid = ref_to_listener_pid(Ref),
+	ListenerPid = ranch_server:lookup_listener(Ref),
 	ok = ranch_listener:set_protocol_options(ListenerPid, ProtoOpts).
-
-%% Internal.
-
--spec ref_to_listener_pid(any()) -> pid().
-ref_to_listener_pid(Ref) ->
-	Children = supervisor:which_children(ranch_sup),
-	{_, ListenerSupPid, _, _} = lists:keyfind(
-		{ranch_listener_sup, Ref}, 1, Children),
-	ListenerSupChildren = supervisor:which_children(ListenerSupPid),
-	{_, ListenerPid, _, _} = lists:keyfind(
-		ranch_listener, 1, ListenerSupChildren),
-	ListenerPid.
diff --git a/src/ranch_listener_sup.erl b/src/ranch_listener_sup.erl
index 63e3ffe..b42732d 100644
--- a/src/ranch_listener_sup.erl
+++ b/src/ranch_listener_sup.erl
@@ -17,21 +17,22 @@
 -behaviour(supervisor).
 
 %% API.
--export([start_link/5]).
+-export([start_link/6]).
 
 %% supervisor.
 -export([init/1]).
 
 %% API.
 
--spec start_link(non_neg_integer(), module(), any(), module(), any())
+-spec start_link(any(), non_neg_integer(), module(), any(), module(), any())
 	-> {ok, pid()}.
-start_link(NbAcceptors, Transport, TransOpts, Protocol, ProtoOpts) ->
+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, [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]}),
diff --git a/src/ranch_server.erl b/src/ranch_server.erl
new file mode 100644
index 0000000..bafdeb5
--- /dev/null
+++ b/src/ranch_server.erl
@@ -0,0 +1,92 @@
+%% Copyright (c) 2012, Loïc Hoguin <[email protected]>
+%%
+%% 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.
+
+%% @private
+-module(ranch_server).
+-behaviour(gen_server).
+
+%% API.
+-export([start_link/0]).
+-export([insert_listener/2]).
+-export([lookup_listener/1]).
+
+%% gen_server.
+-export([init/1]).
+-export([handle_call/3]).
+-export([handle_cast/2]).
+-export([handle_info/2]).
+-export([terminate/2]).
+-export([code_change/3]).
+
+-define(TAB, ?MODULE).
+
+-type key() :: {listener, any()}.
+-record(state, {
+	monitors = [] :: [{{reference(), pid()}, key()}]
+}).
+
+%% API.
+
+%% @doc Start the ranch_server gen_server.
+-spec start_link() -> {ok, pid()}.
+start_link() ->
+	gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
+
+%% @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}),
+	gen_server:cast(?MODULE, {insert_listener, Ref, Pid}).
+
+%% @doc Lookup a listener in the database.
+-spec lookup_listener(any()) -> pid().
+lookup_listener(Ref) ->
+	ets:lookup_element(?TAB, {listener, Ref}, 2).
+
+%% gen_server.
+
+%% @private
+init([]) ->
+	?TAB = ets:new(?TAB, [
+		ordered_set, public, named_table, {read_concurrency, true}]),
+	{ok, #state{}}.
+
+%% @private
+handle_call(_Request, _From, State) ->
+	{reply, ignore, State}.
+
+%% @private
+handle_cast({insert_listener, Ref, Pid}, State=#state{monitors=Monitors}) ->
+	MonitorRef = erlang:monitor(process, Pid),
+	{noreply, State#state{
+		monitors=[{{MonitorRef, Pid}, {listener, Ref}}|Monitors]}};
+handle_cast(_Request, State) ->
+	{noreply, State}.
+
+%% @private
+handle_info({'DOWN', Ref, process, Pid, _}, State=#state{monitors=Monitors}) ->
+	{_, Key} = lists:keyfind({Ref, Pid}, 1, Monitors),
+	true = ets:delete(?TAB, Key),
+	Monitors2 = lists:keydelete({Ref, Pid}, 1, Monitors),
+	{noreply, State#state{monitors=Monitors2}};
+handle_info(_Info, State) ->
+	{noreply, State}.
+
+%% @private
+terminate(_Reason, _State) ->
+	ok.
+
+%% @private
+code_change(_OldVsn, State, _Extra) ->
+	{ok, State}.
diff --git a/src/ranch_sup.erl b/src/ranch_sup.erl
index 5551f12..ff05194 100644
--- a/src/ranch_sup.erl
+++ b/src/ranch_sup.erl
@@ -33,4 +33,8 @@ start_link() ->
 %% supervisor.
 
 init([]) ->
-	{ok, {{one_for_one, 10, 10}, []}}.
+	Procs = [
+		{ranch_server, {ranch_server, start_link, []},
+			permanent, 5000, worker, [ranch_server]}
+	],
+	{ok, {{one_for_one, 10, 10}, Procs}}.
diff --git a/test/acceptor_SUITE.erl b/test/acceptor_SUITE.erl
index bf74a1d..c1313cc 100644
--- a/test/acceptor_SUITE.erl
+++ b/test/acceptor_SUITE.erl
@@ -81,6 +81,8 @@ ssl_echo(Config) ->
 	{ok, <<"SSL Ranch is working!">>} = ssl:recv(Socket, 21, 1000),
 	ok = ranch:stop_listener(ssl_echo),
 	{error, closed} = ssl:recv(Socket, 0, 1000),
+	%% Make sure the listener stopped.
+	{'EXIT', _} = begin catch ranch:get_port(ssl_echo) end,
 	ok.
 
 %% tcp.
@@ -95,4 +97,6 @@ tcp_echo(_) ->
 	{ok, <<"TCP Ranch is working!">>} = gen_tcp:recv(Socket, 21, 1000),
 	ok = ranch:stop_listener(tcp_echo),
 	{error, closed} = gen_tcp:recv(Socket, 0, 1000),
+	%% Make sure the listener stopped.
+	{'EXIT', _} = begin catch ranch:get_port(tcp_echo) end,
 	ok.