From 98143ae182bf7134bf3c1a5b1027bd55e24d60e5 Mon Sep 17 00:00:00 2001
From: Andrew Thompson <andrew@hijacked.us>
Date: Fri, 15 Mar 2013 00:48:42 -0400
Subject: [PATCH] Allow storage of persistant per-process lager metadata

For persistant processes with some immutable metadata (riak vnode and
the vnode ID, for example), implement lager:md/0 and lager:md/1 for
getting/setting such metadata into the process dictionary.

Such metadata is automatically included in any lager message metadata,
so you can just set it in your init() function or whatever and not have
to worry about passing the data around and using it in every lager call.
---
 src/lager.erl               | 28 ++++++++++++++++++++++++++++
 src/lager_transform.erl     |  6 ++++--
 test/lager_test_backend.erl | 27 +++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/lager.erl b/src/lager.erl
index 37feb3a..6d5f13a 100644
--- a/src/lager.erl
+++ b/src/lager.erl
@@ -20,9 +20,12 @@
 
 -include("lager.hrl").
 
+-define(LAGER_MD_KEY, '__lager_metadata').
+
 %% API
 -export([start/0,
         log/3, log/4,
+        md/0, md/1,
         trace/2, trace/3, trace_file/2, trace_file/3, trace_console/1, trace_console/2,
         clear_all_traces/0, stop_trace/1, status/0,
         get_loglevel/1, set_loglevel/2, set_loglevel/3, get_loglevels/0,
@@ -51,6 +54,31 @@ start_ok(App, {error, {not_started, Dep}}) ->
 start_ok(App, {error, Reason}) ->
     erlang:error({app_start_failed, App, Reason}).
 
+%% @doc Get lager metadata for current process
+-spec md() -> [{atom(), any()}].
+md() ->
+    case erlang:get(?LAGER_MD_KEY) of
+        undefined -> [];
+        MD -> MD
+    end.
+
+%% @doc Set lager metadata for current process.
+%% Will badarg if you don't supply a list of {key, value} tuples keyed by atoms.
+-spec md([{atom(), any()},...]) -> ok.
+md(NewMD) when is_list(NewMD) ->
+    %% make sure its actually a real proplist
+    case lists:all(
+            fun({Key, _Value}) when is_atom(Key) -> true;
+                (_) -> false
+            end, NewMD) of
+        true ->
+            erlang:put(?LAGER_MD_KEY, NewMD),
+            ok;
+        false ->
+            erlang:error(badarg)
+    end;
+md(_) ->
+    erlang:error(badarg).
 
 -spec dispatch_log(log_level(), list(), string(), list() | none, pos_integer()) ->  ok | {error, lager_not_running}.
 %% this is the same check that the parse transform bakes into the module at compile time
diff --git a/src/lager_transform.erl b/src/lager_transform.erl
index 30364ee..f610684 100644
--- a/src/lager_transform.erl
+++ b/src/lager_transform.erl
@@ -89,7 +89,9 @@ transform_statement({call, Line, {remote, _Line1, {atom, _Line2, lager},
                         {cons, Line, {tuple, Line, [
                                     {atom, Line, node},
                                     {call, Line, {atom, Line, node}, []}]},
-                         {nil, Line}}}}}},
+                        %% get the metadata with lager:md(), this will always return a list so we can use it as the tail here
+                        {call, Line, {remote, Line, {atom, Line, lager}, {atom, Line, md}}, []}}}}}},
+                            %{nil, Line}}}}}}},
             DefaultAttrs = case erlang:get(application) of
                 undefined ->
                     DefaultAttrs0;
@@ -151,7 +153,7 @@ transform_statement({call, Line, {remote, _Line1, {atom, _Line2, lager},
                         [],
                         %% trick the linter into avoiding a 'term constructed by not used' error:
                         %% (fun() -> {error, lager_not_running} end)();
-                        [{call,9, {'fun',9, {clauses, [{clause,9,[],[], [{tuple,9,[{atom,9,error},{atom,9,lager_not_running}]}]}]}}, []}]},
+                        [{call, Line, {'fun', Line, {clauses, [{clause, Line, [],[], [{tuple, Line, [{atom, Line, error},{atom, Line, lager_not_running}]}]}]}}, []}]},
                     %% If we care about the loglevel, or there's any traces installed, we have do more checking
                     %% {Level, Traces} when (Level band SeverityAsInt) /= 0 orelse Traces /= [] ->
                     {clause, Line,
diff --git a/test/lager_test_backend.erl b/test/lager_test_backend.erl
index 345bc0f..a85b162 100644
--- a/test/lager_test_backend.erl
+++ b/test/lager_test_backend.erl
@@ -424,6 +424,33 @@ lager_test_() ->
                         ?assertEqual({?DEBUG bor ?INFO bor ?NOTICE bor ?WARNING bor ?ERROR bor ?CRITICAL bor ?ALERT bor ?EMERGENCY, []}, lager_config:get(loglevel)),
                         ok
                 end
+            },
+            {"metadata in the process dictionary works",
+                fun() ->
+                        lager:md([{platypus, gravid}, {sloth, hirsute}, {duck, erroneous}]),
+                        lager:info("I sing the animal kingdom electric!"),
+                        {_Level, _Time, _Message, Metadata}  = pop(),
+                        ?assertEqual(gravid, proplists:get_value(platypus, Metadata)),
+                        ?assertEqual(hirsute, proplists:get_value(sloth, Metadata)),
+                        ?assertEqual(erroneous, proplists:get_value(duck, Metadata)),
+                        ?assertEqual(undefined, proplists:get_value(eagle, Metadata)),
+                        lager:md([{platypus, gravid}, {sloth, hirsute}, {eagle, superincumbent}]),
+                        lager:info("I sing the animal kingdom dielectric!"),
+                        {_Level2, _Time2, _Message2, Metadata2}  = pop(),
+                        ?assertEqual(gravid, proplists:get_value(platypus, Metadata2)),
+                        ?assertEqual(hirsute, proplists:get_value(sloth, Metadata2)),
+                        ?assertEqual(undefined, proplists:get_value(duck, Metadata2)),
+                        ?assertEqual(superincumbent, proplists:get_value(eagle, Metadata2)),
+                        ok
+                end
+            },
+            {"can't store invalid metadata",
+                fun() ->
+                        ?assertEqual(ok, lager:md([{platypus, gravid}, {sloth, hirsute}, {duck, erroneous}])),
+                        ?assertError(badarg, lager:md({flamboyant, flamingo})),
+                        ?assertError(badarg, lager:md("zookeeper zephyr")),
+                        ok
+                end
             }
         ]
     }.