Add logic to return filename or the profiled function's return value

src/eflambe.erl

@@ -20,6 +20,8 @@
                 | {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]).
@@ -33,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, binary() | [any()]}  | {error, already_mecked}.
+    capture_return().
 
 capture({Module, Function, Arity}, NumCalls, Options)
   when is_atom(Module), is_atom(Function), is_integer(Arity) ->
     CompleteOptions = merge([{max_calls, NumCalls}|Options], ?DEFAULT_CAPTURE_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
@@ -81,7 +80,7 @@ apply({Module, Function, Args}, Options) when is_atom(Module), is_atom(Function)
     % 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) ->
@@ -93,7 +92,7 @@ apply({Function, Args}, Options) when is_function(Function), is_list(Args) ->
     % Invoke the original function
     Results = erlang:apply(Function, Args),
 
-    ok = eflambe_server:stop_trace(TracePid),
+    {ok, _} = eflambe_server:stop_trace(TracePid),
     Results.
 
 %%%===================================================================

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,
@@ -30,7 +30,8 @@
           calls = 0 :: integer(),
           options = [] :: list(),
           pid_traces = [] :: list(pid_trace()),
-          results = [] :: list()
+          results = [] :: list(),
+          return :: atom()
          }).
 
 % For capture traces we could end up tracing invocations in multiple processes.
@@ -92,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
@@ -146,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,7 +163,15 @@ init([{Module, Function, Arity}, Options]) ->
         ok ->
             MaxCalls = proplists:get_value(max_calls, Options),
             Callback = proplists:get_value(callback, Options),
-            {ok, #state{callback = Callback, module = Module, max_calls = MaxCalls, options = 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()) ->
@@ -190,38 +202,46 @@ 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),
 
-    case {NoTracesRemaining, NewTraces} of
-        {true, []} ->
-            % We are completely finished
-            ok = eflambe_meck:unload(ModuleName),
+            NewState = State#state{pid_traces = NewTraces, results = [FinalResult|Results]},
 
-            io:format("Successful finished trace~n"),
+            case {NoTracesRemaining, NewTraces} of
+                {true, []} ->
+                    % We are completely finished
+                    ok = eflambe_meck:unload(ModuleName),
 
-            % 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}}
+                    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, NewState, {continue, finish}};
+                {_, _} ->
+                    % Still some stuff in flight
+                    {reply, ok, NewState}
+            end;
+        undefined ->
+            {reply, ok, State}
     end.
 
 -spec handle_cast(any(), state()) -> {noreply, state()}.
@@ -230,7 +250,7 @@ handle_cast(_Msg, State) ->
     {noreply, State}.
 
 handle_continue(finish, #state{callback = PidRef, results = Results} = State) ->
-    ok = gen_server:reply(PidRef, Results),
+    ok = gen_server:reply(PidRef, lists:reverse(Results)),
     {stop, normal, State}.
 
 -spec handle_info(Info :: any(), state()) -> {noreply, state()} |

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()}.
 

test/eflambe_server_SUITE.erl

@@ -80,49 +80,55 @@ start_link(_Config) ->
     {ok, Pid} = eflambe_server:start_link({arithmetic, multiply, 2}, [{max_calls, 1}]),
     true = is_pid(Pid),
 
-    eflambe_server:stop_capture_trace(Pid).
+    eflambe_server:stop_capture_trace(Pid, return_value).
 
 start_capture_trace(_Config) ->
-    Options = [{max_calls, 1}, {output_format, plain}],
     Self = self(),
+    Ref = make_ref(),
+    Options = [{callback, {Self, Ref}}, {max_calls, 1}, {output_format, plain}, {return, filename}],
 
     % Returns an ok tuple when eflambe_server is running and arguments are valid
     {ok, Pid} = eflambe_server:start_link({arithmetic, multiply, 2}, Options),
-    {state,arithmetic,1,0, FinalOptions, []} = get_server_state(Pid),
+    {state, {Self, Ref}, arithmetic, 1, 0, FinalOptions, [], [], filename} = get_server_state(Pid),
 
     {ok, true} = eflambe_server:start_capture_trace(Pid),
 
     % Stores trace state
-    {state, arithmetic, 1, 1, FinalOptions, [{pid_trace, Self, _Tracer}]} = State =
-      get_server_state(Pid),
+    State = get_server_state(Pid),
+    {state, {Self, Ref},  arithmetic, 1, 1, FinalOptions, Traces, [], filename} = State,
+    [{pid_trace, Self, _Tracer}] = Traces,
 
     % Returns the same trace data when called twice
     {ok, false} = eflambe_server:start_capture_trace(Pid),
     State = get_server_state(Pid),
 
     % Returns false when trace is stopped but number of calls has already been reached
-    ok = eflambe_server:stop_capture_trace(Pid).
+    ok = eflambe_server:stop_capture_trace(Pid, return_value).
 
 stop_capture_trace(_Config) ->
-    Options = [{max_calls, 2}, {output_format, plain}],
     Self = self(),
+    Ref = make_ref(),
+    Reply = {Self, Ref},
+    Options = [{callback, Reply}, {max_calls, 2}, {output_format, plain}, {return, filename}],
 
     {ok, Pid} = eflambe_server:start_link({arithmetic, multiply, 2}, Options),
 
     % Returns an ok tuple when eflambe_server is running and arguments are valid
     {ok, true} = eflambe_server:start_capture_trace(Pid),
 
     % Stores trace state
-    {state, arithmetic, 2, 1, FinalOptions, [{pid_trace, Self, _}]} = get_server_state(Pid),
+    {state, Reply, arithmetic, 2, 1, FinalOptions, Traces, [], filename} = get_server_state(Pid),
+    [{pid_trace, Self, _}] = Traces,
 
     % When called removes trace from server state
-    ok = eflambe_server:stop_capture_trace(Pid),
-    {state, arithmetic, 2, 1, FinalOptions, []} = get_server_state(Pid),
+    ok = eflambe_server:stop_capture_trace(Pid, return_value),
+    State = get_server_state(Pid),
+    {state, Reply, arithmetic, 2, 1, FinalOptions, [], [_Filename], filename} = State,
 
     % After being called start_capture_trace invoked in the same process returns
     % true to kick off a new trace
     {ok, true} = eflambe_server:start_capture_trace(Pid),
-    ok = eflambe_server:stop_capture_trace(Pid).
+    ok = eflambe_server:stop_capture_trace(Pid, return_value).
 
 get_server_state(Pid) ->
     {status, _, _, State} = sys:get_status(Pid),