gen_tcp2.erl

-module(gen_tcp2).

-export([listen/2, accept/1, accept/2,
         connect/3, connect/4,
         send/2, recv/1, recv/2, recv/3,
         close/1]).

-export([run/0, run/2]).

%% Naive gen_tcp shim over the socket API
listen(Port, Opts) ->
    {ok, LSock} = socket:open(inet, stream, tcp),
    Addr = proplists:get_value(ip, Opts, any),
    {ok, OutPort} = socket:bind(LSock, #{family => inet,
                                          port   => Port,
                                          addr   => Addr}),
    ok = socket:listen(LSock),
    {ok, LSock, OutPort}.

accept(LSock) ->
    accept(LSock, infinity).
accept(LSock, Timeout) ->
    socket:accept(LSock, Timeout).

connect(SAddr, SPort, Opts) ->
    connect(SAddr, SPort, Opts, infinity).
connect(localhost, SPort, Opts, Timeout) ->
    connect(loopback, SPort, Opts, Timeout);
connect(SAddr, SPort, Opts, Timeout) ->
    {ok, Sock} = socket:open(inet, stream, tcp),
    Addr = proplists:get_value(ip, Opts, any),
    {ok, _} = socket:bind(Sock, #{family => inet,
                                  addr   => Addr}),
    ok = socket:connect(Sock, #{family => inet,
                                addr   => SAddr,
                                port   => SPort}, Timeout),
    {ok, Sock}.

send(Sock, Data) ->
    socket:send(Sock, Data).

recv(Sock) ->
    socket:recv(Sock).
recv(Sock, Length) ->
    socket:recv(Sock, Length).
recv(Sock, Length, Timeout) ->
    socket:recv(Sock, Length, Timeout).

close(Sock) ->
    socket:close(Sock).
    
%% This benchmark sends 128 * 1000000 bytes over loopback
run() ->
    N = 1000000,
    {ClientTS, ServerTS, BytesSent} = run(gen_tcp, N),
    io:format("              client             server~n",[]),
    io:format(" gen_tcp: ~10.1f ns/byte ~10.1f ns/byte~n",
              [ClientTS/BytesSent * 1000, ServerTS/BytesSent * 1000]),
    {ClientTS2, ServerTS2, BytesSent} = run(gen_tcp2, N),
    io:format("gen_tcp2: ~10.1f ns/byte ~10.1f ns/byte~n",
              [ClientTS2/BytesSent * 1000, ServerTS2/BytesSent * 1000]),
    io:format("   ratio: ~10.1f %       ~10.1f%~n",
              [ClientTS2 / ClientTS * 100,
               ServerTS2 / ServerTS * 100]),
    ok.

run(Mod, N) ->
    P = self(),
    Bin = <<0:1024>>,
    NList = lists:seq(1,N),
    spawn_link(
      fun() ->
              case Mod:listen(0, [{active, false}, binary]) of
                  {ok, LS} ->
                      {ok, LP} = inet:port(LS);
                  {ok, LS, LP} ->
                      ok
              end,
              P ! LP,
              {ok, AS} = Mod:accept(LS),
              {TS, BytesRecieved} = timer:tc(fun() ->
                                                     recv_loop(Mod, AS)
                                             end),
              Mod:close(LS),
              P ! {done, TS, BytesRecieved}
      end),
    receive Port -> Port end,
    {ok, S} = Mod:connect(localhost, Port, []),
    {ClientTS, ok} =
        timer:tc(fun() ->
                         [ok = Mod:send(S, Bin) || _ <- NList],
                         Mod:close(S)
                 end),
    receive {done, ServerTS, BytesRecieved} -> ServerTS end,
    {ClientTS, ServerTS, BytesRecieved}.

recv_loop(Mod, S) ->
    recv_loop(Mod, S, 0).
recv_loop(Mod, S, N) ->
    case Mod:recv(S, 0) of
        {ok, Bytes} ->
            recv_loop(Mod, S, N + byte_size(Bytes));
        _E ->
            N
    end.