1%%
2%% %CopyrightBegin%
3%%
4%% Copyright Ericsson AB 2010-2015. All Rights Reserved.
5%%
6%% Licensed under the Apache License, Version 2.0 (the "License");
7%% you may not use this file except in compliance with the License.
8%% You may obtain a copy of the License at
9%%
10%%     http://www.apache.org/licenses/LICENSE-2.0
11%%
12%% Unless required by applicable law or agreed to in writing, software
13%% distributed under the License is distributed on an "AS IS" BASIS,
14%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15%% See the License for the specific language governing permissions and
16%% limitations under the License.
17%%
18%% %CopyrightEnd%
19%%
20
21%%
22%% Statistics collector.
23%%
24
25-module(diameter_stats).
26-behaviour(gen_server).
27
28-export([reg/2, reg/1,
29         incr/3, incr/1,
30         read/1,
31         sum/1,
32         flush/1]).
33
34%% supervisor callback
35-export([start_link/0]).
36
37%% gen_server callbacks
38-export([init/1,
39         terminate/2,
40         handle_call/3,
41         handle_cast/2,
42         handle_info/2,
43         code_change/3]).
44
45%% debug
46-export([state/0,
47         uptime/0]).
48
49-include("diameter_internal.hrl").
50
51%% ets table containing 2-tuple stats. reg(Pid, Ref) inserts a {Pid,
52%% Ref}, incr(Counter, X, N) updates the counter keyed at {Counter,
53%% X}, and Pid death causes counters keyed on {Counter, Pid} to be
54%% deleted and added to those keyed on {Counter, Ref}.
55-define(TABLE, ?MODULE).
56
57%% Name of registered server.
58-define(SERVER, ?MODULE).
59
60%% Server state.
61-record(state, {id = diameter_lib:now()}).
62
63-type counter() :: any().
64-type ref() :: any().
65
66%% ---------------------------------------------------------------------------
67%% # reg(Pid, Ref)
68%%
69%% Register a process as a contributor of statistics associated with a
70%% specified term. Statistics can be contributed by specifying either
71%% Pid or Ref as the second argument to incr/3. Statistics contributed
72%% by Pid are folded into the corresponding entry for Ref when the
73%% process dies.
74%% ---------------------------------------------------------------------------
75
76-spec reg(pid(), ref())
77   -> boolean().
78
79reg(Pid, Ref)
80  when is_pid(Pid) ->
81    try
82        call({reg, Pid, Ref})
83    catch
84        exit: _ -> false
85    end.
86
87-spec reg(ref())
88   -> boolean().
89
90reg(Ref) ->
91    reg(self(), Ref).
92
93%% ---------------------------------------------------------------------------
94%% # incr(Counter, Ref, N)
95%%
96%% Increment a counter for the specified contributor.
97%%
98%% Ref will typically be an argument passed to reg/2 but there's
99%% nothing that requires this. Only registered pids can contribute
100%% counters however, otherwise incr/3 is a no-op.
101%% ---------------------------------------------------------------------------
102
103-spec incr(counter(), ref(), integer())
104   -> integer() | false.
105
106incr(Ctr, Ref, N)
107  when is_integer(N) ->
108    update_counter({Ctr, Ref}, N).
109
110incr(Ctr) ->
111    incr(Ctr, self(), 1).
112
113%% ---------------------------------------------------------------------------
114%% # read(Refs)
115%%
116%% Retrieve counters for the specified contributors.
117%% ---------------------------------------------------------------------------
118
119%% Read in the server process to ensure that counters for a dying
120%% contributor aren't folded concurrently with select.
121
122-spec read([ref()])
123   -> [{ref(), [{counter(), integer()}]}].
124
125read(Refs)
126  when is_list(Refs) ->
127    try call({read, Refs, false}) of
128        L -> to_refdict(L)
129    catch
130        exit: _ -> []
131    end.
132
133read(Refs, B) ->
134    MatchSpec = [{{{'_', '$1'}, '_'},
135                  [?ORCOND([{'=:=', '$1', {const, R}}
136                            || R <- Refs])],
137                  ['$_']}],
138    L = ets:select(?TABLE, MatchSpec),
139    B andalso delete(L),
140    L.
141
142to_refdict(L) ->
143    lists:foldl(fun append/2, orddict:new(), L).
144
145%% Order both references and counters in the returned list.
146append({{Ctr, Ref}, N}, Dict) ->
147    orddict:update(Ref,
148                   fun(D) -> orddict:store(Ctr, N, D) end,
149                   [{Ctr, N}],
150                   Dict).
151
152%% ---------------------------------------------------------------------------
153%% # sum(Refs)
154%%
155%% Retrieve counters summed over all contributors for each term.
156%% ---------------------------------------------------------------------------
157
158-spec sum([ref()])
159   -> [{ref(), [{counter(), integer()}]}].
160
161sum(Refs)
162  when is_list(Refs) ->
163    try call({read, Refs}) of
164        L -> [{R, to_ctrdict(Cs)} || {R, [_|_] = Cs} <- L]
165    catch
166        exit: _ -> []
167    end.
168
169read_refs(Refs) ->
170    [{R, readr(R)} || R <- Refs].
171
172readr(Ref) ->
173    MatchSpec = [{{{'_', '$1'}, '_'},
174                  [?ORCOND([{'=:=', '$1', {const, R}}
175                            || R <- [Ref | pids(Ref)]])],
176                  ['$_']}],
177    ets:select(?TABLE, MatchSpec).
178
179pids(Ref) ->
180    MatchSpec = [{{'$1', '$2'},
181                  [{'=:=', '$2', {const, Ref}}],
182                  ['$1']}],
183    ets:select(?TABLE, MatchSpec).
184
185to_ctrdict(L) ->
186    lists:foldl(fun({{C,_}, N}, D) -> orddict:update_counter(C, N, D) end,
187                orddict:new(),
188                L).
189
190%% ---------------------------------------------------------------------------
191%% # flush(Refs)
192%%
193%% Retrieve and delete statistics for the specified contributors.
194%% ---------------------------------------------------------------------------
195
196-spec flush([ref()])
197   -> [{ref(), {counter(), integer()}}].
198
199flush(Refs) ->
200    try call({read, Refs, true}) of
201        L -> to_refdict(L)
202    catch
203        exit: _ -> []
204    end.
205
206%% ===========================================================================
207
208start_link() ->
209    ServerName = {local, ?SERVER},
210    Module     = ?MODULE,
211    Args       = [],
212    Options    = [{spawn_opt, diameter_lib:spawn_opts(server, [])}],
213    gen_server:start_link(ServerName, Module, Args, Options).
214
215state() ->
216    call(state).
217
218uptime() ->
219    call(uptime).
220
221%% ----------------------------------------------------------
222%% # init/1
223%% ----------------------------------------------------------
224
225init([]) ->
226    ets:new(?TABLE, [named_table, set, public, {write_concurrency, true}]),
227    {ok, #state{}}.
228
229%% ----------------------------------------------------------
230%% # handle_call/3
231%% ----------------------------------------------------------
232
233handle_call(state, _, State) ->
234    {reply, State, State};
235
236handle_call(uptime, _, #state{id = Time} = State) ->
237    {reply, diameter_lib:now_diff(Time), State};
238
239handle_call({incr, T}, _, State) ->
240    {reply, update_counter(T), State};
241
242handle_call({reg, Pid, Ref}, _From, State) ->
243    B = ets:insert_new(?TABLE, {Pid, Ref}),
244    B andalso erlang:monitor(process, Pid),
245    {reply, B, State};
246
247handle_call({read, Refs, Del}, _From, State) ->
248    {reply, read(Refs, Del), State};
249
250handle_call({read, Refs}, _, State) ->
251    {reply, read_refs(Refs), State};
252
253handle_call(Req, From, State) ->
254    ?UNEXPECTED([Req, From]),
255    {reply, nok, State}.
256
257%% ----------------------------------------------------------
258%% # handle_cast/2
259%% ----------------------------------------------------------
260
261handle_cast(Msg, State) ->
262    ?UNEXPECTED([Msg]),
263    {noreply, State}.
264
265%% ----------------------------------------------------------
266%% # handle_info/2
267%% ----------------------------------------------------------
268
269handle_info({'DOWN', _MRef, process, Pid, _}, State) ->
270    down(Pid),
271    {noreply, State};
272
273handle_info(Info, State) ->
274    ?UNEXPECTED([Info]),
275    {noreply, State}.
276
277%% ----------------------------------------------------------
278%% # terminate/2
279%% ----------------------------------------------------------
280
281terminate(_Reason, _State) ->
282    ok.
283
284%% ----------------------------------------------------------
285%% # code_change/3
286%% ----------------------------------------------------------
287
288code_change(_OldVsn, State, _Extra) ->
289    {ok, State}.
290
291%% ===========================================================================
292
293%% down/1
294
295down(Pid) ->
296    down(lookup(Pid), ets:match_object(?TABLE, {{'_', Pid}, '_'})).
297
298down([{_, Ref} = T], L) ->
299    fold(Ref, L),
300    delete([T|L]);
301down([], L) -> %% flushed
302    delete(L).
303
304%% Fold pid-based entries into ref-based ones.
305fold(Ref, L) ->
306    lists:foreach(fun({{K, _}, V}) -> update_counter({{K, Ref}, V}) end, L).
307
308%% update_counter/2
309%%
310%% From an arbitrary process. Call to the server process to insert a
311%% new element if the counter doesn't exists so that two processes
312%% don't insert simultaneously.
313
314update_counter(Key, N) ->
315    try
316        ets:update_counter(?TABLE, Key, N)
317    catch
318        error: badarg ->
319            call({incr, {Key, N}})
320    end.
321
322%% update_counter/1
323%%
324%% From the server process, when update_counter/2 failed due to a
325%% non-existent entry.
326
327update_counter({{_Ctr, Ref} = Key, N} = T) ->
328    try
329        ets:update_counter(?TABLE, Key, N)
330    catch
331        error: badarg ->
332            (not is_pid(Ref) orelse ets:member(?TABLE, Ref))
333                andalso begin insert(T), N end
334    end.
335
336insert(T) ->
337    ets:insert(?TABLE, T).
338
339lookup(Key) ->
340    ets:lookup(?TABLE, Key).
341
342delete(Objs) ->
343    lists:foreach(fun({K,_}) -> ets:delete(?TABLE, K) end, Objs).
344
345%% call/1
346
347call(Request) ->
348    gen_server:call(?SERVER, Request, infinity).
349