|
| 1 | +%% Copyright Gojek |
| 2 | + |
| 3 | +-module(vmq_enhanced_auth_metrics). |
| 4 | + |
| 5 | +-behaviour(gen_server). |
| 6 | + |
| 7 | +%% API |
| 8 | +-export([ |
| 9 | + start_link/0, |
| 10 | + metrics/0, |
| 11 | + incr/1, |
| 12 | + incr/2 |
| 13 | +]). |
| 14 | + |
| 15 | +%% gen_server callbacks |
| 16 | +-export([ |
| 17 | + init/1, |
| 18 | + handle_call/3, |
| 19 | + handle_cast/2, |
| 20 | + handle_info/2, |
| 21 | + terminate/2, |
| 22 | + code_change/3 |
| 23 | +]). |
| 24 | + |
| 25 | +-define(SERVER, ?MODULE). |
| 26 | + |
| 27 | +-include_lib("vmq_enhanced_auth/src/vmq_enhanced_auth.hrl"). |
| 28 | + |
| 29 | +-record(state, {}). |
| 30 | +-record(metric_def, { |
| 31 | + type :: atom(), |
| 32 | + labels :: [metric_label()], |
| 33 | + id :: metric_id(), |
| 34 | + name :: atom(), |
| 35 | + description :: undefined | binary() |
| 36 | +}). |
| 37 | + |
| 38 | +-type metric_def() :: #metric_def{}. |
| 39 | +-type metric_val() :: {Id :: metric_id(), Val :: any()}. |
| 40 | +-type metric_label() :: {atom(), string()}. |
| 41 | +-type metric_id() :: |
| 42 | + atom() |
| 43 | + | {atom(), non_neg_integer() | atom()} |
| 44 | + | {atom(), atom(), atom()} |
| 45 | + | [{atom(), [{atom(), any()}]}]. |
| 46 | + |
| 47 | +%%%=================================================================== |
| 48 | +%%% API |
| 49 | +%%%=================================================================== |
| 50 | + |
| 51 | +-spec start_link() -> 'ignore' | {'error', _} | {'ok', pid()}. |
| 52 | +start_link() -> |
| 53 | + gen_server:start_link({local, ?SERVER}, ?MODULE, [], []). |
| 54 | + |
| 55 | +-spec metrics() -> [{metric_def(), non_neg_integer()}]. |
| 56 | +metrics() -> |
| 57 | + MetricDefs = metrics_defs(), |
| 58 | + MetricValues = metric_values(), |
| 59 | + |
| 60 | + %% Create id->metric def map |
| 61 | + IdDef = lists:foldl( |
| 62 | + fun(#metric_def{id = Id} = MD, Acc) -> |
| 63 | + maps:put(Id, MD, Acc) |
| 64 | + end, |
| 65 | + #{}, |
| 66 | + MetricDefs |
| 67 | + ), |
| 68 | + |
| 69 | + %% Merge metrics definitions with values and filter based on labels. |
| 70 | + Metrics = lists:filtermap( |
| 71 | + fun({Id, Val}) -> |
| 72 | + case maps:find(Id, IdDef) of |
| 73 | + {ok, #metric_def{ |
| 74 | + type = Type, id = Id, name = Name, labels = GotLabels, description = Description |
| 75 | + }} -> |
| 76 | + {true, {Type, GotLabels, Id, Name, Description, Val}}; |
| 77 | + error -> |
| 78 | + %% this could happen if metric definitions does |
| 79 | + %% not correspond to the ids returned with the |
| 80 | + %% metrics values. |
| 81 | + lager:warning("unknown metrics id: ~p", [Id]), |
| 82 | + false |
| 83 | + end |
| 84 | + end, |
| 85 | + MetricValues |
| 86 | + ), |
| 87 | + Metrics. |
| 88 | + |
| 89 | +-spec incr(any()) -> 'ok'. |
| 90 | +incr(Entry) -> |
| 91 | + incr_item(Entry, 1). |
| 92 | + |
| 93 | +-spec incr(any(), non_neg_integer()) -> 'ok'. |
| 94 | +incr(Entry, N) -> |
| 95 | + incr_item(Entry, N). |
| 96 | + |
| 97 | +%%%=================================================================== |
| 98 | +%%% gen_server callbacks |
| 99 | +%%%=================================================================== |
| 100 | + |
| 101 | +-spec init(_) -> {ok, #state{}}. |
| 102 | +init([]) -> |
| 103 | + AllEntries = [Id || #metric_def{id = Id} <- metrics_defs()], |
| 104 | + NumEntries = length(AllEntries), |
| 105 | + %% Sanity check where it is checked that there is a one-to-one |
| 106 | + %% mapping between atomics indexes and metrics identifiers by 1) |
| 107 | + %% checking that all metric identifiers have an atomics index and |
| 108 | + %% 2) that there are as many indexes as there are metrics and 3) |
| 109 | + %% that there are no index duplicates. |
| 110 | + Idxs = lists:map(fun(Id) -> met2idx(Id) end, AllEntries), |
| 111 | + NumEntries = length(lists:sort(Idxs)), |
| 112 | + NumEntries = length(lists:usort(Idxs)), |
| 113 | + |
| 114 | + %% only alloc a new atomics array if one doesn't already exist! |
| 115 | + case catch persistent_term:get(?MODULE) of |
| 116 | + {'EXIT', {badarg, _}} -> |
| 117 | + %% allocate twice the number of entries to make it possible to add |
| 118 | + %% new metrics during a hot code upgrade. |
| 119 | + ARef = atomics:new(2 * NumEntries, [{signed, false}]), |
| 120 | + persistent_term:put(?MODULE, ARef); |
| 121 | + _ExistingRef -> |
| 122 | + ok |
| 123 | + end, |
| 124 | + {ok, #state{}}. |
| 125 | + |
| 126 | +handle_call(_Request, _From, State) -> |
| 127 | + Reply = ok, |
| 128 | + {reply, Reply, State}. |
| 129 | + |
| 130 | +handle_cast(_Msg, State) -> |
| 131 | + {noreply, State}. |
| 132 | + |
| 133 | +handle_info(_Info, State) -> |
| 134 | + {noreply, State}. |
| 135 | + |
| 136 | +terminate(_Reason, _State) -> |
| 137 | + ok. |
| 138 | + |
| 139 | +code_change(_OldVsn, State, _Extra) -> |
| 140 | + {ok, State}. |
| 141 | + |
| 142 | +%%%=================================================================== |
| 143 | +%%% Internal functions |
| 144 | +%%%=================================================================== |
| 145 | + |
| 146 | +%% don't do the update |
| 147 | +incr_item(_, 0) -> |
| 148 | + ok; |
| 149 | +incr_item(Entry, Val) when Val > 0 -> |
| 150 | + ARef = |
| 151 | + case get(vmq_enhanced_auth_atomics_ref) of |
| 152 | + undefined -> |
| 153 | + Ref = persistent_term:get(?MODULE), |
| 154 | + put(vmq_enhanced_auth_atomics_ref, Ref), |
| 155 | + Ref; |
| 156 | + Ref -> |
| 157 | + Ref |
| 158 | + end, |
| 159 | + atomics:add(ARef, met2idx(Entry), Val). |
| 160 | + |
| 161 | +-spec metric_values() -> [metric_val()]. |
| 162 | +metric_values() -> |
| 163 | + lists:map( |
| 164 | + fun(#metric_def{id = Id}) -> |
| 165 | + try counter_val(Id) of |
| 166 | + Value -> {Id, Value} |
| 167 | + catch |
| 168 | + _:_ -> {Id, 0} |
| 169 | + end |
| 170 | + end, |
| 171 | + metrics_defs() |
| 172 | + ). |
| 173 | + |
| 174 | +-spec metrics_defs() -> [metric_def()]. |
| 175 | +metrics_defs() -> |
| 176 | + [ |
| 177 | + m( |
| 178 | + counter, |
| 179 | + [ |
| 180 | + {reason, atom_to_list(?INVALID_SIGNATURE)} |
| 181 | + ], |
| 182 | + {?REGISTER_AUTH_ERROR, ?INVALID_SIGNATURE}, |
| 183 | + ?REGISTER_AUTH_ERROR, |
| 184 | + <<"The number of times the auth_on_register hook returned error due to invalid_signature.">> |
| 185 | + ), |
| 186 | + m( |
| 187 | + counter, |
| 188 | + [ |
| 189 | + {reason, atom_to_list(?MISSING_RID)} |
| 190 | + ], |
| 191 | + {?REGISTER_AUTH_ERROR, ?MISSING_RID}, |
| 192 | + ?REGISTER_AUTH_ERROR, |
| 193 | + <<"The number of times the auth_on_register hook returned error due to missing_rid.">> |
| 194 | + ), |
| 195 | + m( |
| 196 | + counter, |
| 197 | + [ |
| 198 | + {reason, atom_to_list(?USERNAME_RID_MISMATCH)} |
| 199 | + ], |
| 200 | + {?REGISTER_AUTH_ERROR, ?USERNAME_RID_MISMATCH}, |
| 201 | + ?REGISTER_AUTH_ERROR, |
| 202 | + <<"The number of times the auth_on_register hook returned error due to username_rid_mismatch.">> |
| 203 | + ) |
| 204 | + ]. |
| 205 | + |
| 206 | +-spec m(atom(), [metric_label()], metric_id(), atom(), 'undefined' | binary()) -> metric_def(). |
| 207 | +m(Type, Labels, UniqueId, Name, Description) -> |
| 208 | + #metric_def{ |
| 209 | + type = Type, |
| 210 | + labels = Labels, |
| 211 | + id = UniqueId, |
| 212 | + name = Name, |
| 213 | + description = Description |
| 214 | + }. |
| 215 | + |
| 216 | +counter_val(Entry) -> |
| 217 | + ARef = persistent_term:get(?MODULE), |
| 218 | + atomics:get(ARef, met2idx(Entry)). |
| 219 | + |
| 220 | +met2idx({?REGISTER_AUTH_ERROR, ?INVALID_SIGNATURE}) -> 1; |
| 221 | +met2idx({?REGISTER_AUTH_ERROR, ?USERNAME_RID_MISMATCH}) -> 2; |
| 222 | +met2idx({?REGISTER_AUTH_ERROR, ?MISSING_RID}) -> 3. |
0 commit comments