lib/etorrent-1.0/src/etorrent_rate_mgr.erl

   1 %%%-------------------------------------------------------------------
   2 %%% File    : etorrent_rate_mgr.erl
   3 %%% Author  : Jesper Louis Andersen <jlouis@ogre.home>
   4 %%% Description : Rate management process
   5 %%%
   6 %%% Created : 17 Jul 2008 by Jesper Louis Andersen <jlouis@ogre.home>
   7 %%%-------------------------------------------------------------------
   8 -module(etorrent_rate_mgr).
   9
  10 -include("peer_state.hrl").
  11 -include("rate_mgr.hrl").
  12 -include("etorrent_rate.hrl").
  13
  14 -behaviour(gen_server).
  15
  16 -define(DEFAULT_SNUB_TIME, 30).
  17
  18 %% API
  19 -export([start_link/0,
  20
  21          choke/2, unchoke/2, interested/2, not_interested/2,
  22          local_choke/2, local_unchoke/2,
  23
  24          recv_rate/5, recv_rate/4, send_rate/4,
  25
  26          snubbed/2,
  27
  28          fetch_recv_rate/2,
  29          fetch_send_rate/2,
  30          select_state/2,
  31
  32          global_rate/0]).
  33
  34 %% gen_server callbacks
  35 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
  36          terminate/2, code_change/3]).
  37
  38 -record(state, { recv,
  39                  send,
  40                  state,
  41
  42                  global_recv,
  43                  global_send}).
  44
  45 -define(SERVER, ?MODULE).
  46
  47 %%====================================================================
  48 %% API
  49 %%====================================================================
  50 %%--------------------------------------------------------------------
  51 %% Function: start_link() -> {ok,Pid} | ignore | {error,Error}
  52 %% Description: Starts the server
  53 %%--------------------------------------------------------------------
  54 start_link() ->
  55     gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).
  56
  57 %% Send state information
  58 choke(Id, Pid) -> gen_server:cast(?SERVER, {choke, Id, Pid}).
  59 unchoke(Id, Pid) -> gen_server:cast(?SERVER, {unchoke, Id, Pid}).
  60 interested(Id, Pid) -> gen_server:cast(?SERVER, {interested, Id, Pid}).
  61 not_interested(Id, Pid) -> gen_server:cast(?SERVER, {not_interested, Id, Pid}).
  62 local_choke(Id, Pid) -> gen_server:cast(?SERVER, {local_choke, Id, Pid}).
  63 local_unchoke(Id, Pid) -> gen_server:cast(?SERVER, {local_unchoke, Id, Pid}).
  64
  65 select_state(Id, Who) ->
  66     case ets:lookup(etorrent_peer_state, {Id, Who}) of
  67         [] ->
  68             #peer_state { }; % Pick defaults
  69         [P] ->
  70             P
  71     end.
  72
  73 snubbed(Id, Who) ->
  74     T = etorrent_rate:now_secs(),
  75     case ets:lookup(etorrent_recv_state, {Id, Who}) of
  76         [] ->
  77             false;
  78         [#rate_mgr { last_got = unknown }] ->
  79             false;
  80         [#rate_mgr { last_got = U}] ->
  81             T - U > ?DEFAULT_SNUB_TIME
  82     end.
  83
  84
  85 fetch_recv_rate(Id, Pid) -> fetch_rate(etorrent_recv_state, Id, Pid).
  86 fetch_send_rate(Id, Pid) -> fetch_rate(etorrent_send_state, Id, Pid).
  87
  88 recv_rate(Id, Pid, Rate, Amount) ->
  89     recv_rate(Id, Pid, Rate, Amount, normal).
  90
  91 recv_rate(Id, Pid, Rate, Amount, Update) ->
  92     gen_server:cast(?SERVER, {recv_rate, Id, Pid, Rate, Amount, Update}).
  93
  94 send_rate(Id, Pid, Rate, Amount) ->
  95     gen_server:cast(?SERVER, {send_rate, Id, Pid, Rate, Amount, normal}).
  96
  97 global_rate() ->
  98     gen_server:call(?SERVER, global_rate).
  99
 100 %%====================================================================
 101 %% gen_server callbacks
 102 %%====================================================================
 103
 104 %%--------------------------------------------------------------------
 105 %% Function: init(Args) -> {ok, State} |
 106 %%                         {ok, State, Timeout} |
 107 %%                         ignore               |
 108 %%                         {stop, Reason}
 109 %% Description: Initiates the server
 110 %%--------------------------------------------------------------------
 111 init([]) ->
 112     process_flag(trap_exit, true),
 113     RTid = ets:new(etorrent_recv_state, [set, protected, named_table,
 114                                          {keypos, #rate_mgr.pid}]),
 115     STid = ets:new(etorrent_send_state, [set, protected, named_table,
 116                                          {keypos, #rate_mgr.pid}]),
 117     StTid = ets:new(etorrent_peer_state, [set, protected, named_table,
 118                                          {keypos, #peer_state.pid}]),
 119     {ok, #state{ recv = RTid, send = STid, state = StTid,
 120                  global_recv = etorrent_rate:init(?RATE_FUDGE),
 121                  global_send = etorrent_rate:init(?RATE_FUDGE)}}.
 122
 123 %%--------------------------------------------------------------------
 124 %% Function: %% handle_call(Request, From, State) -> {reply, Reply, State} |
 125 %%                                      {reply, Reply, State, Timeout} |
 126 %%                                      {noreply, State} |
 127 %%                                      {noreply, State, Timeout} |
 128 %%                                      {stop, Reason, Reply, State} |
 129 %%                                      {stop, Reason, State}
 130 %% Description: Handling call messages
 131 %%--------------------------------------------------------------------
 132 handle_call(global_rate, _From, S) ->
 133     #state { global_recv = GR, global_send = GS } = S,
 134     Reply = { GR#peer_rate.rate, GS#peer_rate.rate },
 135     {reply, Reply, S};
 136 handle_call(_Request, _From, State) ->
 137     Reply = ok,
 138     {reply, Reply, State}.
 139
 140 %%--------------------------------------------------------------------
 141 %% Function: handle_cast(Msg, State) -> {noreply, State} |
 142 %%                                      {noreply, State, Timeout} |
 143 %%                                      {stop, Reason, State}
 144 %% Description: Handling cast messages
 145 %%--------------------------------------------------------------------
 146 handle_cast({What, Id, Pid}, S) ->
 147     ok = alter_state(What, Id, Pid),
 148     {noreply, S};
 149 handle_cast({What, Id, Who, Rate, Amount, Update}, S) ->
 150     NS = alter_state(What, Id, Who, Rate, Amount, Update, S),
 151     {noreply, NS};
 152 handle_cast(_Msg, State) ->
 153     {noreply, State}.
 154
 155 %%--------------------------------------------------------------------
 156 %% Function: handle_info(Info, State) -> {noreply, State} |
 157 %%                                       {noreply, State, Timeout} |
 158 %%                                       {stop, Reason, State}
 159 %% Description: Handling all non call/cast messages
 160 %%--------------------------------------------------------------------
 161 handle_info({'DOWN', _Ref, process, Pid, _Reason}, S) ->
 162     true = ets:match_delete(etorrent_recv_state, #rate_mgr { pid = {'_', Pid}, _='_'}),
 163     true = ets:match_delete(etorrent_send_state, #rate_mgr { pid = {'_', Pid}, _='_'}),
 164     true = ets:match_delete(etorrent_peer_state, #peer_state { pid = {'_', Pid}, _='_'}),
 165     {noreply, S};
 166 handle_info(_Info, State) ->
 167     {noreply, State}.
 168
 169 %%--------------------------------------------------------------------
 170 %% Function: terminate(Reason, State) -> void()
 171 %% Description: This function is called by a gen_server when it is about to
 172 %% terminate. It should be the opposite of Module:init/1 and do any necessary
 173 %% cleaning up. When it returns, the gen_server terminates with Reason.
 174 %% The return value is ignored.
 175 %%--------------------------------------------------------------------
 176 terminate(_Reason, S) ->
 177     true = ets:delete(S#state.recv),
 178     true = ets:delete(S#state.send),
 179     true = ets:delete(S#state.state),
 180     ok.
 181
 182
 183 %%--------------------------------------------------------------------
 184 %% Func: code_change(OldVsn, State, Extra) -> {ok, NewState}
 185 %% Description: Convert process state when code is changed
 186 %%--------------------------------------------------------------------
 187 code_change(_OldVsn, State, _Extra) ->
 188     {ok, State}.
 189
 190 %%--------------------------------------------------------------------
 191 %%% Internal functions
 192 %%--------------------------------------------------------------------
 193
 194 alter_state(What, Id, Pid) ->
 195     _R = case ets:lookup(etorrent_peer_state, {Id, Pid}) of
 196         [] ->
 197             ets:insert(etorrent_peer_state,
 198               alter_record(What,
 199                            #peer_state { pid = {Id, Pid},
 200                                          choke_state = choked,
 201                                          interest_state = not_interested,
 202                                          local_choke = true})),
 203             erlang:monitor(process, Pid);
 204         [R] ->
 205             ets:insert(etorrent_peer_state,
 206                        alter_record(What, R))
 207     end,
 208     ok.
 209
 210 alter_record(What, R) ->
 211     case What of
 212         choke ->
 213             R#peer_state { choke_state = choked };
 214         unchoke ->
 215             R#peer_state { choke_state = unchoked };
 216         interested ->
 217             R#peer_state { interest_state = interested };
 218         not_interested ->
 219             R#peer_state { interest_state = not_interested };
 220         local_choke ->
 221             R#peer_state { local_choke = true };
 222         local_unchoke ->
 223             R#peer_state { local_choke = false}
 224     end.
 225
 226 alter_state(What, Id, Who, Rate, Amount, Update, S) ->
 227     {T, NS} = case What of
 228                   recv_rate -> NR = etorrent_rate:update(
 229                                       S#state.global_recv,
 230                                       Amount),
 231                                { etorrent_recv_state,
 232                                  S#state { global_recv = NR }};
 233                   send_rate -> NR = etorrent_rate:update(
 234                                       S#state.global_send,
 235                                       Amount),
 236                                { etorrent_send_state,
 237                                  S#state { global_send = NR }}
 238               end,
 239     _R = case ets:lookup(T, {Id, Who}) of
 240         [] ->
 241             ets:insert(T,
 242               #rate_mgr { pid = {Id, Who},
 243                           last_got = case Update of
 244                                          normal -> unknown;
 245                                          last_update -> etorrent_rate:now_secs()
 246                                      end,
 247                           rate = Rate }),
 248             erlang:monitor(process, Who);
 249         [R] ->
 250             ets:insert(T, R#rate_mgr { rate = Rate })
 251     end,
 252     NS.
 253
 254 fetch_rate(Where, Id, Pid) ->
 255     case ets:lookup(Where, {Id, Pid}) of
 256         [] ->
 257             none;
 258         [R] -> R#rate_mgr.rate
 259     end.
 260