Improved version of test_tcp

test_tcp.erl

-module(test_tcp).

-export([run_tcp/3, recv_tcp/3, send_tcp/4]).

-define(SOCKET_OPTS, [{active, false}, binary, {delay_send, true}, {packet, 4}]).

run_tcp(Port, P, N) ->
    spawn_link(fun() ->
        recv_tcp(Port, P, N)
    end),
    timer:sleep(1000),
    send_tcp(localhost, Port, P, N).

recv_tcp(Port, P, N) ->
    {ok, ListenSocket} = gen_tcp:listen(Port, ?SOCKET_OPTS),
    lists:foreach(fun(Idx) ->
        spawn_monitor(fun() ->
            {ok, Socket} = gen_tcp:accept(ListenSocket),
            io:format("receiving ~p/~p~n", [Idx, P]),
            ok = inet:setopts(Socket, [{active, 10000}]),
            recv_n(Socket, N),
            io:format("receiving ~p/~p done~n", [Idx, P]),
            ok = gen_tcp:close(Socket)
        end)
    end, lists:seq(1, P)),
    waitdown(P),
    io:format("closing listen socket~n"),
    ok = gen_tcp:close(ListenSocket).

recv_n(_, 0) -> ok;
recv_n(Socket, N) ->
    receive
      {tcp, Socket, _} -> recv_n(Socket, N-1);
      {tcp_passive, Socket} ->
         ok = inet:setopts(Socket, [{active, 10000}]),
         recv_n(Socket, N);
      {tcp_closed, Socket} ->
        io:format("Closed!~n"),
        exit(closed);
      Otherwise ->
        io:format("Unknown: ~p~n", [Otherwise])
    end.

send_tcp(Host, Port, P, N) ->
    lists:foreach(fun(Idx) ->
        spawn_monitor(fun() ->
            {ok, Socket} = gen_tcp:connect(Host, Port, ?SOCKET_OPTS),
            io:format("sending ~p/~p~n", [Idx, P]),
            send_n(Socket, N),
            io:format("sending ~p/~p done~n", [Idx, P]),
            ok = gen_tcp:close(Socket)
        end)
    end, lists:seq(1, P)),
    waitdown(P).

send_n(_, 0) -> ok;
send_n(Socket, N) ->
    gen_tcp:send(Socket, term_to_binary(N)),
    send_n(Socket, N-1).

waitdown(0) -> ok;
waitdown(Count) ->
    receive
        {'DOWN', _, process, _, normal} ->
            waitdown(Count-1)
    end.