feat: process hook tracing

src/ejabberd_hooks.erl

@@ -46,12 +46,23 @@
 	 handle_info/2,
 	 terminate/2]).
 
+
+-export(
+    [
+        get_tracing_options/3,
+        trace_off/3,
+        trace_on/5
+    ]
+).
+
 -include("logger.hrl").
 -include("ejabberd_stacktrace.hrl").
 
 -record(state, {}).
 -type hook() :: {Seq :: integer(), Module :: atom(), Function :: atom() | fun()}.
 
+-define(TRACE_HOOK_KEY, '$trace_hook').
+-define(TIMING_KEY, '$trace_hook_timer').
 %%%----------------------------------------------------------------------
 %%% API
 %%%----------------------------------------------------------------------
@@ -104,7 +115,18 @@ run(Hook, Args) ->
 run(Hook, Host, Args) ->
     try ets:lookup(hooks, {Hook, Host}) of
 	[{_, Ls}] ->
+        case erlang:get(?TRACE_HOOK_KEY) of
+            undefined ->
                 run1(Ls, Hook, Args);
+            TracingHooksOpts ->
+                case do_get_tracing_options(Hook, Host, TracingHooksOpts) of
+                    undefined ->
+                        run1(Ls, Hook, Args);
+                    TracingOpts ->
+                        foreach_start_hook_tracing(TracingOpts, Hook, Host, Args),
+                        run2(Ls, Hook, Args, Host, TracingOpts)
+                end
+        end;
 	[] ->
 	    ok
     catch _:badarg ->
@@ -124,13 +146,75 @@ run_fold(Hook, Val, Args) ->
 run_fold(Hook, Host, Val, Args) ->
     try ets:lookup(hooks, {Hook, Host}) of
 	[{_, Ls}] ->
+        case erlang:get(?TRACE_HOOK_KEY) of
+            undefined ->
                 run_fold1(Ls, Hook, Val, Args);
+            TracingHooksOpts ->
+                case do_get_tracing_options(Hook, Host, TracingHooksOpts) of
+                    undefined ->
+                        run_fold1(Ls, Hook, Val, Args);
+                    TracingOpts ->
+                        fold_start_hook_tracing(TracingOpts, Hook, Host, [Val | Args]),
+                        run_fold2(Ls, Hook, Val, Args, Host, TracingOpts)
+                end
+        end;
 	[] ->
 	    Val
     catch _:badarg ->
 	    Val
     end.
 
+get_tracing_options(Hook, Host, Pid) when Pid == erlang:self() ->
+    do_get_tracing_options(Hook, Host, erlang:get(?TRACE_HOOK_KEY));
+get_tracing_options(Hook, Host, Pid) when erlang:is_pid(Pid) ->
+    case erlang:process_info(Pid, dictionary) of
+        {_, DictPropList} ->
+            case lists:keyfind(?TRACE_HOOK_KEY, 1, DictPropList) of
+                {_, TracingHooksOpts} ->
+                    do_get_tracing_options(Hook, Host, TracingHooksOpts);
+                _ ->
+                    undefined
+            end;
+        _ ->
+            undefined
+    end.
+
+trace_on(Hook, Host, Pid, #{}=Opts, Timeout) when Pid == erlang:self() ->
+    do_trace_on(Hook, Host, Opts, Timeout);
+trace_on(Hook, Host, Proc, #{}=Opts, Timeout) ->
+    try sys:replace_state(
+        Proc,
+        fun(State) ->
+            do_trace_on(Hook, Host, Opts, Timeout),
+            State
+        end,
+        15000
+    ) of
+        _ -> % process state
+            ok
+    catch
+        _:Reason ->
+            {error, Reason}
+    end.
+
+trace_off(Hook, Host, Pid) when Pid == erlang:self() ->
+    do_trace_off(Hook, Host);
+trace_off(Hook, Host, Proc) ->
+    try sys:replace_state(
+        Proc,
+        fun(State) ->
+            do_trace_off(Hook, Host),
+            State
+        end,
+        15000
+    ) of
+        _ -> % process state
+            ok
+    catch
+        _:Reason ->
+            {error, Reason}
+    end.
+
 %%%----------------------------------------------------------------------
 %%% Callback functions from gen_server
 %%%----------------------------------------------------------------------
@@ -247,3 +331,323 @@ safe_apply(Hook, Module, Function, Args) ->
 			misc:format_exception(2, E, R, Stack)|Args]),
 	    'EXIT'
     end.
