bench.escript

#!/usr/bin/env escript

-include_lib("kernel/include/inet.hrl").
-include_lib("kernel/include/inet_sctp.hrl").

main([Protocol, IP, Port, Len, Cnt]) ->
    bench(list_to_atom(Protocol), inet:parse_address(IP),
	  list_to_integer(Port), list_to_integer(Len), list_to_integer(Cnt));
main(_) ->
    io:format("invalid arguments~n~n"),
    usage(),
    ok.

usage() ->
    io:format("Usage: bench < tcp | sctp > IP Port Len Count~n").

bench(Protocol, {ok, IP}, Port, Len, Cnt)
  when Protocol == tcp; Protocol == udp; Protocol == sctp ->
    io:format("bench(~p, ~p, ~p, ~p, ~p)~n", [Protocol, IP, Port, Len, Cnt]),
    Data = crypto:strong_rand_bytes(Len),
    Time = start_bench(Protocol, IP, Port, Data, Cnt),
    io:format("~p,~p,~p,~p~n", [Protocol, Len, Cnt, Time]),
    ok;

bench(_, _, _, _, _) ->
    io:format("invalid arguments~n~n"),
    usage(),
    ok.

start_bench(sctp, IP, Port, Data, Cnt) ->
    {ok, Socket} = gen_sctp:open(),
    SOpts = [{sctp_initmsg,#sctp_initmsg{num_ostreams=5}},
	    {recbuf, 1024000}, {sndbuf, 1024000}],

    {ok, Assoc} = gen_sctp:connect(Socket, IP, Port, SOpts),
    {Time, _} = timer:tc(fun() -> sctp_loop(Socket, Assoc, Data, Cnt) end),
    Time;

start_bench(tcp, IP, Port, Data, Cnt) ->
    SOpts = [{recbuf, 1024000}, {sndbuf, 1024000}],

    {ok, Socket} = gen_tcp:connect(IP, Port, SOpts),
    {Time, _} = timer:tc(fun() -> tcp_loop(Socket, Data, Cnt) end),
    Time;

start_bench(udp, IP, Port, Data, Cnt) ->
    SOpts = [{recbuf, 1024000}, {sndbuf, 1024000}],

    {ok, Socket} = gen_udp:open(0, SOpts),
    {Time, _} = timer:tc(fun() -> udp_loop(Socket, IP, Port, Data, Cnt) end),
    Time.

sctp_loop(Socket, _Assoc, _, 0) ->
    gen_sctp:close(Socket),
    ok;
sctp_loop(Socket, Assoc, Data, Cnt) ->
    ok = sctp_send(Socket, Assoc, 0, Data),
    sctp_loop(Socket, Assoc, Data, Cnt - 1).

sctp_send(Socket, Assoc, Stream, Data) ->
    case gen_sctp:send(Socket, Assoc, Stream, Data) of
	ok ->
	    ok;
	{error, eagain} ->
	    sctp_send(Socket, Assoc, Stream, Data);
	{error, Reason} ->
	    {error, Reason}
    end.

tcp_loop(Socket, _, 0) ->
    gen_tcp:close(Socket),
    ok;
tcp_loop(Socket, Data, Cnt) ->
    ok = gen_tcp:send(Socket, Data),
    tcp_loop(Socket, Data, Cnt - 1).

udp_loop(Socket, _IP, _Port, _, 0) ->
    gen_udp:close(Socket),
    ok;
udp_loop(Socket, IP, Port, Data, Cnt) ->
    ok = gen_udp:send(Socket, IP, Port, Data),
    udp_loop(Socket, IP, Port, Data, Cnt - 1).