Make eflambe:capture/3 synchronous, return flamegraph filename

README.adoc

@@ -80,6 +80,10 @@ ok
 
 === Options
 
+* `{return, value | flamegraph | filename}` - What return to the calling code. Defaults to `value` for `apply/2` function and `filename` for `capture/3` function.
+  * `value` - return the return value of the function being profiled. Implies the flamegraph data will be written to a file.
+  * `flamegraph` - return the flamegraph data in the format specified instead of writing it to a file.
+  * `filename` - return the filename of the file the flamegraph data was written to. Implies the flamegraph data will be written to a file.
 * `{output_directory, Dir}` - Specify the output directory to write trace data to. Default is project root, and relative paths are relative to the project root.
 * `{output_format, Format}` - Specify the output format. Unfortunately I was not able to implement all the formatters I would have liked to during Spawnfest. Currently only the `brendan_gregg` formatter is available. +++<del>+++`Format` must be one of `perf`, `brendan_gregg`, or `svg`. Default is `brendan_gregg`.+++</del>+++
 * `{open, Program}` - Specify the program to load the trace output in after output is generated. `Program` must be one of `speedscope` or `hotspot`. By default this option is not set. `hotspot` is not useful right now as it cannot load files generated by the `brendan_gregg` formatter.

src/eflambe.erl

@@ -14,9 +14,18 @@
 -type mfa_fun() :: {atom(), atom(), list()} | {fun(), list()}.
 
 -type program() :: hotspot | speedscope.
--type option() :: {output_directory, binary()} | {output_format, brendan_gregg} | {open, program()}.
+-type option() :: {return, value | flamegraph | filename}
+                | {output_directory, binary()}
+                | {output_format, brendan_gregg}
+                | {open, program()}.
 -type options() :: [option()].
 
+-type capture_return() :: {ok, [any()]} | {error, already_mecked}.
+
+-define(DEFAULT_OPTIONS, [{output_format, brendan_gregg}]).
+-define(DEFAULT_APPLY_OPTIONS, [{return, value}|?DEFAULT_OPTIONS]).
+-define(DEFAULT_CAPTURE_OPTIONS, [{return, filename}|?DEFAULT_OPTIONS]).
+
 %%--------------------------------------------------------------------
 %% @doc
 %% Starts capturing of function call data for any invocation of the specified
@@ -26,29 +35,26 @@
 %% @end
 %%--------------------------------------------------------------------
 
--spec capture(MFA :: mfa()) -> ok.
+-spec capture(MFA :: mfa()) -> capture_return().
 
 capture(MFA) ->
     capture(MFA, 1).
 
--spec capture(MFA :: mfa(), NumCalls :: pos_integer()) -> ok.
+-spec capture(MFA :: mfa(), NumCalls :: pos_integer()) -> capture_return().
 
 capture(MFA, NumCalls) ->
     capture(MFA, NumCalls, []).
 
 -spec capture(MFA :: mfa(), NumCalls :: pos_integer(), Options :: options()) ->
-    ok | {error, already_mecked}.
+    capture_return().
 
 capture({Module, Function, Arity}, NumCalls, Options)
   when is_atom(Module), is_atom(Function), is_integer(Arity) ->
-    setup_for_trace(),
+    CompleteOptions = merge([{max_calls, NumCalls}|Options], ?DEFAULT_CAPTURE_OPTIONS),
 
-    CompleteOptions = [{max_calls, NumCalls}|Options],
+    setup_for_trace(),
 
-    case eflambe_sup:start_trace({Module, Function, Arity}, CompleteOptions) of
-        {ok, _Pid} -> ok;
-        {error, _} = Error -> Error
-    end.
+    eflambe_sup:start_trace({Module, Function, Arity}, CompleteOptions).
 
 %%--------------------------------------------------------------------
 %% @doc
@@ -66,23 +72,27 @@ apply(Function) ->
 -spec apply(Function :: mfa_fun(), Options :: options()) -> any().
 
 apply({Module, Function, Args}, Options) when is_atom(Module), is_atom(Function), is_list(Args) ->
