1% Licensed under the Apache License, Version 2.0 (the "License"); you may not
2% use this file except in compliance with the License. You may obtain a copy of
3% the License at
4%
5%   http://www.apache.org/licenses/LICENSE-2.0
6%
7% Unless required by applicable law or agreed to in writing, software
8% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10% License for the specific language governing permissions and limitations under
11% the License.
12
13-module(couch_ref_counter).
14-behaviour(gen_server).
15
16-export([start/1, init/1, terminate/2, handle_call/3, handle_cast/2, code_change/3, handle_info/2]).
17-export([drop/1,drop/2,add/1,add/2,count/1]).
18
19start(ChildProcs) ->
20    gen_server:start(couch_ref_counter, {self(), ChildProcs}, []).
21
22
23drop(RefCounterPid) ->
24    drop(RefCounterPid, self()).
25
26drop(RefCounterPid, Pid) ->
27    gen_server:call(RefCounterPid, {drop, Pid}, infinity).
28
29
30add(RefCounterPid) ->
31    add(RefCounterPid, self()).
32
33add(RefCounterPid, Pid) ->
34    gen_server:call(RefCounterPid, {add, Pid}, infinity).
35
36count(RefCounterPid) ->
37    gen_server:call(RefCounterPid, count).
38
39% server functions
40
41-record(srv,
42    {
43    referrers=dict:new(), % a dict of each ref counting proc.
44    child_procs=[]
45    }).
46
47init({Pid, ChildProcs}) ->
48    [link(ChildProc) || ChildProc <- ChildProcs],
49    Referrers = dict:from_list([{Pid, {erlang:monitor(process, Pid), 1}}]),
50    {ok, #srv{referrers=Referrers, child_procs=ChildProcs}}.
51
52
53terminate(_Reason, #srv{child_procs=ChildProcs}) ->
54    [couch_util:shutdown_sync(Pid) || Pid <- ChildProcs],
55    ok.
56
57
58handle_call({add, Pid},_From, #srv{referrers=Referrers}=Srv) ->
59    Referrers2 =
60    case dict:find(Pid, Referrers) of
61    error ->
62        dict:store(Pid, {erlang:monitor(process, Pid), 1}, Referrers);
63    {ok, {MonRef, RefCnt}} ->
64        dict:store(Pid, {MonRef, RefCnt + 1}, Referrers)
65    end,
66    {reply, ok, Srv#srv{referrers=Referrers2}};
67handle_call(count, _From, Srv) ->
68    {monitors, Monitors} =  process_info(self(), monitors),
69    {reply, length(Monitors), Srv};
70handle_call({drop, Pid}, _From, #srv{referrers=Referrers}=Srv) ->
71    Referrers2 =
72    case dict:find(Pid, Referrers) of
73    {ok, {MonRef, 1}} ->
74        erlang:demonitor(MonRef, [flush]),
75        dict:erase(Pid, Referrers);
76    {ok, {MonRef, Num}} ->
77        dict:store(Pid, {MonRef, Num-1}, Referrers);
78    error ->
79        Referrers
80    end,
81    Srv2 = Srv#srv{referrers=Referrers2},
82    case should_close() of
83    true ->
84        {stop,normal,ok,Srv2};
85    false ->
86        {reply, ok, Srv2}
87    end.
88
89handle_cast(Msg, _Srv)->
90    exit({unknown_msg,Msg}).
91
92
93code_change(_OldVsn, State, _Extra) ->
94    {ok, State}.
95
96handle_info({'DOWN', MonRef, _, Pid, _}, #srv{referrers=Referrers}=Srv) ->
97    {ok, {MonRef, _RefCount}} = dict:find(Pid, Referrers),
98    Srv2 = Srv#srv{referrers=dict:erase(Pid, Referrers)},
99    case should_close() of
100    true ->
101        {stop,normal,Srv2};
102    false ->
103        {noreply,Srv2}
104    end.
105
106
107should_close() ->
108    case process_info(self(), monitors) of
109    {monitors, []} ->   true;
110    _ ->                false
111    end.
112