+
+%%%----------------------------------------------------------------------
+%%% Internal tracing functions
+%%%----------------------------------------------------------------------
+
+do_trace_on(Hook, Host, Opts, undefined) ->
+    case erlang:get(?TRACE_HOOK_KEY) of
+        _ when Hook == '_' andalso Host == <<"*">> ->
+            % Trace everything:
+            erlang:put(?TRACE_HOOK_KEY, #{'_' => #{<<"*">> => Opts}});
+        #{'_' := #{<<"*">> := _}} -> % Already tracing everything
+            % Update Opts:
+            erlang:put(?TRACE_HOOK_KEY, #{'_' => #{<<"*">> => Opts}});
+        #{'_' := HostOpts} when Hook == '_' -> % Already Tracing everything for some hosts
+            % Add/Update Host and Opts:
+            erlang:put(?TRACE_HOOK_KEY, #{'_' => HostOpts#{Host => Opts}});
+        #{'_' := _} -> % Already tracing everything and Hook is not '_'
+            ok;
+        #{} when Hook == '_' ->
+            % Remove other hooks by just adding '_':
+            erlang:put(?TRACE_HOOK_KEY, #{'_' => #{Host => Opts}});
+        #{}=TraceHooksOpts when Host == <<"*">> -> % Want to trace a hook for all hosts
+            erlang:put(?TRACE_HOOK_KEY, TraceHooksOpts#{Hook => #{Host => Opts}});
+        #{}=TraceHooksOpts ->
+            case maps:get(Hook, TraceHooksOpts, #{}) of
+                #{<<"*">> := _} -> % Already tracing this hook for all hosts
+                    ok;
+                HostOpts ->
+                    erlang:put(?TRACE_HOOK_KEY, TraceHooksOpts#{Hook => HostOpts#{Host => Opts}})
+            end;
+        undefined ->
+            erlang:put(?TRACE_HOOK_KEY, #{Hook => #{Host => Opts}})
+    end,
+    ok;
+do_trace_on(Hook, Host, Opts, Timeout) -> % Trace myself `Timeout` time
+    ParentPid = erlang:self(),
+    try erlang:spawn(
+        fun() ->
+            MonitorRef = erlang:monitor(process, ParentPid),
+            receive
+                {_, MonitorRef, _, _, _} ->
+                    ok
+            after Timeout ->
+                trace_off(Hook, Host, ParentPid)
+            end,
+            erlang:exit(normal)
+        end
+    ) of
+        _ ->
+            do_trace_on(Hook, Host, Opts, undefined) % ok
+    catch
+        _:Reason -> % system_limit
+            {error, Reason}
+    end.
+
+do_trace_off(Hook, Host) ->
+    case erlang:get(?TRACE_HOOK_KEY) of
+        _ when Hook == '_' andalso Host == <<"*">> ->
+            % Remove all tracing:
+            erlang:erase(?TRACE_HOOK_KEY);
+        #{'_' := HostOpts} when Hook == '_' -> % Already tracing all hooks
+            % Remove Host:
+            HostOpts2 = maps:remove(Host, HostOpts),
+            if
+                HostOpts2 == #{} ->
+                    % Remove all tracing:
+                    erlang:erase(?TRACE_HOOK_KEY);
+                true ->
+                    erlang:put(?TRACE_HOOK_KEY, #{'_' => HostOpts2})
+            end;
+        #{}=TraceHooksOpts when Host == <<"*">> ->
+            % Remove tracing of this hook for all hosts:
+            TraceHooksOpts2 = maps:remove(Hook, TraceHooksOpts),
+            if
+                TraceHooksOpts2 == #{} ->
+                    % Remove all tracing:
+                    erlang:erase(?TRACE_HOOK_KEY);
+                true ->
+                    erlang:put(?TRACE_HOOK_KEY, TraceHooksOpts2)
+            end;
+        #{}=TraceHooksOpts ->
+            case maps:get(Hook, TraceHooksOpts, undefined) of
+                #{}=HostOpts ->
+                    NewHostOpts = maps:remove(Host, HostOpts),
+                    if
+                        NewHostOpts == #{} ->
+                            % Remove hook:
+                            erlang:put(?TRACE_HOOK_KEY, maps:remove(Hook, TraceHooksOpts));
+                        true ->
+                            erlang:put(?TRACE_HOOK_KEY, TraceHooksOpts#{Hook => NewHostOpts})
+                    end;
+                _ ->
+                    ok
+            end;
+        undefined ->
+            ok
+    end,
+    ok.
+
+do_get_tracing_options(Hook, Host, MaybeMap) ->
+    case MaybeMap of
+        undefined ->
+            undefined;
+        #{'_' := #{<<"*">> := Opts}} -> % Tracing everything
+            Opts;
+        #{'_' := HostOpts} -> % Tracing all hooks for some hosts
+            maps:get(Host, HostOpts, undefined);
+        #{}=TraceHooksOpts ->
+            HostOpts = maps:get(Hook, TraceHooksOpts, #{}),
+            case maps:get(Host, HostOpts, undefined) of
+                undefined ->
+                    maps:get(<<"*">>, HostOpts, undefined);
+                Opts ->
+                    Opts
+            end
+    end.
+
+run2([], Hook, Args, Host, Opts) ->
+    foreach_stop_hook_tracing(Opts, Hook, Host, Args, undefined),
+    ok;
+run2([{Seq, Module, Function} | Ls], Hook, Args, Host, TracingOpts) ->
+    foreach_start_callback_tracing(TracingOpts, Hook, Host, Module, Function, Args, Seq),
+    Res = safe_apply(Hook, Module, Function, Args),
+    foreach_stop_callback_tracing(TracingOpts, Hook, Host, Module, Function, Args, Seq, Res),
+    case Res of
+        'EXIT' ->
+            run2(Ls, Hook, Args, Host, TracingOpts);
+        stop ->
+            foreach_stop_hook_tracing(TracingOpts, Hook, Host, Args, {Module, Function, Seq, Ls}),
+            ok;
+        _ ->
+            run2(Ls, Hook, Args, Host, TracingOpts)
+    end.
+
+run_fold2([], Hook, Val, Args, Host, Opts) ->
+    fold_stop_hook_tracing(Opts, Hook, Host, [Val | Args], undefined),
+    Val;
+run_fold2([{Seq, Module, Function} | Ls], Hook, Val, Args, Host, TracingOpts) ->
+    fold_start_callback_tracing(TracingOpts, Hook, Host, Module, Function, [Val | Args], Seq),
+    Res = safe_apply(Hook, Module, Function, [Val | Args]),
+    fold_stop_callback_tracing(TracingOpts, Hook, Host, Module, Function, [Val | Args], Seq, Res),
+    case Res of
+        'EXIT' ->
+            run_fold2(Ls, Hook, Val, Args, Host, TracingOpts);
+        stop ->
+            fold_stop_hook_tracing(TracingOpts, Hook, Host, [Val | Args], {Module, Function, Seq, {old, Val}, Ls}),
+            Val;
+        {stop, NewVal} ->
+            fold_stop_hook_tracing(TracingOpts, Hook, Host, [Val | Args], {Module, Function, Seq, {new, NewVal}, Ls}),
+            NewVal;
+        NewVal ->
+            run_fold2(Ls, Hook, NewVal, Args, Host, TracingOpts)
+    end.
+
+foreach_start_hook_tracing(TracingOpts, Hook, Host, Args) ->
+    run_event_handler(TracingOpts, Hook, Host, start_hook, [Args], foreach).
+
+foreach_stop_hook_tracing(TracingOpts, Hook, Host, Args, BreakCallback) ->
+    run_event_handler(TracingOpts, Hook, Host, stop_hook, [Args, BreakCallback], foreach).
+
+foreach_start_callback_tracing(TracingOpts, Hook, Host, Mod, Func, Args, Seq) ->
+    run_event_handler(TracingOpts, Hook, Host, start_callback, [Mod, Func, Args, Seq], foreach).
+
+foreach_stop_callback_tracing(TracingOpts, Hook, Host, Mod, Func, Args, Seq, Res) ->
+    run_event_handler(TracingOpts, Hook, Host, stop_callback, [Mod, Func, Args, Seq, Res], foreach).
+
+fold_start_hook_tracing(TracingOpts, Hook, Host, Args) ->
+    run_event_handler(TracingOpts, Hook, Host, start_hook, [Args], fold).
+
+fold_stop_hook_tracing(TracingOpts, Hook, Host, Args, BreakCallback) ->
+    run_event_handler(TracingOpts, Hook, Host, stop_hook, [Args, BreakCallback], fold).
+
+fold_start_callback_tracing(TracingOpts, Hook, Host, Mod, Func, Args, Seq) ->
+    run_event_handler(TracingOpts, Hook, Host, start_callback, [Mod, Func, Args, Seq], fold).
+
+fold_stop_callback_tracing(TracingOpts, Hook, Host, Mod, Func, Args, Seq, Res) ->
+    run_event_handler(TracingOpts, Hook, Host, stop_callback, [Mod, Func, Args, Seq, Res], fold).
+
+run_event_handler(TracingOpts, Hook, Host, Event, EventArgs, RunType) ->
+    EventHandler = maps:get(event_handler, TracingOpts, fun default_tracing_event_handler/6),
+    EventHandlerOpts = maps:get(event_handler_options, TracingOpts, #{}),
+    try
+        if
+            erlang:is_function(EventHandler) ->
+                erlang:apply(EventHandler, [Event, EventArgs, RunType, Hook, Host, EventHandlerOpts]);
+            true ->
+                EventHandler:handle_hook_tracing_event(Event, EventArgs, RunType, Hook, Host, EventHandlerOpts)
+        end
+    of
+        _ ->
+            ok
+    catch
+        ?EX_RULE(E, R, St) ->
+            ?ERROR_MSG("Tracing event handler exception(~0p): ~0p: ~0p", [E, R, St]),
+            ok
+    end.
+
+default_tracing_event_handler(Event, EventArgs, RunType, Hook, Host, EventHandlerOpts) ->
+    % For now:
+    case maps:get(timing, EventHandlerOpts, true) of
+        true ->
+            tracing_timing_event_handler(Event, EventArgs, RunType, Hook, Host, EventHandlerOpts);
+        _ ->
+            ok
+    end.
+
+tracing_timing_event_handler(start_hook, _, _, Hook, Host, EventHandlerOpts) ->
+    HookStart = erlang:system_time(nanosecond),
+    tracing_output(EventHandlerOpts, "(~0p|~ts) Timing started\n", [Hook, Host]),
+    case erlang:get(?TIMING_KEY) of
+        #{}=TimingMap ->
+            erlang:put(?TIMING_KEY, TimingMap#{{Hook, Host} => {HookStart, []}});
+        _ ->
+            erlang:put(?TIMING_KEY, #{{Hook, Host} => {HookStart, []}})
+    end,
+    ok;
+tracing_timing_event_handler(stop_hook, _, _, Hook, Host, EventHandlerOpts) ->
+    HookStop = erlang:system_time(nanosecond),
+    TimingMap = #{} = erlang:get(?TIMING_KEY),
+    {HookStart, CallbackList} = maps:get({Hook, Host}, TimingMap),
+    if
+        erlang:length(CallbackList) < 2 -> % We don't need sorted timing result
+            ok;
+        true ->
+            {CallbackListText, CallbackListTotal} =
+                lists:foldl(
+                    fun({Mod, Func, Arity, Diff}, {CallbackText, Total}) ->
+                        {
+                            CallbackText
+                                ++ "\n\t"
+                                ++ mfa_string({Mod, Func, Arity})
+                                ++ " -> "
+                                ++ erlang:integer_to_list(Diff)
+                                ++ "ns",
+                            Total + Diff
+                        }
+                    end,
+                    {"", 0},
+                    lists:keysort(4, CallbackList)
+                ),
+            tracing_output(
+                EventHandlerOpts,
+                "(~0p|~ts) All callbacks got ~0pns to run. Sorted running time:"
+                    ++ CallbackListText
+                    ++ "\n",
+                [Hook, Host, CallbackListTotal]
+            ),
+            tracing_output(
+                EventHandlerOpts,
+                "(~0p|~ts) Time calculations for all callbacks got ~0pns\n",
+                [Hook, Host, (HookStop - HookStart) - CallbackListTotal]
+            )
+    end,
+    tracing_output(EventHandlerOpts, "(~0p|~ts) Timing stopped\n", [Hook, Host]),
+    NewTimingMap = maps:remove({Hook, Host}, TimingMap),
+    if
+        NewTimingMap == #{} ->
+            erlang:erase(?TIMING_KEY);
+        true ->
+            erlang:put(?TIMING_KEY, NewTimingMap)
+    end,
+    ok;
+tracing_timing_event_handler(start_callback, [Mod, Func, Args | _], _, Hook, Host, _) ->
+    CallbackStart = erlang:system_time(nanosecond),
+    TimingMap = #{} = erlang:get(?TIMING_KEY),
+    {HookStart, Callbacks} = maps:get({Hook, Host}, TimingMap),
+    erlang:put(
+        ?TIMING_KEY,
+        TimingMap#{
+            {Hook, Host} => {HookStart, [{Mod, Func, erlang:length(Args), CallbackStart} | Callbacks]}
+        }
+    ),
+    ok;
+tracing_timing_event_handler(stop_callback, _, _, Hook, Host, EventHandlerOpts) ->
+    CallbackStop = erlang:system_time(nanosecond),
+    TimingMap = #{} = erlang:get(?TIMING_KEY),
+    {HookStart, [{Mod, Func, Arity, CallbackStart} | Callbacks]} = maps:get({Hook, Host}, TimingMap),
+    tracing_output(
+        EventHandlerOpts,
+        "(~0p|~ts) " ++
+            mfa_string({Mod, Func, Arity}) ++
+            " got " ++
+            erlang:integer_to_list(CallbackStop - CallbackStart) ++
+            "ns\n",
+        [Hook, Host]
+    ),
+    erlang:put(
+        ?TIMING_KEY,
+        TimingMap#{
+            {Hook, Host} => {HookStart, [{Mod, Func, Arity, CallbackStop - CallbackStart} | Callbacks]}
+        }
+    ),
+    ok.
+
+tracing_output(#{output_function := OutputF}, Text, Args) ->
+    try
+        OutputF(Text, Args)
+    of
+        _ ->
+            ok
+    catch
+        ?EX_RULE(E, R, St) ->
+            ?ERROR_MSG("Tracing output function exception(~0p): ~0p: ~0p", [E, R, St]),
+            ok
+    end;
+tracing_output(#{output_log_level := Output}, Text, Args) ->
+    if
+        Output == debug ->
+            ?DEBUG(Text, Args);
+        true -> % info
+            ?INFO_MSG(Text, Args)
+    end,
+    ok;
+tracing_output(Opts, Text, Args) ->
+    tracing_output(Opts#{output_log_level => info}, Text, Args).
+
+mfa_string({_, Fun, _}) when erlang:is_function(Fun) ->
+    io_lib:format("~0p", [Fun]);
+mfa_string({Mod, Func, Arity}) ->
+    erlang:atom_to_list(Mod) ++ ":" ++ erlang:atom_to_list(Func) ++ "/" ++ erlang:integer_to_list(Arity).