+    CompleteOptions = merge(Options, ?DEFAULT_APPLY_OPTIONS),
+
     setup_for_trace(),
-    {ok, TracePid} = eflambe_server:start_trace(Options),
+    {ok, TracePid} = eflambe_server:start_trace(CompleteOptions),
 
     % Invoke the original function
     Results = erlang:apply(Module, Function, Args),
 
-    ok = eflambe_server:stop_trace(TracePid),
+    {ok, _} = eflambe_server:stop_trace(TracePid),
     Results;
 
 apply({Function, Args}, Options) when is_function(Function), is_list(Args) ->
+    CompleteOptions = merge(Options, ?DEFAULT_APPLY_OPTIONS),
+
     setup_for_trace(),
-    {ok, TracePid} = eflambe_server:start_trace(Options),
+    {ok, TracePid} = eflambe_server:start_trace(CompleteOptions),
 
     % Invoke the original function
     Results = erlang:apply(Function, Args),
 
-    ok = eflambe_server:stop_trace(TracePid),
+    {ok, _} = eflambe_server:stop_trace(TracePid),
     Results.
 
 %%%===================================================================
@@ -91,3 +101,12 @@ apply({Function, Args}, Options) when is_function(Function), is_list(Args) ->
 
 setup_for_trace() ->
     application:ensure_all_started(eflambe).
+
+% https://stackoverflow.com/questions/21873644/combine-merge-two-erlang-lists
+merge(In1, In2) ->
+    Combined = In1 ++ In2,
+    Fun = fun(Key) ->
+                  [FinalValue|_] = proplists:get_all_values(Key, Combined),
+                  {Key, FinalValue}
+          end,
+    lists:map(Fun, proplists:get_keys(Combined)).

src/eflambe_server.erl

@@ -9,7 +9,7 @@
 -behaviour(gen_server).
 
 %% API
