lib/etorrent-1.0/src/etorrent_piece_mgr.erl

   1 %%%-------------------------------------------------------------------
   2 %%% File    : etorrent_piece_mgr.erl
   3 %%% Author  : Jesper Louis Andersen <jlouis@ogre.home>
   4 %%% Description : Piece Manager code
   5 %%%
   6 %%% Created : 21 Jul 2008 by Jesper Louis Andersen <jlouis@ogre.home>
   7 %%%-------------------------------------------------------------------
   8 -module(etorrent_piece_mgr).
   9
  10 -include_lib("stdlib/include/qlc.hrl").
  11 -include("etorrent_piece.hrl").
  12
  13 -behaviour(gen_server).
  14
  15 %% API
  16 -export([start_link/0, delete/1, decrease_missing_chunks/2, statechange/3,
  17          fetched/2, bitfield/1, select/1, select/2, valid/2, interesting/2,
  18          num_not_fetched/1, check_interest/2, add_pieces/2, chunk/3]).
  19
  20 -export([fetched/1]).
  21
  22 %% gen_server callbacks
  23 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
  24          terminate/2, code_change/3]).
  25
  26 -record(state, {}).
  27 -define(SERVER, ?MODULE).
  28
  29 %%====================================================================
  30 %% API
  31 %%====================================================================
  32 %%--------------------------------------------------------------------
  33 %% Function: start_link() -> {ok,Pid} | ignore | {error,Error}
  34 %% Description: Starts the server
  35 %%--------------------------------------------------------------------
  36 start_link() ->
  37     gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).
  38
  39 %%--------------------------------------------------------------------
  40 %% Function: delete(Id) -> ok
  41 %% Description: Rip out the pieces identified by torrent with Id
  42 %%--------------------------------------------------------------------
  43 delete(Id) ->
  44     gen_server:cast(?SERVER, {delete, Id}).
  45 %%--------------------------------------------------------------------
  46 %% Function: add_pieces(Id, FPList) -> void()
  47 %% Args:  Id ::= integer() - torrent id
  48 %%        FPList ::= [{Hash, Files, Done}]
  49 %% Description: Add a list of pieces to the database.
  50 %%--------------------------------------------------------------------
  51 add_pieces(Id, Pieces) ->
  52     gen_server:call(?SERVER, {add_pieces, Id, Pieces}).
  53
  54 chunk(Id, Idx, N) ->
  55     gen_server:call(?SERVER, {chunk, Id, Idx, N}).
  56
  57 %%--------------------------------------------------------------------
  58 %% Function: t_decrease_missing_chunks/2
  59 %% Args:       Id ::= integer() - torrent id
  60 %%             PieceNum ::= integer() - piece index
  61 %% Description: Decrease missing chunks for the {Id, PieceNum} pair.
  62 %%--------------------------------------------------------------------
  63 decrease_missing_chunks(Id, Idx) ->
  64     gen_server:call(?SERVER, {decrease_missing, Id, Idx}).
  65
  66 %%--------------------------------------------------------------------
  67 %% Function: statechange(Id, PieceNumber, S) -> ok
  68 %% Description: Update the {Id, PieceNumber} pair to have state S
  69 %%--------------------------------------------------------------------
  70 statechange(Id, PN, State) ->
  71     gen_server:call(?SERVER, {statechange, Id, PN, State}).
  72
  73 fetched(Id, Idx) ->
  74     [R] = ets:lookup(etorrent_piece_tbl, {Id, Idx}),
  75     R#piece.state =:= fetched.
  76
  77 fetched(Id) ->
  78     MH = #piece { state = fetched, idpn = {Id, '$1'}, _ = '_'},
  79     ets:select(etorrent_piece_tbl,
  80                [{MH, [], ['$1']}]).
  81
  82 bitfield(Id) when is_integer(Id) ->
  83     NP = etorrent_torrent:num_pieces(Id),
  84     Fetched = fetched(Id),
  85     etorrent_peer_communication:construct_bitfield(
  86       NP,
  87       gb_sets:from_list(Fetched)).
  88
  89 %%--------------------------------------------------------------------
  90 %% Function: num_not_fetched(Id) -> integer()
  91 %% Description: Return the number of not_fetched pieces for torrent Id.
  92 %%--------------------------------------------------------------------
  93 num_not_fetched(Id) when is_integer(Id) ->
  94     Q = qlc:q([R#piece.piece_number ||
  95                   R <- ets:table(etorrent_piece_tbl),
  96                   R#piece.id =:= Id,
  97                   R#piece.state =:= not_fetched]),
  98     length(qlc:e(Q)).
  99
 100 %%--------------------------------------------------------------------
 101 %% Function: check_interest(Id, PieceSet) -> interested | not_interested | invalid_piece
 102 %% Description: Given a set of pieces, return if we are interested in any of them.
 103 %%--------------------------------------------------------------------
 104 check_interest(Id, PieceSet) when is_integer(Id) ->
 105     It = gb_sets:iterator(PieceSet),
 106     find_interest_piece(Id, gb_sets:next(It)).
 107
 108 %%--------------------------------------------------------------------
 109 %% Function: select(Id) -> [#piece]
 110 %% Description: Return all pieces for a torrent id
 111 %%--------------------------------------------------------------------
 112 select(Id) ->
 113     MatchHead = #piece { idpn = {Id, '_'}, _ = '_'},
 114     ets:select(etorrent_piece_tbl, [{MatchHead, [], ['$_']}]).
 115
 116 %%--------------------------------------------------------------------
 117 %% Function: select(Id, PieceNumber) -> [#piece]
 118 %% Description: Return the piece PieceNumber for the Id torrent
 119 %%--------------------------------------------------------------------
 120 select(Id, PN) ->
 121     ets:lookup(etorrent_piece_tbl, {Id, PN}).
 122
 123 %%--------------------------------------------------------------------
 124 %% Function: valid(Id, PieceNumber) -> bool()
 125 %% Description: Is the piece valid for this torrent?
 126 %%--------------------------------------------------------------------
 127 valid(Id, Pn) when is_integer(Id) ->
 128     case ets:lookup(etorrent_piece_tbl, {Id, Pn}) of
 129         [] -> false;
 130         [_] -> true
 131     end.
 132
 133 %%--------------------------------------------------------------------
 134 %% Function: interesting(Id, Pn) -> bool()
 135 %% Description: Is the piece interesting?
 136 %%--------------------------------------------------------------------
 137 interesting(Id, Pn) when is_integer(Id) ->
 138     [P] = ets:lookup(etorrent_piece_tbl, {Id, Pn}),
 139     case P#piece.state of
 140         fetched ->
 141             false;
 142         _ ->
 143             true
 144     end.
 145
 146 %%====================================================================
 147 %% gen_server callbacks
 148 %%====================================================================
 149
 150 %%--------------------------------------------------------------------
 151 %% Function: init(Args) -> {ok, State} |
 152 %%                         {ok, State, Timeout} |
 153 %%                         ignore               |
 154 %%                         {stop, Reason}
 155 %% Description: Initiates the server
 156 %%--------------------------------------------------------------------
 157 init([]) ->
 158     process_flag(trap_exit, true),
 159     _Tid = ets:new(etorrent_piece_tbl, [set, protected, named_table,
 160                                         {keypos, #piece.idpn}]),
 161     {ok, #state{}}.
 162
 163 %%--------------------------------------------------------------------
 164 %% Function: %% handle_call(Request, From, State) -> {reply, Reply, State} |
 165 %%                                      {reply, Reply, State, Timeout} |
 166 %%                                      {noreply, State} |
 167 %%                                      {noreply, State, Timeout} |
 168 %%                                      {stop, Reason, Reply, State} |
 169 %%                                      {stop, Reason, State}
 170 %% Description: Handling call messages
 171 %%--------------------------------------------------------------------
 172 handle_call({add_pieces, Id, Pieces}, _From, S) ->
 173     lists:foreach(
 174       fun({PN, Hash, Files, State}) ->
 175               ets:insert(etorrent_piece_tbl,
 176                          #piece { idpn = {Id, PN},
 177                                   id = Id,
 178                                   piece_number = PN,
 179                                   hash = Hash,
 180                                   files = Files,
 181                                   state = State})
 182       end,
 183       Pieces),
 184     {reply, ok, S};
 185 handle_call({chunk, Id, Idx, N}, _From, S) ->
 186     [P] = ets:lookup(etorrent_piece_tbl, {Id, Idx}),
 187     ets:insert(etorrent_piece_tbl, P#piece { state = chunked, left = N}),
 188     {reply, ok, S};
 189 handle_call({decrease_missing, Id, Idx}, _From, S) ->
 190     case ets:update_counter(etorrent_piece_tbl, {Id, Idx},
 191                             {#piece.left, -1}) of
 192         0 -> {reply, full, S};
 193         N when is_integer(N) -> {reply, ok, S}
 194     end;
 195 handle_call({statechange, Id, Idx, State}, _From, S) ->
 196     [P] = ets:lookup(etorrent_piece_tbl, {Id, Idx}),
 197     ets:insert(etorrent_piece_tbl, P#piece { state = State }),
 198     {reply, ok, S};
 199 handle_call({delete, Id}, _From, S) ->
 200     MatchHead = #piece { idpn = {Id, '_'}, _ = '_'},
 201     ets:select_delete(etorrent_piece_tbl, [{MatchHead, [], [true]}]),
 202     {reply, ok, S};
 203 handle_call(_Request, _From, State) ->
 204     Reply = ok,
 205     {reply, Reply, State}.
 206
 207 %%--------------------------------------------------------------------
 208 %% Function: handle_cast(Msg, State) -> {noreply, State} |
 209 %%                                      {noreply, State, Timeout} |
 210 %%                                      {stop, Reason, State}
 211 %% Description: Handling cast messages
 212 %%--------------------------------------------------------------------
 213 handle_cast(_Msg, State) ->
 214     {noreply, State}.
 215
 216 %%--------------------------------------------------------------------
 217 %% Function: handle_info(Info, State) -> {noreply, State} |
 218 %%                                       {noreply, State, Timeout} |
 219 %%                                       {stop, Reason, State}
 220 %% Description: Handling all non call/cast messages
 221 %%--------------------------------------------------------------------
 222 handle_info(_Info, State) ->
 223     {noreply, State}.
 224
 225 %%--------------------------------------------------------------------
 226 %% Function: terminate(Reason, State) -> void()
 227 %% Description: This function is called by a gen_server when it is about to
 228 %% terminate. It should be the opposite of Module:init/1 and do any necessary
 229 %% cleaning up. When it returns, the gen_server terminates with Reason.
 230 %% The return value is ignored.
 231 %%--------------------------------------------------------------------
 232 terminate(_Reason, _State) ->
 233     ets:delete(etorrent_piece_tbl),
 234     ok.
 235
 236 %%--------------------------------------------------------------------
 237 %% Func: code_change(OldVsn, State, Extra) -> {ok, NewState}
 238 %% Description: Convert process state when code is changed
 239 %%--------------------------------------------------------------------
 240 code_change(_OldVsn, State, _Extra) ->
 241     {ok, State}.
 242
 243 %%--------------------------------------------------------------------
 244 %%% Internal functions
 245 %%--------------------------------------------------------------------
 246 find_interest_piece(_Id, none) ->
 247     not_interested;
 248 find_interest_piece(Id, {Pn, Next}) ->
 249     case ets:lookup(etorrent_piece_tbl, {Id, Pn}) of
 250         [] ->
 251             invalid_piece;
 252         [P] when is_record(P, piece) ->
 253             case P#piece.state of
 254                 fetched ->
 255                     find_interest_piece(Id, gb_sets:next(Next));
 256                 _Other ->
 257                     interested
 258             end
 259     end.