Carbon driver for basho bench

gistfile1.erl

-module(basho_bench_driver_carbon).

% basho bench needs only two functions
-export([new/1, run/4]).

-include("basho_bench.hrl").

-record(state, {
        host,
        port,
        keep_connect,
        batch_size,
        socket
    }).

new(_Id) ->
    % Fetch configurations
    Host = basho_bench_config:get(carbon_server, "127.0.0.1"),
    Port = basho_bench_config:get(carbon_port, 2003),
    KeepConnect = basho_bench_config:get(carbon_keep_connect, false),
    BatchSize = basho_bench_config:get(carbon_batch_size, 100),
    % You can initialize what you need or start applications
    % Return a state
    {ok, #state{
            host=Host,
            port=Port,
            keep_connect=KeepConnect,
            batch_size=BatchSize,
            socket=nil}}.

carbon(State = #state{host=Host, port=Port, socket=nil}, Message) ->
    % Not connected. Just connect and play again
    case gen_tcp:connect(Host, Port, [list, {packet, 0}]) of
        {ok, Sock} ->
            carbon(State#state{socket=Sock}, Message);
        Error ->
            Error
    end;

carbon(State = #state{socket=Socket, keep_connect=KeepConnect}, Message) ->
    case gen_tcp:send(Socket, Message) of
        ok ->
            case KeepConnect of
                % close connection or keep it open?
                true ->
                    {ok, State};
                _ ->
                    ok = gen_tcp:close(Socket),
                    {ok, State#state{socket=nil}}
            end;
        Error ->
            Error % let it crash
    end.

concat(0, _Keygen, _Ts, List) ->
    List;

concat(Size, KeyGen, Ts, List) ->
    % Carbon text protocol is simple : key timestamp value\n
    Msg = io_lib:format("pim.pam.poum.~p ~p ~p~n", [KeyGen(), Ts,  random:uniform(1000)]),
    % Erlang way to concatenate lines : recursivity
    concat(Size -1, KeyGen, Ts, [Msg | List]).

run(set, KeyGen, _ValueGen, State = #state{batch_size=BatchSize}) ->
    % set : for now, only one action, setting values.
    % KeyGen : helper for generating key name. I use a pareto distribution in the conf file.
    % _ValueGen : helper provides binary, I don't need that, random integer is enough.
    % State : the memory of this process.
    {Mega, Sec, _Micro} = now(),
    % Build message
    Msg = concat(BatchSize, KeyGen, Mega * 1000 + Sec, []),

    % Send message
    case carbon(State, Msg) of
        {error, E} ->
            {error, E, State};
        OK ->
            OK
    end.