--export([start_link/2, start_capture_trace/1, stop_capture_trace/1, start_trace/1, stop_trace/1]).
+-export([start_link/2, start_capture_trace/1, stop_capture_trace/2, start_trace/1, stop_trace/1]).
 
 %% gen_server callbacks
 -export([init/1,
@@ -20,16 +20,18 @@
 
 -include_lib("kernel/include/logger.hrl").
 
--define(DEFAULT_OPTIONS, [{output_format, brendan_gregg}]).
 -define(FLAGS, [call, return_to, running, procs, garbage_collection, arity,
                 timestamp, set_on_spawn]).
 
 -record(state, {
+          callback :: {pid(), reference()},
           module :: atom(),
           max_calls :: integer(),
           calls = 0 :: integer(),
           options = [] :: list(),
-          pid_traces = [] :: list(pid_trace())
+          pid_traces = [] :: list(pid_trace()),
+          results = [] :: list(),
+          return :: atom()
          }).
 
 % For capture traces we could end up tracing invocations in multiple processes.
@@ -91,8 +93,11 @@ start_capture_trace(ServerPid) ->
 %%
 %% @end
 %%--------------------------------------------------------------------
-stop_capture_trace(ServerPid) ->
-    gen_server:call(ServerPid, stop_trace).
+
+-spec stop_capture_trace(pid(), any()) -> ok.
+
+stop_capture_trace(ServerPid, Return) ->
+    gen_server:call(ServerPid, {stop_trace, Return}).
 
 %%--------------------------------------------------------------------
 %% @doc
@@ -132,9 +137,6 @@ stop_trace(TracerPid) ->
 -spec init(Args :: list(tracer_options())) -> {ok, state()}.
 
 init([{Module, Function, Arity}, Options]) ->
-    % Generate complete list of options by falling back to default list
-    FinalOptions = merge(Options, ?DEFAULT_OPTIONS),
-
     % If necessary, set up meck to wrap original function in tracing code that
     % calls out to this process. We'll need to inject the pid of this tracer
     % (self()) into the wrapper function so we can track invocation and return
@@ -148,7 +150,7 @@ init([{Module, Function, Arity}, Options]) ->
 
         case StartedNew of
             true ->
-                eflambe_server:stop_capture_trace(ServerPid);
+                eflambe_server:stop_capture_trace(ServerPid, Results);
             false ->
                 ok
         end,
@@ -159,8 +161,17 @@ init([{Module, Function, Arity}, Options]) ->
         {error, already_mecked} ->
             {stop, already_mecked};
         ok ->
-            MaxCalls = proplists:get_value(max_calls, FinalOptions),
-            {ok, #state{module = Module, max_calls = MaxCalls, options = FinalOptions}}
+            MaxCalls = proplists:get_value(max_calls, Options),
+            Callback = proplists:get_value(callback, Options),
+            Return = proplists:get_value(return, Options),
+            InitialState = #state{
+                              callback = Callback,
+                              module = Module,
+                              max_calls = MaxCalls,
+                              options = Options,
+                              return = Return
+                             },
+            {ok, InitialState}
     end.
 
 -spec handle_call(Request :: any(), from(), state()) ->
@@ -191,46 +202,55 @@ handle_call(start_trace, {FromPid, _}, #state{max_calls = MaxCalls, calls = Call
             {reply, {ok, false}, State}
     end;
 
-handle_call(stop_trace, {FromPid, _}, #state{max_calls = MaxCalls, calls = Calls,
+handle_call({stop_trace, Return}, {FromPid, _}, #state{max_calls = MaxCalls, calls = Calls,
                                              module = ModuleName,
-                                             pid_traces = PidTraces} = State) ->
+                                             results = Results,
+                                             return = ReturnOption} = State) ->
     % Stop tracing immediately!
     stop_erlang_trace(FromPid),
 
-    NewTraces = case get_pid_trace(State, FromPid) of
+    case get_pid_trace(State, FromPid) of
         #pid_trace{tracer_pid = TracerPid} ->
-            eflambe_tracer:finish(TracerPid),
+            {ok, Result} = eflambe_tracer:finish(TracerPid),
+
+            FinalResult = case ReturnOption of
+                              value -> Return;
+                              _ -> Result
+                          end,
 
             % Generate new list of traces
-            remove_pid_trace(State, FromPid);
-        undefined ->
-            PidTraces
-    end,
+            NewTraces = remove_pid_trace(State, FromPid),
 
-    NoTracesRemaining = (MaxCalls =:= Calls),
+            NoTracesRemaining = (MaxCalls =:= Calls),
+
+            NewState = State#state{pid_traces = NewTraces, results = [FinalResult|Results]},
 
-    case {NoTracesRemaining, NewTraces} of
-        {true, []} ->
-            % We are completely finished
-            ok = eflambe_meck:unload(ModuleName),
+            case {NoTracesRemaining, NewTraces} of
+                {true, []} ->
+                    % We are completely finished
+                    ok = eflambe_meck:unload(ModuleName),
 
-            io:format("Successful finished trace~n"),
+                    io:format("Successful finished trace~n"),
 
-            % The only reason we don't stop here is because this is a call and
-            % the linked call would crash as well. This feels kind of wrong so
-            % I may revisit this
-            {reply, ok, State, {continue, finish}};
-        {_, _} ->
-            % Still some stuff in flight
-            {reply, ok, State#state{pid_traces = NewTraces}}
+                    % The only reason we don't stop here is because this is a call and
+                    % the linked call would crash as well. This feels kind of wrong so
+                    % I may revisit this
+                    {reply, ok, NewState, {continue, finish}};
+                {_, _} ->
+                    % Still some stuff in flight
+                    {reply, ok, NewState}
+            end;
+        undefined ->
+            {reply, ok, State}
     end.
 
 -spec handle_cast(any(), state()) -> {noreply, state()}.
 
 handle_cast(_Msg, State) ->
     {noreply, State}.
 
-handle_continue(finish, State) ->
+handle_continue(finish, #state{callback = PidRef, results = Results} = State) ->
+    ok = gen_server:reply(PidRef, lists:reverse(Results)),
     {stop, normal, State}.
 
 -spec handle_info(Info :: any(), state()) -> {noreply, state()} |
@@ -264,15 +284,6 @@ lookup_fun(Pid) ->
         (_) -> false
     end.
 
-% https://stackoverflow.com/questions/21873644/combine-merge-two-erlang-lists
-merge(In1, In2) ->
-    Combined = In1 ++ In2,
-    Fun = fun(Key) ->
-                  [FinalValue|_] = proplists:get_all_values(Key, Combined),
-                  {Key, FinalValue}
-          end,
-    lists:map(Fun, proplists:get_keys(Combined)).
-
 -spec start_erlang_trace(PidToTrace :: pid(), TracerPid :: pid()) -> integer().
 
 start_erlang_trace(PidToTrace, TracerPid) ->

src/eflambe_sup.erl

@@ -20,10 +20,26 @@
 start_link() ->
     supervisor:start_link({local, ?SERVER}, ?MODULE, []).
 
--spec start_trace(MFA :: mfa(), Options :: list()) -> {ok, pid()} | {error, already_mecked}.
+-spec start_trace(MFA :: mfa(), Options :: list()) -> {ok, [any()]} | {error, already_mecked}.
 
 start_trace(MFA, Options) ->
-    supervisor:start_child(?SERVER, [MFA, Options]).
+    % This feels hacky but we MAY need the gen_server to send us all the results
+    % right before it shuts down. Is there a gen_server or proc_lib function
+    % that provides a better abstraction for one-off messages FROM a gen_server
+    % like this?
+    %
+    % It's as the initial start_child start_link call should be treated as a
+    % regular gen_server call and we need to wait on a response.
+    Ref = make_ref(),
+    case supervisor:start_child(?SERVER, [MFA, [{callback, {self(), Ref}}|Options]]) of
+        {ok, _Child} ->
+            receive
+                {Ref, Msg} ->
+                    {ok, Msg}
+            end;
+        {error, _} = Error -> Error
+    end.
+
 
 %%%===================================================================
 %%% supervisor callbacks

src/eflambe_tracer.erl

@@ -31,7 +31,8 @@
           impl :: atom(),
           impl_state :: any(),
           options :: eflambe:options(),
-          filename :: file:filename_all()
+          filename :: file:filename_all(),
+          return :: atom()
          }).
 
 -type state() :: #state{}.
@@ -81,27 +82,45 @@ init([Options]) ->
     Filename = generate_filename(Ext),
     FullFilename = filename:join([output_directory(Options), Filename]),
 
+    Return = proplists:get_value(return, FinalOptions),
+
     % Initialize implementation state
     {ok, State} = erlang:apply(Impl, init, [FullFilename, Options]),
-    {ok, #state{impl = Impl, impl_state = State, options = FinalOptions, filename = FullFilename}}.
+    InitialState = #state{
+                      impl = Impl,
+                      impl_state = State,
+                      options = FinalOptions,
+                      filename = FullFilename,
+                      return = Return
+                     },
+    {ok, InitialState}.
 
 -spec handle_call(Request :: any(), from(), state()) ->
                                   {reply, Reply :: any(), state()} |
                                   {reply, Reply :: any(), state(), {continue, finish}}.
 
 handle_call(finish, _From, #state{impl = Impl, impl_state = ImplState, options = Options,
-                                 filename = Filename} = State) ->
+                                 filename = Filename, return = Return} = State) ->
     % Format the trace data and write to file
     {ok, _FinalImplState} = erlang:apply(Impl, finalize, [Options, ImplState]),
 
     % Open flamegraph viewer if specified
     maybe_open_in_program(Options, Filename),
     io:format("Output filename: ~s~n", [Filename]),
 
+    Results = case Return of
+                  value ->
+                      % Values are only readily available in the tracer so we
+                      % let the eflambe_server populate this type of return value
+                      undefined;
+                  filename -> Filename;
+                  flamegraph -> ok % TODO
+              end,
+
     % The only reason we don't stop here is because this is a call and the
     % linked call would crash as well. This feels kind of wrong so I may revisit
     % this
-    {reply, ok, State, {continue, finish}}.
+    {reply, {ok, Results}, State, {continue, finish}}.
 
 -spec handle_cast(any(), state()) -> {noreply, state()}.