partisan

-module(partisan_telemetry).


-export([span/3]).



%% =============================================================================
%% API
%% =============================================================================



%% -----------------------------------------------------------------------------
%% @doc This is a re-write of {@link telemetry:span/3} that provides an async
%% implementation of {@link telemetry:execute/3}.
%%
%% The telemetry library's implementation calls the handlers sequentially
%% with the start event before calling `SpanFunction' and then again with the
%% end or exception event after calling it.
%%
%% This is bad as users might inadvertedly attach too many handlers and or slow
%% handlers, affecting the observed execution time of the Partisan function
%% that wants to use this function.
%% @end
%% -----------------------------------------------------------------------------
-spec span(
    telemetry:event_prefix(),
    telemetry:event_metadata(),
    telemetry:span_function()) -> telemetry:span_result().

span(EventPrefix, StartMeta, SpanFunction) ->
    StartTime = erlang:monotonic_time(),
    SysTime = erlang:system_time(),

    try {_, #{}} = SpanFunction() of
        {Result, StopMeta} ->
            StopTime = erlang:monotonic_time(),
            ok = execute(
                EventPrefix, StartMeta, StopMeta, SysTime, StartTime, StopMeta
            ),
            Result
    catch
        ?WITH_STACKTRACE(Class, Reason, Stacktrace)
            StopTime = erlang:monotonic_time(),
            ok = execute(
                EventPrefix, StartMeta, StopMeta, SysTime, StartTime, StopMeta
            ),
            erlang:raise(Class, Reason, Stacktrace)
    end.




%% =============================================================================
%% PRIVATE
%% =============================================================================



%% @private
execute(EventPrefix, StartMeta, StopMeta, SysTime, StartTime, StopTime) ->
    spawn(fun() ->
        DefaultCtx = erlang:make_ref(),

        ok = execute(
            EventPrefix ++ [start],
            #{monotonic_time => StartTime, system_time => SysTime},
            merge_ctx(StartMeta, DefaultCtx)
        ),

        ok = execute(
            EventPrefix ++ [stop],
            #{duration => StopTime - StartTime, monotonic_time => StopTime},
            merge_ctx(StopMeta, DefaultCtx)
        )
    end.


%% -----------------------------------------------------------------------------
%% @private
%% @doc
%% @end
%% -----------------------------------------------------------------------------
-spec merge_ctx(event_metadata(), any()) -> event_metadata().

merge_ctx(#{telemetry_span_context := _} = Metadata, _Ctx) ->
    Metadata;

merge_ctx(Metadata, Ctx) ->
    Metadata#{telemetry_span_context => Ctx}.