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 { interest_state = not_interested,
  69                           choke_state = choked,
  70                           local_choke = true };
  71         [P] ->
  72             P
  73     end.
  74
  75 snubbed(Id, Who) ->
  76     T = etorrent_rate:now_secs(),
  77     case ets:lookup(etorrent_recv_state, {Id, Who}) of
  78         [] ->
  79             false;
  80         [#rate_mgr { last_got = unknown }] ->
  81             false;
  82         [#rate_mgr { last_got = U}] ->
  83             T - U > ?DEFAULT_SNUB_TIME
  84     end.
  85
  86
  87 fetch_recv_rate(Id, Pid) -> fetch_rate(etorrent_recv_state, Id, Pid).
  88 fetch_send_rate(Id, Pid) -> fetch_rate(etorrent_send_state, Id, Pid).
  89
  90 recv_rate(Id, Pid, Rate, Amount) ->
  91     recv_rate(Id, Pid, Rate, Amount, normal).
  92
  93 recv_rate(Id, Pid, Rate, Amount, Update) ->
  94     gen_server:cast(?SERVER, {recv_rate, Id, Pid, Rate, Amount, Update}).
  95
  96 send_rate(Id, Pid, Rate, Amount) ->
  97     gen_server:cast(?SERVER, {send_rate, Id, Pid, Rate, Amount, normal}).
  98
  99 global_rate() ->
 100     gen_server:call(?SERVER, global_rate).
 101
 102 %%====================================================================
 103 %% gen_server callbacks
 104 %%====================================================================
 105
 106 %%--------------------------------------------------------------------
 107 %% Function: init(Args) -> {ok, State} |
 108 %%                         {ok, State, Timeout} |
 109 %%                         ignore               |
 110 %%                         {stop, Reason}
 111 %% Description: Initiates the server
 112 %%--------------------------------------------------------------------
 113 init([]) ->
 114     process_flag(trap_exit, true),
 115     RTid = ets:new(etorrent_recv_state, [set, protected, named_table,
 116                                          {keypos, #rate_mgr.pid}]),
 117     STid = ets:new(etorrent_send_state, [set, protected, named_table,
 118                                          {keypos, #rate_mgr.pid}]),
 119     StTid = ets:new(etorrent_peer_state, [set, protected, named_table,
 120                                          {keypos, #peer_state.pid}]),
 121     {ok, #state{ recv = RTid, send = STid, state = StTid,
 122                  global_recv = etorrent_rate:init(?RATE_FUDGE),
 123                  global_send = etorrent_rate:init(?RATE_FUDGE)}}.
 124
 125 %%--------------------------------------------------------------------
 126 %% Function: %% handle_call(Request, From, State) -> {reply, Reply, State} |
 127 %%                                      {reply, Reply, State, Timeout} |
 128 %%                                      {noreply, State} |
 129 %%                                      {noreply, State, Timeout} |
 130 %%                                      {stop, Reason, Reply, State} |
 131 %%                                      {stop, Reason, State}
 132 %% Description: Handling call messages
 133 %%--------------------------------------------------------------------
 134 handle_call(global_rate, _From, S) ->
 135     #state { global_recv = GR, global_send = GS } = S,
 136     Reply = { GR#peer_rate.rate, GS#peer_rate.rate },
 137     {reply, Reply, S};
 138 handle_call(_Request, _From, State) ->
 139     Reply = ok,
 140     {reply, Reply, State}.
 141
 142 %%--------------------------------------------------------------------
 143 %% Function: handle_cast(Msg, State) -> {noreply, State} |
 144 %%                                      {noreply, State, Timeout} |
 145 %%                                      {stop, Reason, State}
 146 %% Description: Handling cast messages
 147 %%--------------------------------------------------------------------
 148 handle_cast({What, Id, Pid}, S) ->
 149     ok = alter_state(What, Id, Pid),
 150     {noreply, S};
 151 handle_cast({What, Id, Who, Rate, Amount, Update}, S) ->
 152     NS = alter_state(What, Id, Who, Rate, Amount, Update, S),
 153     {noreply, NS};
 154 handle_cast(_Msg, State) ->
 155     {noreply, State}.
 156
 157 %%--------------------------------------------------------------------
 158 %% Function: handle_info(Info, State) -> {noreply, State} |
 159 %%                                       {noreply, State, Timeout} |
 160 %%                                       {stop, Reason, State}
 161 %% Description: Handling all non call/cast messages
 162 %%--------------------------------------------------------------------
 163 handle_info({'DOWN', _Ref, process, Pid, _Reason}, S) ->
 164     true = ets:match_delete(etorrent_recv_state, #rate_mgr { pid = {'_', Pid}, _='_'}),
 165     true = ets:match_delete(etorrent_send_state, #rate_mgr { pid = {'_', Pid}, _='_'}),
 166     true = ets:match_delete(etorrent_peer_state, #peer_state { pid = {'_', Pid}, _='_'}),
 167     {noreply, S};
 168 handle_info(_Info, State) ->
 169     {noreply, State}.
 170
 171 %%--------------------------------------------------------------------
 172 %% Function: terminate(Reason, State) -> void()
 173 %% Description: This function is called by a gen_server when it is about to
 174 %% terminate. It should be the opposite of Module:init/1 and do any necessary
 175 %% cleaning up. When it returns, the gen_server terminates with Reason.
 176 %% The return value is ignored.
 177 %%--------------------------------------------------------------------
 178 terminate(_Reason, S) ->
 179     true = ets:delete(S#state.recv),
 180     true = ets:delete(S#state.send),
 181     true = ets:delete(S#state.state),
 182     ok.
 183
 184
 185 %%--------------------------------------------------------------------
 186 %% Func: code_change(OldVsn, State, Extra) -> {ok, NewState}
 187 %% Description: Convert process state when code is changed
 188 %%--------------------------------------------------------------------
 189 code_change(_OldVsn, State, _Extra) ->
 190     {ok, State}.
 191
 192 %%--------------------------------------------------------------------
 193 %%% Internal functions
 194 %%--------------------------------------------------------------------
 195
 196 alter_state(What, Id, Pid) ->
 197     _R = case ets:lookup(etorrent_peer_state, {Id, Pid}) of
 198         [] ->
 199             ets:insert(etorrent_peer_state,
 200               alter_record(What,
 201                            #peer_state { pid = {Id, Pid},
 202                                          choke_state = choked,
 203                                          interest_state = not_interested,
 204                                          local_choke = true})),
 205             erlang:monitor(process, Pid);
 206         [R] ->
 207             ets:insert(etorrent_peer_state,
 208                        alter_record(What, R))
 209     end,
 210     ok.
 211
 212 alter_record(What, R) ->
 213     case What of
 214         choke ->
 215             R#peer_state { choke_state = choked };
 216         unchoke ->
 217             R#peer_state { choke_state = unchoked };
 218         interested ->
 219             R#peer_state { interest_state = interested };
 220         not_interested ->
 221             R#peer_state { interest_state = not_interested };
 222         local_choke ->
 223             R#peer_state { local_choke = true };
 224         local_unchoke ->
 225             R#peer_state { local_choke = false}
 226     end.
 227
 228 alter_state(What, Id, Who, Rate, Amount, Update, S) ->
 229     {T, NS} = case What of
 230                   recv_rate -> NR = etorrent_rate:update(
 231                                       S#state.global_recv,
 232                                       Amount),
 233                                { etorrent_recv_state,
 234                                  S#state { global_recv = NR }};
 235                   send_rate -> NR = etorrent_rate:update(
 236                                       S#state.global_send,
 237                                       Amount),
 238                                { etorrent_send_state,
 239                                  S#state { global_send = NR }}
 240               end,
 241     _R = case ets:lookup(T, {Id, Who}) of
 242         [] ->
 243             ets:insert(T,
 244               #rate_mgr { pid = {Id, Who},
 245                           last_got = case Update of
 246                                          normal -> unknown;
 247                                          last_update -> etorrent_rate:now_secs()
 248                                      end,
 249                           rate = Rate }),
 250             erlang:monitor(process, Who);
 251         [R] ->
 252             ets:insert(T, R#rate_mgr { rate = Rate })
 253     end,
 254     NS.
 255
 256 fetch_rate(Where, Id, Pid) ->
 257     case ets:lookup(Where, {Id, Pid}) of
 258         [] ->
 259             none;
 260         [R] -> R#rate_mgr.rate
 261     end.
 262