Add a simple implementation of max connections

src/mochiweb_socket_server.erl

@@ -18,7 +18,6 @@
         {port,
          loop,
          name=undefined,
-         %% NOTE: This is currently ignored.
          max=2048,
          ip=any,
          listen=null,
@@ -78,7 +77,16 @@ parse_options(State=#mochiweb_socket_server{}) ->
 parse_options(Options) ->
     parse_options(Options, #mochiweb_socket_server{}).
 
-parse_options([], State) ->
+parse_options([], State=#mochiweb_socket_server{acceptor_pool_size=PoolSize,
+                                                max=Max}) ->
+    case Max < PoolSize of
+        true ->
+            error_logger:info_report([{warning, "max is set lower than acceptor_pool_size"},
+                                      {max, Max},
+                                      {acceptor_pool_size, PoolSize}]);
+        false ->
+            ok
+    end,
     State;
 parse_options([{name, L} | Rest], State) when is_list(L) ->
     Name = {local, list_to_atom(L)},
@@ -117,8 +125,6 @@ parse_options([{acceptor_pool_size, Max} | Rest], State) ->
     parse_options(Rest,
                   State#mochiweb_socket_server{acceptor_pool_size=MaxInt});
 parse_options([{max, Max} | Rest], State) ->
-    error_logger:info_report([{warning, "TODO: max is currently unsupported"},
-                              {max, Max}]),
     MaxInt = ensure_int(Max),
     parse_options(Rest, State#mochiweb_socket_server{max=MaxInt});
 parse_options([{ssl, Ssl} | Rest], State) when is_boolean(Ssl) ->
@@ -206,6 +212,8 @@ listen(Port, Opts, State=#mochiweb_socket_server{ssl=Ssl, ssl_opts=SslOpts}) ->
 
 do_get(port, #mochiweb_socket_server{port=Port}) ->
     Port;
+do_get(waiting_acceptors, #mochiweb_socket_server{acceptor_pool=Pool}) ->
+    sets:size(Pool);
 do_get(active_sockets, #mochiweb_socket_server{active_sockets=ActiveSockets}) ->
     ActiveSockets.
 
@@ -273,16 +281,33 @@ code_change(_OldVsn, State, _Extra) ->
 
 recycle_acceptor(Pid, State=#mochiweb_socket_server{
                         acceptor_pool=Pool,
+                        acceptor_pool_size=PoolSize,
                         listen=Listen,
                         loop=Loop,
+                        max=Max,
                         active_sockets=ActiveSockets}) ->
     case sets:is_element(Pid, Pool) of
         true ->
-            Acceptor = mochiweb_acceptor:start_link(self(), Listen, Loop),
-            Pool1 = sets:add_element(Acceptor, sets:del_element(Pid, Pool)),
-            State#mochiweb_socket_server{acceptor_pool=Pool1};
+            Pool1 = sets:del_element(Pid, Pool),
+            case ActiveSockets + sets:size(Pool1) < Max of
+                true ->
+                    Acceptor = mochiweb_acceptor:start_link(self(), Listen, Loop),
+                    Pool2 = sets:add_element(Acceptor, Pool1),
+                    State#mochiweb_socket_server{acceptor_pool=Pool2};
+                false ->
+                    State#mochiweb_socket_server{acceptor_pool=Pool1}
+            end;
         false ->
-            State#mochiweb_socket_server{active_sockets=ActiveSockets - 1}
+            Pool1 =
+                case sets:size(Pool) < PoolSize of
+                    true ->
+                        Acceptor = mochiweb_acceptor:start_link(self(), Listen, Loop),
+                        sets:add_element(Acceptor, Pool);
+                    false ->
+                        Pool
+                end,
+            State#mochiweb_socket_server{active_sockets=ActiveSockets - 1,
+                                         acceptor_pool=Pool1}
     end.
 
 handle_info(Msg, State) when ?is_old_state(State) ->

src/mochiweb_socket_server_tests.erl

@@ -0,0 +1,155 @@
+-module(mochiweb_socket_server_tests).
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+socket_server(Opts, ServerFun) ->
+    ServerOpts = [{ip, "127.0.0.1"}, {port, 0}, {backlog, 5}, {loop, ServerFun}],
+    {ok, Server} = mochiweb_socket_server:start(ServerOpts ++ Opts),
+    Port = mochiweb_socket_server:get(Server, port),
+    {Server, Port}.
+
+echo_loop(Socket) ->
+    ok = mochiweb_socket:setopts(Socket, [{active, once}]),
+    receive
+        {_Protocol, _, Data} ->
+            gen_tcp:send(Socket, Data),
+            echo_loop(Socket);
+        {tcp_closed, Socket} ->
+            ok
+    end.
+
+start_client_conns(Port, NumClients, ClientFun, ClientArgs, Tester) ->
+    Opts = [binary, {active, false}, {packet, 1}],
+    lists:foreach(fun (_N) ->
+                          case gen_tcp:connect("127.0.0.1", Port, Opts) of
+                              {ok, Socket} ->
+                                  spawn_link(fun() -> ClientFun(Socket, ClientArgs) end);
+                              {error, E} ->
+                                  Tester ! {client_conn_error, E}
+                          end
+                  end, lists:seq(1, NumClients)).
+
+client_fun(_Socket, []) -> ok;
+client_fun(Socket, [{close_sock} | Cmds]) ->
+    mochiweb_socket:close(Socket),
+    client_fun(Socket, Cmds);
+client_fun(Socket, [{send_pid, To} | Cmds]) ->
+    To ! {client, self()},
+    client_fun(Socket, Cmds);
+client_fun(Socket, [{send, Data, Tester} | Cmds]) ->
+    case gen_tcp:send(Socket, Data) of
+        ok -> ok;
+        {error, E} -> Tester ! {client_send_error, self(), E}
+    end,
+    client_fun(Socket, Cmds);
+client_fun(Socket, [{recv, Length, Timeout, Tester} | Cmds]) ->
+    case gen_tcp:recv(Socket, Length, Timeout) of
+        {ok, _} -> ok;
+        {error, E} -> Tester ! {client_recv_error, self(), E}
+    end,
+    client_fun(Socket, Cmds);
+client_fun(Socket, [{wait_msg, Msg} | Cmds]) ->
+    receive
+        M when M =:= Msg -> ok
+    end,
+    client_fun(Socket, Cmds);
+client_fun(Socket, [{send_msg, Msg, To} | Cmds]) ->
+    To ! {Msg, self()},
+    client_fun(Socket, Cmds).
+
+test_basic_accept(Max, PoolSize, NumClients, ReportTo) ->
+    Tester = self(),
+
+    ServerOpts = [{max, Max}, {acceptor_pool_size, PoolSize}],
+    ServerLoop =
+        fun (Socket) ->
+                Tester ! {server_accepted, self()},
+                mochiweb_socket:setopts(Socket, [{packet, 1}]),
+                echo_loop(Socket)
+        end,
+    {Server, Port} = socket_server(ServerOpts, ServerLoop),
+
+    Data = <<"data">>,
+    Timeout = 2000,
+    ClientCmds = [{send_pid, Tester}, {wait_msg, go},
+                  {send, Data, Tester}, {recv, size(Data), Timeout, Tester},
+                  {close_sock}, {send_msg, done, Tester}],
+    start_client_conns(Port, NumClients, fun client_fun/2, ClientCmds, Tester),
+    % Wait a bit for the connections to complete.
+    timer:sleep(2500),
+    ConnectLoop =
+        fun (Loop, Connected, Accepted, Errors) ->
+                case (length(Connected) + Errors >= NumClients) of
+                    true -> {Connected, Accepted};
+                    false ->
+                        receive
+                            {server_accepted, ServerPid} ->
+                                Loop(Loop, Connected, [ServerPid | Accepted], Errors);
+                            {client, ClientPid} ->
+                                Loop(Loop, [ClientPid | Connected], Accepted, Errors);
+                            {client_conn_error, _E} ->
+                                Loop(Loop, Connected, Accepted, Errors + 1)
+                        end
+                end
+        end,
+    {Connected, Accepted} = ConnectLoop(ConnectLoop, [], [], 0),
+
+    ActiveAfterConnect = mochiweb_socket_server:get(Server, active_sockets),
+    WaitingAfterConnect = mochiweb_socket_server:get(Server, waiting_acceptors),
+
+    lists:foreach(fun(Pid) -> Pid ! go end, Connected),
+    WaitLoop =
+        fun (Loop, Done) ->
+                case (length(Done) >= length(Connected)) of
+                    true ->
+                        ok;
+                    false ->
+                        receive
+                            {done, From} ->
+                                Loop(Loop, [From | Done])
+                        end
+                end
+        end,
+    ok = WaitLoop(WaitLoop, []),
+
+    % Wait a little for the server side connections to get their
+    % socket close notifications.
+    timer:sleep(1500),
+
+    ActiveAtEnd = mochiweb_socket_server:get(Server, active_sockets),
+    WaitingAtEnd = mochiweb_socket_server:get(Server, waiting_acceptors),
+    mochiweb_socket_server:stop(Server),
+
+    ReportTo ! {result, {length(Accepted),
+                         ActiveAfterConnect, WaitingAfterConnect,
+                         ActiveAtEnd, WaitingAtEnd}}.
+
+normal_acceptor_test_fun() ->
+    %        {Max, PoolSize, NumClients,
+    %         {ExpectedAccepts,
+    %          ExpectedActiveAfterConnect, ExpectedWaitingAfterConnect,
+    %          ExpectedActiveAtEnd, ExpectedWaitingAtEnd}}
+    Tests = [{3, 1, 1,   {1,   1, 1,   0, 1}},
+             {3, 1, 2,   {2,   2, 1,   0, 1}},
+             {3, 1, 3,   {3,   3, 0,   0, 1}},
+             {3, 3, 3,   {3,   3, 0,   0, 3}},
+             {1, 3, 3,   {3,   3, 0,   0, 3}}, % Max is overridden to PoolSize
+             {3, 2, 10,  {3,   3, 0,   0, 2}}
+            ],
+    [fun () ->
+             Self = self(),
+             spawn(fun () ->
+                           test_basic_accept(Max, PoolSize, NumClients, Self)
+                   end),
+             Result = receive {result, R} -> R end,
+             ?assertEqual(Expected, Result)
+     end || {Max, PoolSize, NumClients, Expected} <- Tests].
+
+-define(LARGE_TIMEOUT, 40).
+
+normal_acceptor_test_() ->
+    Tests = normal_acceptor_test_fun(),
+    {timeout, ?LARGE_TIMEOUT, Tests}.
+
+-endif.