ktrace.erl

-module(ktrace).
-compile(export_all).

-record(ps, {stack = [],
             trace = array:new(),
             depth = 0}).

-record(state, {pinfo = dict:new(),
                ts}).

c() ->
    State = #state{},
    dbg:trace_client(ip, {"localhost", 4711}, {fun ?MODULE:t/2, State}).

s() ->
    %% s(all).
    s(vnodes).


s(vnodes) ->
    VNs = [Pid || {_,_,Pid} <- riak_core_vnode_manager:all_vnodes(riak_kv_vnode)],
    dbg:tracer(port, dbg:trace_port(ip,{4711,10000000})),
    [dbg:p(VN, [call, arity, return_to, timestamp]) || VN <- VNs],
    dbg:tpl('_', '_', []),
    erlang:trace_pattern({riak_kv_vnode, do_get, 4},
                         [{['_',{'$1','$2'},'_','_'],
                           [],
                           [{message,{{'$1','$2'}}}]}],
                         [local]),
    ok;

s(Who) ->
    dbg:tracer(port, dbg:trace_port(ip,4711)),
    dbg:p(Who, [call, arity, return_to, timestamp]),
    dbg:tpl('_', '_', []).

ss() ->
    dbg:stop_clear().

t({trace_ts, Pid, Type, Args, TS}, State) ->
    t({trace, Pid, Type, Args}, State#state{ts=TS});
t({trace_ts, Pid, Type, Args1, Args2, TS}, State) ->
    t({trace, Pid, Type, Args1, Args2}, State#state{ts=TS});

t({trace, Pid, Type, Args}, State) ->
    tt(Type, Pid, Args, State);
t({trace, Pid, Type, Args1, Args2}, State) ->
    tt(Type, Pid, {Args1, Args2}, State);
t({trace, Pid, Type, Args1, Args2, Args3}, State) ->
    tt(Type, Pid, {Args1, Args2, Args3}, State);
t({drop, Num}, State) ->
    io:format("dropped: ~p~n", [Num]),
    State;
t(end_of_trace, State) ->
    State.

tt(call, Pid, MFA={_M,_F,_A}, State) ->
    %% State = maybe_reset_stack(Pid, MFA, State0),
    {Pos, State1} = add_trace(Pid, MFA, undefined, State),
    State2 = push_stack(Pid, MFA, Pos, State1),
    State2;
tt(call, Pid, {MFA={_M,_F,_A}, Extra}, State0) ->
    State = maybe_reset_stack(Pid, MFA, State0),
    {Pos, State1} = add_trace(Pid, MFA, Extra, State),
    State2 = push_stack(Pid, MFA, Pos, State1),
    State2;
tt(return_to, Pid, MFA, State0) ->
    State = maybe_reset_stack(Pid, MFA, State0),
    {Fun, State2} = pop_stack_until(Pid, MFA, State),
    State3 = rrr(Pid, Fun, State2),
    State3;
tt(_, _Pid, _, State) ->
    State.

rrr(_, undefined, State) ->
    State;
rrr(Pid, {_MFA={_M,_F,_}, _Then, Pos, _Depth}, State) ->
    State2 = update_trace(Pid, Pos, State), 
    State2;

rrr(_Pid, {_MFA, _TS, _, _}, State) ->
    State.

update_trace(Pid, Pos, State) ->
    PS = pinfo(Pid, State),
    T1 = PS#ps.trace,
    {MFA, Then, _, _, Depth, Extra} = array:get(Pos, T1),
    Diff = timer:now_diff(State#state.ts, Then),
    T2 = array:set(Pos, {MFA, Then, State#state.ts, Diff, Depth, Extra}, T1),
    PS2 = PS#ps{trace=T2},
    PI = dict:store(Pid, PS2, State#state.pinfo),
    State#state{pinfo=PI}.

add_trace(Pid, MFA, Extra, State=#state{pinfo=SS,ts=TS}) ->
    DT = array_append({MFA,TS,TS,unknown,0,Extra}, array:new()),
    SS2 = dict:update(Pid,
                      fun(PS=#ps{trace=T,depth=Depth}) ->
                              T2 = array_append({MFA,TS,TS,unknown,Depth,Extra},T),
                              PS#ps{trace=T2}
                      end,
                      #ps{trace=DT}, SS),
    State2 = State#state{pinfo=SS2},
    PS = pinfo(Pid, State2),
    Pos = array:size(PS#ps.trace) - 1,
    {Pos, State2}.

push_stack(Pid, MFA, Pos, State=#state{pinfo=SS,ts=TS}) ->
    SS2 = dict:update(Pid,
                      fun(PS=#ps{stack=S,depth=Depth}) ->
                              S2=[{MFA,TS,Pos,Depth}|S],
                              PS#ps{stack=S2,depth=Depth+1}
                      end,
                      #ps{stack=[{MFA,TS,Pos,0}]}, SS),
    State#state{pinfo=SS2}.

pop_stack_until(Pid, MFA, State=#state{pinfo=SS}) ->
    PS = pinfo(Pid, State),
    case PS#ps.stack of
        [] ->
            Stack = [],
            Top = undefined;
        [H|T] ->
            Stack = T,
            Top = H
    end,
    Stack2 = lists:dropwhile(fun({MFA2,_,_,_}) ->
                                     MFA2 /= MFA
                             end, Stack),
    Depth = case Stack2 of
                [] ->
                    0;
                [{_,_,_,D}|_] ->
                    D
            end,
    PS2 = PS#ps{stack=Stack2, depth=Depth+1},
    SS2 = dict:store(Pid, PS2, SS),
    State2 = State#state{pinfo=SS2},
    {Top, State2}.

maybe_reset_stack(Pid, {gen_fsm, handle_msg,_}, State=#state{pinfo=SS}) ->
    PI = pinfo(Pid, State),
    Trace = PI#ps.trace,
    Size = array:size(Trace),
    case Size > 1 of
        true ->
            {_, Start, _, _, _, _} = array:get(1, Trace),
            End = State#state.ts;
        false ->
            Start = {0,0,0},
            End = {0,0,0}
    end,

    Diff = timer:now_diff(End, Start),
    if Diff > 5000 ->
    %% if Diff > 1 ->
            io:format("*** ~p  :: ~p :: ~p~n", [Diff, Pid, node(Pid)]),
            print_trace(PI#ps.trace);
       true ->
            ok
    end,
    PI2 = PI#ps{stack=[], trace=array:new(), depth=0},
    SS2 = dict:store(Pid, PI2, SS),
    State#state{pinfo=SS2};
maybe_reset_stack(_Pid, _MFA, State) ->
    State.

print_trace(Trace) ->
    array:foldl(fun(_, {MFA,_,_,Time,Depth, Extra}, _) ->
                        ExtraOut = case Extra of
                                       undefined ->
                                           "";
                                       _ ->
                                           io_lib:format(" :: ~p", [Extra])
                                   end,
                        case Time of
                            unknown ->
                                io:format("~10s  |  ~*s~p~s~n", ["", Depth, "", MFA, ExtraOut]);
                            _ ->
                                io:format("~10s  |  ~*s~p~s~n", [integer_to_list(Time), Depth, "", MFA, ExtraOut])
                        end
                end, ok, Trace).

get_stack(Pid, State) ->
    PS = pinfo(Pid, State),
    [MFA || {MFA, _TS, _, _} <- PS#ps.stack].

get_stack2(Pid, State) ->
    PS = pinfo(Pid, State),
    PS#ps.stack.

dict_find(Key, Default, Dict) ->
    case dict:find(Key, Dict) of
        {ok, Value} ->
            Value;
        _ ->
            Default
    end.

pinfo(Pid, #state{pinfo=PI}) ->
    dict_find(Pid, #ps{}, PI).

array_append(Item, A) ->
    array:set(array:size(A), Item, A).

%% Connect to relevant Riak console, add ktrace to it's code path
code:add_path(PATH_TO_KTRACE).
l(ktrace).

%% Active tracing on this node
ktrace:s().

%% Start a separate erl node (easiest to just launch in ktrace dir)
%% Also, start trace client
l(ktrace).
ktrace:c().

%% Do stuff to the Riak node that triggers K/V vnode requests. 
%% Should see trace output printed to the trace client erl console