905 lines
30 KiB
Rust
905 lines
30 KiB
Rust
use std::{
|
|
collections::{HashMap, HashSet},
|
|
fs::File,
|
|
net::SocketAddr,
|
|
sync::{
|
|
atomic::{AtomicU64, Ordering},
|
|
Arc,
|
|
},
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use anyhow::Context;
|
|
use futures::{stream::FuturesUnordered, StreamExt};
|
|
use log::{debug, info, trace, warn};
|
|
use parking_lot::{Mutex, RwLock};
|
|
use tokio::{sync::mpsc::UnboundedSender, time::timeout};
|
|
|
|
use crate::{
|
|
buffers::{ByteBuf, ByteString},
|
|
chunk_tracker::{ChunkMarkingResult, ChunkTracker},
|
|
clone_to_owned::CloneToOwned,
|
|
file_ops::FileOps,
|
|
lengths::{Lengths, ValidPieceIndex},
|
|
peer_binary_protocol::{Handshake, Message, MessageOwned, Piece, Request},
|
|
peer_connection::{PeerConnection, PeerConnectionHandler, WriterRequest},
|
|
peer_state::{InflightRequest, LivePeerState, PeerState},
|
|
spawn_utils::{spawn, BlockingSpawner},
|
|
torrent_metainfo::TorrentMetaV1Owned,
|
|
type_aliases::{PeerHandle, Sha1, BF},
|
|
};
|
|
|
|
pub struct InflightPiece {
|
|
pub peer: PeerHandle,
|
|
pub started: Instant,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct PeerStates {
|
|
states: HashMap<PeerHandle, PeerState>,
|
|
seen: HashSet<SocketAddr>,
|
|
inflight_pieces: HashMap<ValidPieceIndex, InflightPiece>,
|
|
tx: HashMap<PeerHandle, Arc<tokio::sync::mpsc::UnboundedSender<WriterRequest>>>,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct AggregatePeerStats {
|
|
pub connecting: usize,
|
|
pub live: usize,
|
|
}
|
|
|
|
impl PeerStates {
|
|
pub fn stats(&self) -> AggregatePeerStats {
|
|
self.states
|
|
.values()
|
|
.fold(AggregatePeerStats::default(), |mut s, p| {
|
|
match p {
|
|
PeerState::Connecting(_) => s.connecting += 1,
|
|
PeerState::Live(_) => s.live += 1,
|
|
};
|
|
s
|
|
})
|
|
}
|
|
pub fn add_if_not_seen(
|
|
&mut self,
|
|
addr: SocketAddr,
|
|
tx: UnboundedSender<WriterRequest>,
|
|
) -> Option<PeerHandle> {
|
|
if self.seen.contains(&addr) {
|
|
return None;
|
|
}
|
|
let handle = self.add(addr, tx)?;
|
|
self.seen.insert(addr);
|
|
Some(handle)
|
|
}
|
|
pub fn seen(&self) -> &HashSet<SocketAddr> {
|
|
&self.seen
|
|
}
|
|
pub fn get_live(&self, handle: PeerHandle) -> Option<&LivePeerState> {
|
|
if let PeerState::Live(ref l) = self.states.get(&handle)? {
|
|
return Some(l);
|
|
}
|
|
None
|
|
}
|
|
pub fn get_live_mut(&mut self, handle: PeerHandle) -> Option<&mut LivePeerState> {
|
|
if let PeerState::Live(ref mut l) = self.states.get_mut(&handle)? {
|
|
return Some(l);
|
|
}
|
|
None
|
|
}
|
|
pub fn try_get_live_mut(&mut self, handle: PeerHandle) -> anyhow::Result<&mut LivePeerState> {
|
|
self.get_live_mut(handle)
|
|
.ok_or_else(|| anyhow::anyhow!("peer dropped"))
|
|
}
|
|
pub fn add(
|
|
&mut self,
|
|
addr: SocketAddr,
|
|
tx: UnboundedSender<WriterRequest>,
|
|
) -> Option<PeerHandle> {
|
|
let handle = addr;
|
|
if self.states.contains_key(&addr) {
|
|
return None;
|
|
}
|
|
self.states.insert(handle, PeerState::Connecting(addr));
|
|
self.tx.insert(handle, Arc::new(tx));
|
|
Some(handle)
|
|
}
|
|
pub fn drop_peer(&mut self, handle: PeerHandle) -> Option<PeerState> {
|
|
let result = self.states.remove(&handle);
|
|
self.tx.remove(&handle);
|
|
result
|
|
}
|
|
pub fn mark_i_am_choked(&mut self, handle: PeerHandle, is_choked: bool) -> Option<bool> {
|
|
let live = self.get_live_mut(handle)?;
|
|
let prev = live.i_am_choked;
|
|
live.i_am_choked = is_choked;
|
|
Some(prev)
|
|
}
|
|
pub fn mark_peer_interested(
|
|
&mut self,
|
|
handle: PeerHandle,
|
|
is_interested: bool,
|
|
) -> Option<bool> {
|
|
let live = self.get_live_mut(handle)?;
|
|
let prev = live.peer_interested;
|
|
live.peer_interested = is_interested;
|
|
Some(prev)
|
|
}
|
|
pub fn update_bitfield_from_vec(
|
|
&mut self,
|
|
handle: PeerHandle,
|
|
bitfield: Vec<u8>,
|
|
) -> Option<Option<BF>> {
|
|
let live = self.get_live_mut(handle)?;
|
|
let bitfield = BF::from_vec(bitfield);
|
|
let prev = live.bitfield.take();
|
|
live.bitfield = Some(bitfield);
|
|
Some(prev)
|
|
}
|
|
pub fn clone_tx(&self, handle: PeerHandle) -> Option<Arc<UnboundedSender<WriterRequest>>> {
|
|
Some(self.tx.get(&handle)?.clone())
|
|
}
|
|
pub fn remove_inflight_piece(&mut self, piece: ValidPieceIndex) -> Option<InflightPiece> {
|
|
self.inflight_pieces.remove(&piece)
|
|
}
|
|
}
|
|
|
|
pub struct TorrentStateLocked {
|
|
pub peers: PeerStates,
|
|
pub chunks: ChunkTracker,
|
|
}
|
|
|
|
#[derive(Default, Debug)]
|
|
pub struct AtomicStats {
|
|
pub have: AtomicU64,
|
|
pub downloaded_and_checked: AtomicU64,
|
|
pub uploaded: AtomicU64,
|
|
pub fetched_bytes: AtomicU64,
|
|
|
|
pub downloaded_pieces: AtomicU64,
|
|
pub total_piece_download_ms: AtomicU64,
|
|
}
|
|
|
|
impl AtomicStats {
|
|
pub fn average_piece_download_time(&self) -> Option<Duration> {
|
|
let d = self.downloaded_pieces.load(Ordering::Relaxed);
|
|
let t = self.total_piece_download_ms.load(Ordering::Relaxed);
|
|
if d == 0 {
|
|
return None;
|
|
}
|
|
Some(Duration::from_secs_f64(t as f64 / d as f64 / 1000f64))
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct StatsSnapshot {
|
|
pub have_bytes: u64,
|
|
pub downloaded_and_checked_bytes: u64,
|
|
pub downloaded_and_checked_pieces: u64,
|
|
pub fetched_bytes: u64,
|
|
pub uploaded_bytes: u64,
|
|
pub initially_needed_bytes: u64,
|
|
pub remaining_bytes: u64,
|
|
pub live_peers: u32,
|
|
pub seen_peers: u32,
|
|
pub connecting_peers: u32,
|
|
pub time: Instant,
|
|
}
|
|
|
|
pub struct TorrentState {
|
|
pub torrent: TorrentMetaV1Owned,
|
|
pub locked: Arc<RwLock<TorrentStateLocked>>,
|
|
pub files: Vec<Arc<Mutex<File>>>,
|
|
pub info_hash: [u8; 20],
|
|
pub peer_id: [u8; 20],
|
|
pub lengths: Lengths,
|
|
pub needed: u64,
|
|
pub stats: AtomicStats,
|
|
|
|
pub spawner: BlockingSpawner,
|
|
}
|
|
|
|
impl TorrentState {
|
|
pub fn file_ops(&self) -> FileOps<'_, Sha1> {
|
|
FileOps::new(&self.torrent, &self.files, &self.lengths)
|
|
}
|
|
|
|
pub fn get_next_needed_piece(&self, peer_handle: PeerHandle) -> Option<ValidPieceIndex> {
|
|
let g = self.locked.read();
|
|
let bf = g.peers.get_live(peer_handle)?.bitfield.as_ref()?;
|
|
for n in g.chunks.get_needed_pieces().iter_ones() {
|
|
if bf.get(n).map(|v| *v) == Some(true) {
|
|
// in theory it should be safe without validation, but whatever.
|
|
return self.lengths.validate_piece_index(n as u32);
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
pub fn am_i_choked(&self, peer_handle: PeerHandle) -> Option<bool> {
|
|
self.locked
|
|
.read()
|
|
.peers
|
|
.get_live(peer_handle)
|
|
.map(|l| l.i_am_choked)
|
|
}
|
|
|
|
pub fn reserve_next_needed_piece(&self, peer_handle: PeerHandle) -> Option<ValidPieceIndex> {
|
|
if self.am_i_choked(peer_handle)? {
|
|
warn!("we are choked by {}, can't reserve next piece", peer_handle);
|
|
return None;
|
|
}
|
|
let mut g = self.locked.write();
|
|
let n = {
|
|
let mut n_opt = None;
|
|
let bf = g.peers.get_live(peer_handle)?.bitfield.as_ref()?;
|
|
for n in g.chunks.get_needed_pieces().iter_ones() {
|
|
if bf.get(n).map(|v| *v) == Some(true) {
|
|
n_opt = Some(n);
|
|
break;
|
|
}
|
|
}
|
|
|
|
self.lengths.validate_piece_index(n_opt? as u32)?
|
|
};
|
|
g.peers.inflight_pieces.insert(
|
|
n,
|
|
InflightPiece {
|
|
peer: peer_handle,
|
|
started: Instant::now(),
|
|
},
|
|
);
|
|
g.chunks.reserve_needed_piece(n);
|
|
Some(n)
|
|
}
|
|
|
|
pub fn am_i_interested_in_peer(&self, handle: PeerHandle) -> bool {
|
|
self.get_next_needed_piece(handle).is_some()
|
|
}
|
|
|
|
pub fn try_steal_old_slow_piece(&self, handle: PeerHandle) -> Option<ValidPieceIndex> {
|
|
let total = self.stats.downloaded_pieces.load(Ordering::Relaxed);
|
|
|
|
// heuristic for not enough precision in average time
|
|
if total < 20 {
|
|
return None;
|
|
}
|
|
let avg_time = self.stats.average_piece_download_time()?;
|
|
|
|
let mut g = self.locked.write();
|
|
let (idx, elapsed, piece_req) = g
|
|
.peers
|
|
.inflight_pieces
|
|
.iter_mut()
|
|
// don't steal from myself
|
|
.filter(|(_, r)| r.peer != handle)
|
|
.map(|(p, r)| (p, r.started.elapsed(), r))
|
|
.max_by_key(|(_, e, _)| *e)?;
|
|
|
|
// heuristic for "too slow peer"
|
|
if elapsed > avg_time * 10 {
|
|
debug!(
|
|
"{} will steal piece {} from {}: elapsed time {:?}, avg piece time: {:?}",
|
|
handle, idx, piece_req.peer, elapsed, avg_time
|
|
);
|
|
piece_req.peer = handle;
|
|
piece_req.started = Instant::now();
|
|
return Some(*idx);
|
|
}
|
|
None
|
|
}
|
|
|
|
pub fn try_steal_piece(&self, handle: PeerHandle) -> Option<ValidPieceIndex> {
|
|
let mut rng = rand::thread_rng();
|
|
use rand::seq::IteratorRandom;
|
|
let g = self.locked.read();
|
|
let pl = g.peers.get_live(handle)?;
|
|
g.peers
|
|
.inflight_pieces
|
|
.keys()
|
|
.filter(|p| !pl.inflight_requests.iter().any(|req| req.piece == **p))
|
|
.choose(&mut rng)
|
|
.copied()
|
|
}
|
|
|
|
pub fn set_peer_live(&self, handle: PeerHandle, h: Handshake) {
|
|
let mut g = self.locked.write();
|
|
match g.peers.states.get_mut(&handle) {
|
|
Some(s @ &mut PeerState::Connecting(_)) => {
|
|
*s = PeerState::Live(LivePeerState::new(h.peer_id));
|
|
}
|
|
_ => {
|
|
warn!("peer {} was in wrong state", handle);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn drop_peer(&self, handle: PeerHandle) -> bool {
|
|
let mut g = self.locked.write();
|
|
let peer = match g.peers.drop_peer(handle) {
|
|
Some(peer) => peer,
|
|
None => return false,
|
|
};
|
|
match peer {
|
|
PeerState::Connecting(_) => {}
|
|
PeerState::Live(l) => {
|
|
for req in l.inflight_requests {
|
|
g.chunks.mark_chunk_request_cancelled(req.piece, req.chunk);
|
|
}
|
|
}
|
|
}
|
|
true
|
|
}
|
|
|
|
pub fn get_uploaded(&self) -> u64 {
|
|
self.stats.uploaded.load(Ordering::Relaxed)
|
|
}
|
|
pub fn get_downloaded(&self) -> u64 {
|
|
self.stats.downloaded_and_checked.load(Ordering::Relaxed)
|
|
}
|
|
|
|
pub fn get_left_to_download(&self) -> u64 {
|
|
self.needed - self.get_downloaded()
|
|
}
|
|
|
|
pub fn maybe_transmit_haves(&self, index: ValidPieceIndex) {
|
|
let mut futures = Vec::new();
|
|
|
|
let g = self.locked.read();
|
|
for (handle, peer_state) in g.peers.states.iter() {
|
|
match peer_state {
|
|
PeerState::Live(live) => {
|
|
if !live.peer_interested {
|
|
continue;
|
|
}
|
|
|
|
if live
|
|
.bitfield
|
|
.as_ref()
|
|
.and_then(|b| b.get(index.get() as usize).map(|v| *v))
|
|
.unwrap_or(false)
|
|
{
|
|
continue;
|
|
}
|
|
|
|
let tx = match g.peers.tx.get(handle) {
|
|
Some(tx) => tx,
|
|
None => continue,
|
|
};
|
|
let tx = Arc::downgrade(tx);
|
|
futures.push(async move {
|
|
if let Some(tx) = tx.upgrade() {
|
|
if tx
|
|
.send(WriterRequest::Message(Message::Have(index.get())))
|
|
.is_err()
|
|
{
|
|
// whatever
|
|
}
|
|
}
|
|
});
|
|
}
|
|
_ => continue,
|
|
}
|
|
}
|
|
|
|
if futures.is_empty() {
|
|
trace!("no peers to transmit Have={} to, saving some work", index);
|
|
return;
|
|
}
|
|
|
|
let mut unordered: FuturesUnordered<_> = futures.into_iter().collect();
|
|
spawn(
|
|
format!("transmit_haves(piece={}, count={})", index, unordered.len()),
|
|
async move {
|
|
while unordered.next().await.is_some() {}
|
|
Ok(())
|
|
},
|
|
);
|
|
}
|
|
|
|
pub fn add_peer_if_not_seen(self: &Arc<Self>, addr: SocketAddr) -> bool {
|
|
let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<WriterRequest>();
|
|
let handle = match self.locked.write().peers.add_if_not_seen(addr, out_tx) {
|
|
Some(handle) => handle,
|
|
None => return false,
|
|
};
|
|
|
|
let handler = PeerHandler {
|
|
addr,
|
|
state: self.clone(),
|
|
spawner: self.spawner,
|
|
};
|
|
let peer_connection =
|
|
PeerConnection::new(addr, self.torrent.info_hash, self.peer_id, handler);
|
|
spawn(format!("manage_peer({})", handle), async move {
|
|
if let Err(e) = peer_connection.manage_peer(out_rx).await {
|
|
debug!("error managing peer {}: {:#}", handle, e)
|
|
};
|
|
peer_connection.into_handler().state.drop_peer(handle);
|
|
Ok::<_, anyhow::Error>(())
|
|
});
|
|
true
|
|
}
|
|
|
|
pub fn stats_snapshot(&self) -> StatsSnapshot {
|
|
let g = self.locked.read();
|
|
use Ordering::*;
|
|
let (live, connecting) =
|
|
g.peers
|
|
.states
|
|
.values()
|
|
.fold((0u32, 0u32), |(live, connecting), p| match p {
|
|
PeerState::Connecting(_) => (live, connecting + 1),
|
|
PeerState::Live(_) => (live + 1, connecting),
|
|
});
|
|
let downloaded = self.stats.downloaded_and_checked.load(Relaxed);
|
|
let remaining = self.needed - downloaded;
|
|
StatsSnapshot {
|
|
have_bytes: self.stats.have.load(Relaxed),
|
|
downloaded_and_checked_bytes: downloaded,
|
|
downloaded_and_checked_pieces: self.stats.downloaded_pieces.load(Relaxed),
|
|
fetched_bytes: self.stats.fetched_bytes.load(Relaxed),
|
|
uploaded_bytes: self.stats.fetched_bytes.load(Relaxed),
|
|
live_peers: live,
|
|
seen_peers: g.peers.seen.len() as u32,
|
|
connecting_peers: connecting,
|
|
time: Instant::now(),
|
|
initially_needed_bytes: self.needed,
|
|
remaining_bytes: remaining,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct PeerHandler {
|
|
state: Arc<TorrentState>,
|
|
addr: SocketAddr,
|
|
spawner: BlockingSpawner,
|
|
}
|
|
|
|
impl PeerConnectionHandler for PeerHandler {
|
|
fn on_received_message(&self, message: Message<ByteBuf<'_>>) -> anyhow::Result<()> {
|
|
match message {
|
|
Message::Request(request) => {
|
|
self.on_download_request(self.addr, request)
|
|
.with_context(|| {
|
|
format!("error handling download request from {}", self.addr)
|
|
})?;
|
|
}
|
|
Message::Bitfield(b) => self.on_bitfield(self.addr, b.clone_to_owned())?,
|
|
Message::Choke => self.on_i_am_choked(self.addr),
|
|
Message::Unchoke => self.on_i_am_unchoked(self.addr),
|
|
Message::Interested => self.on_peer_interested(self.addr),
|
|
Message::Piece(piece) => {
|
|
self.on_received_piece(self.addr, piece)
|
|
.context("error in on_received_piece()")?;
|
|
}
|
|
Message::KeepAlive => {
|
|
debug!("keepalive received from {}", self.addr);
|
|
}
|
|
Message::Have(h) => self.on_have(self.addr, h),
|
|
Message::NotInterested => {
|
|
info!("received \"not interested\", but we don't care yet")
|
|
}
|
|
message => {
|
|
warn!(
|
|
"{}: received unsupported message {:?}, ignoring",
|
|
self.addr, message
|
|
);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn get_have_bytes(&self) -> u64 {
|
|
self.state.stats.have.load(Ordering::Relaxed)
|
|
}
|
|
|
|
fn serialize_bitfield_message_to_buf(&self, buf: &mut Vec<u8>) -> Option<usize> {
|
|
let g = self.state.locked.read();
|
|
let msg = Message::Bitfield(ByteBuf(g.chunks.get_have_pieces().as_raw_slice()));
|
|
let len = msg.serialize(buf, None).unwrap();
|
|
debug!("sending to {}: {:?}, length={}", self.addr, &msg, len);
|
|
Some(len)
|
|
}
|
|
|
|
fn on_handshake(&self, handshake: Handshake) {
|
|
self.state.set_peer_live(self.addr, handshake)
|
|
}
|
|
|
|
fn on_uploaded_bytes(&self, bytes: u32) {
|
|
self.state
|
|
.stats
|
|
.uploaded
|
|
.fetch_add(bytes as u64, Ordering::Relaxed);
|
|
}
|
|
|
|
fn read_chunk(&self, chunk: &crate::lengths::ChunkInfo, buf: &mut [u8]) -> anyhow::Result<()> {
|
|
self.state.file_ops().read_chunk(self.addr, chunk, buf)
|
|
}
|
|
|
|
fn on_extended_handshake(
|
|
&self,
|
|
extended_handshake: &crate::peer_binary_protocol::ExtendedHandshake<ByteBuf>,
|
|
) {
|
|
}
|
|
}
|
|
|
|
impl PeerHandler {
|
|
fn on_download_request(&self, peer_handle: PeerHandle, request: Request) -> anyhow::Result<()> {
|
|
let piece_index = match self.state.lengths.validate_piece_index(request.index) {
|
|
Some(p) => p,
|
|
None => {
|
|
anyhow::bail!(
|
|
"{}: received {:?}, but it is not a valid chunk request (piece index is invalid). Ignoring.",
|
|
peer_handle, request
|
|
);
|
|
}
|
|
};
|
|
let chunk_info = match self.state.lengths.chunk_info_from_received_data(
|
|
piece_index,
|
|
request.begin,
|
|
request.length,
|
|
) {
|
|
Some(d) => d,
|
|
None => {
|
|
anyhow::bail!(
|
|
"{}: received {:?}, but it is not a valid chunk request (chunk data is invalid). Ignoring.",
|
|
peer_handle, request
|
|
);
|
|
}
|
|
};
|
|
|
|
let tx = {
|
|
let g = self.state.locked.read();
|
|
if !g.chunks.is_chunk_ready_to_upload(&chunk_info) {
|
|
anyhow::bail!(
|
|
"got request for a chunk that is not ready to upload. chunk {:?}",
|
|
&chunk_info
|
|
);
|
|
}
|
|
|
|
g.peers.clone_tx(peer_handle).ok_or_else(|| {
|
|
anyhow::anyhow!(
|
|
"peer {} died, dropping chunk that it requested",
|
|
peer_handle
|
|
)
|
|
})?
|
|
};
|
|
|
|
// TODO: this is not super efficient as it does copying multiple times.
|
|
// Theoretically, this could be done in the sending code, so that it reads straight into
|
|
// the send buffer.
|
|
let request = WriterRequest::ReadChunkRequest(chunk_info);
|
|
debug!("sending to {}: {:?}", peer_handle, &request);
|
|
Ok::<_, anyhow::Error>(tx.send(request)?)
|
|
}
|
|
|
|
fn on_have(&self, handle: PeerHandle, have: u32) {
|
|
if let Some(bitfield) = self
|
|
.state
|
|
.locked
|
|
.write()
|
|
.peers
|
|
.get_live_mut(handle)
|
|
.and_then(|l| l.bitfield.as_mut())
|
|
{
|
|
debug!("{}: updated bitfield with have={}", handle, have);
|
|
bitfield.set(have as usize, true)
|
|
}
|
|
}
|
|
|
|
fn on_bitfield(&self, handle: PeerHandle, bitfield: ByteString) -> anyhow::Result<()> {
|
|
if bitfield.len() != self.state.lengths.piece_bitfield_bytes() as usize {
|
|
anyhow::bail!(
|
|
"dropping {} as its bitfield has unexpected size. Got {}, expected {}",
|
|
handle,
|
|
bitfield.len(),
|
|
self.state.lengths.piece_bitfield_bytes(),
|
|
);
|
|
}
|
|
self.state
|
|
.locked
|
|
.write()
|
|
.peers
|
|
.update_bitfield_from_vec(handle, bitfield.0);
|
|
|
|
if !self.state.am_i_interested_in_peer(handle) {
|
|
let tx = self
|
|
.state
|
|
.locked
|
|
.read()
|
|
.peers
|
|
.clone_tx(handle)
|
|
.ok_or_else(|| anyhow::anyhow!("peer closed"))?;
|
|
tx.send(WriterRequest::Message(MessageOwned::Unchoke))
|
|
.context("peer dropped")?;
|
|
tx.send(WriterRequest::Message(MessageOwned::NotInterested))
|
|
.context("peer dropped")?;
|
|
return Ok(());
|
|
}
|
|
|
|
// Additional spawn per peer, not good.
|
|
spawn(
|
|
format!("peer_chunk_requester({})", handle),
|
|
self.clone().task_peer_chunk_requester(handle),
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
async fn task_peer_chunk_requester(self, handle: PeerHandle) -> anyhow::Result<()> {
|
|
let tx = match self.state.locked.read().peers.clone_tx(handle) {
|
|
Some(tx) => tx,
|
|
None => return Ok(()),
|
|
};
|
|
tx.send(WriterRequest::Message(MessageOwned::Unchoke))
|
|
.context("peer dropped")?;
|
|
tx.send(WriterRequest::Message(MessageOwned::Interested))
|
|
.context("peer dropped")?;
|
|
|
|
self.requester(handle).await?;
|
|
Ok::<_, anyhow::Error>(())
|
|
}
|
|
|
|
fn on_i_am_choked(&self, handle: PeerHandle) {
|
|
warn!("we are choked by {}", handle);
|
|
self.state
|
|
.locked
|
|
.write()
|
|
.peers
|
|
.mark_i_am_choked(handle, true);
|
|
}
|
|
|
|
fn on_peer_interested(&self, handle: PeerHandle) {
|
|
debug!("peer {} is interested", handle);
|
|
self.state
|
|
.locked
|
|
.write()
|
|
.peers
|
|
.mark_peer_interested(handle, true);
|
|
}
|
|
|
|
async fn requester(self, handle: PeerHandle) -> anyhow::Result<()> {
|
|
let notify = match self.state.locked.read().peers.get_live(handle) {
|
|
Some(l) => l.have_notify.clone(),
|
|
None => return Ok(()),
|
|
};
|
|
|
|
// TODO: this might dangle, same below.
|
|
#[allow(unused_must_use)]
|
|
{
|
|
timeout(Duration::from_secs(60), notify.notified()).await;
|
|
}
|
|
|
|
loop {
|
|
match self.state.am_i_choked(handle) {
|
|
Some(true) => {
|
|
warn!("we are choked by {}, can't reserve next piece", handle);
|
|
#[allow(unused_must_use)]
|
|
{
|
|
timeout(Duration::from_secs(60), notify.notified()).await;
|
|
}
|
|
continue;
|
|
}
|
|
Some(false) => {}
|
|
None => return Ok(()),
|
|
}
|
|
|
|
let next = match self.state.try_steal_old_slow_piece(handle) {
|
|
Some(next) => next,
|
|
None => match self.state.reserve_next_needed_piece(handle) {
|
|
Some(next) => next,
|
|
None => {
|
|
if self.state.get_left_to_download() == 0 {
|
|
debug!("{}: nothing left to download, closing requester", handle);
|
|
return Ok(());
|
|
}
|
|
|
|
if let Some(piece) = self.state.try_steal_piece(handle) {
|
|
debug!("{}: stole a piece {}", handle, piece);
|
|
piece
|
|
} else {
|
|
debug!("no pieces to request from {}", handle);
|
|
#[allow(unused_must_use)]
|
|
{
|
|
timeout(Duration::from_secs(60), notify.notified()).await;
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
},
|
|
};
|
|
|
|
let tx = match self.state.locked.read().peers.clone_tx(handle) {
|
|
Some(tx) => tx,
|
|
None => return Ok(()),
|
|
};
|
|
let sem = match self.state.locked.read().peers.get_live(handle) {
|
|
Some(live) => live.requests_sem.clone(),
|
|
None => return Ok(()),
|
|
};
|
|
for chunk in self.state.lengths.iter_chunk_infos(next) {
|
|
if self.state.locked.read().chunks.is_chunk_downloaded(&chunk) {
|
|
continue;
|
|
}
|
|
if !self
|
|
.state
|
|
.locked
|
|
.write()
|
|
.peers
|
|
.try_get_live_mut(handle)?
|
|
.inflight_requests
|
|
.insert(InflightRequest::from(&chunk))
|
|
{
|
|
warn!(
|
|
"{}: probably a bug, we already requested {:?}",
|
|
handle, chunk
|
|
);
|
|
continue;
|
|
}
|
|
|
|
let request = Request {
|
|
index: next.get(),
|
|
begin: chunk.offset,
|
|
length: chunk.size,
|
|
};
|
|
sem.acquire().await?.forget();
|
|
|
|
tx.send(WriterRequest::Message(MessageOwned::Request(request)))
|
|
.context("peer dropped")?;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn on_i_am_unchoked(&self, handle: PeerHandle) {
|
|
debug!("we are unchoked by {}", handle);
|
|
let mut g = self.state.locked.write();
|
|
let live = match g.peers.get_live_mut(handle) {
|
|
Some(live) => live,
|
|
None => return,
|
|
};
|
|
live.i_am_choked = false;
|
|
live.have_notify.notify_waiters();
|
|
live.requests_sem.add_permits(16);
|
|
}
|
|
|
|
fn on_received_piece(&self, handle: PeerHandle, piece: Piece<ByteBuf>) -> anyhow::Result<()> {
|
|
let chunk_info = match self.state.lengths.chunk_info_from_received_piece(&piece) {
|
|
Some(i) => i,
|
|
None => {
|
|
anyhow::bail!(
|
|
"peer {} sent us a piece that is invalid {:?}",
|
|
handle,
|
|
&piece,
|
|
);
|
|
}
|
|
};
|
|
|
|
let mut g = self.state.locked.write();
|
|
let h = g.peers.try_get_live_mut(handle)?;
|
|
h.requests_sem.add_permits(1);
|
|
|
|
self.state
|
|
.stats
|
|
.fetched_bytes
|
|
.fetch_add(piece.block.len() as u64, Ordering::Relaxed);
|
|
|
|
if !h
|
|
.inflight_requests
|
|
.remove(&InflightRequest::from(&chunk_info))
|
|
{
|
|
anyhow::bail!(
|
|
"peer {} sent us a piece that we did not ask it for. Requested pieces: {:?}. Got: {:?}", handle, &h.inflight_requests, &piece,
|
|
);
|
|
}
|
|
|
|
let full_piece_download_time = match g.chunks.mark_chunk_downloaded(&piece) {
|
|
Some(ChunkMarkingResult::Completed) => {
|
|
debug!(
|
|
"piece={} done by {}, will write and checksum",
|
|
piece.index, handle
|
|
);
|
|
// This will prevent others from stealing it.
|
|
g.peers
|
|
.remove_inflight_piece(chunk_info.piece_index)
|
|
.map(|t| t.started.elapsed())
|
|
}
|
|
Some(ChunkMarkingResult::PreviouslyCompleted) => {
|
|
// TODO: we might need to send cancellations here.
|
|
debug!(
|
|
"piece={} was done by someone else {}, ignoring",
|
|
piece.index, handle
|
|
);
|
|
return Ok(());
|
|
}
|
|
Some(ChunkMarkingResult::NotCompleted) => None,
|
|
None => {
|
|
anyhow::bail!(
|
|
"bogus data received from {}: {:?}, cannot map this to a chunk, dropping peer",
|
|
handle,
|
|
piece
|
|
);
|
|
}
|
|
};
|
|
|
|
// to prevent deadlocks.
|
|
drop(g);
|
|
|
|
self.spawner
|
|
.spawn_block_in_place(move || {
|
|
let index = piece.index;
|
|
|
|
// TODO: in theory we should unmark the piece as downloaded here. But if there was a disk error, what
|
|
// should we really do? If we unmark it, it will get requested forever...
|
|
//
|
|
// So let's just unwrap and abort.
|
|
self.state
|
|
.file_ops()
|
|
.write_chunk(handle, &piece, &chunk_info)
|
|
.expect("expected to be able to write to disk");
|
|
|
|
let full_piece_download_time = match full_piece_download_time {
|
|
Some(t) => t,
|
|
None => return Ok(()),
|
|
};
|
|
|
|
match self
|
|
.state
|
|
.file_ops()
|
|
.check_piece(handle, chunk_info.piece_index, &chunk_info)
|
|
.with_context(|| format!("error checking piece={}", index))?
|
|
{
|
|
true => {
|
|
let piece_len =
|
|
self.state.lengths.piece_length(chunk_info.piece_index) as u64;
|
|
self.state
|
|
.stats
|
|
.downloaded_and_checked
|
|
.fetch_add(piece_len, Ordering::Relaxed);
|
|
self.state
|
|
.stats
|
|
.have
|
|
.fetch_add(piece_len, Ordering::Relaxed);
|
|
self.state
|
|
.stats
|
|
.downloaded_pieces
|
|
.fetch_add(1, Ordering::Relaxed);
|
|
self.state
|
|
.stats
|
|
.downloaded_pieces
|
|
.fetch_add(1, Ordering::Relaxed);
|
|
self.state.stats.total_piece_download_ms.fetch_add(
|
|
full_piece_download_time.as_millis() as u64,
|
|
Ordering::Relaxed,
|
|
);
|
|
self.state
|
|
.locked
|
|
.write()
|
|
.chunks
|
|
.mark_piece_downloaded(chunk_info.piece_index);
|
|
|
|
debug!(
|
|
"piece={} successfully downloaded and verified from {}",
|
|
index, handle
|
|
);
|
|
|
|
self.state.maybe_transmit_haves(chunk_info.piece_index);
|
|
}
|
|
false => {
|
|
warn!(
|
|
"checksum for piece={} did not validate, came from {}",
|
|
index, handle
|
|
);
|
|
self.state
|
|
.locked
|
|
.write()
|
|
.chunks
|
|
.mark_piece_broken(chunk_info.piece_index);
|
|
}
|
|
};
|
|
Ok::<_, anyhow::Error>(())
|
|
})
|
|
.with_context(|| format!("error processing received chunk {:?}", chunk_info))?;
|
|
Ok(())
|
|
}
|
|
}
|