Split the huge torrent manager into 3 classes. All public access members but ok for the start
This commit is contained in:
Igor Katson
parent
bde18ab734
commit
0d85eb9f2a
5 changed files with 1167 additions and 1091 deletions
|
|
@ -8,7 +8,9 @@ pub mod peer_binary_protocol;
|
|||
pub mod peer_connection;
|
||||
pub mod peer_id;
|
||||
pub mod serde_bencode;
|
||||
pub mod spawn_utils;
|
||||
pub mod torrent_manager;
|
||||
pub mod torrent_metainfo;
|
||||
pub mod torrent_state;
|
||||
pub mod tracker_comms;
|
||||
pub mod type_aliases;
|
||||
|
|
|
|||
|
|
@ -1 +1,585 @@
|
|||
// TODO
|
||||
use std::{
|
||||
net::SocketAddr,
|
||||
sync::{atomic::Ordering, Arc},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use anyhow::Context;
|
||||
use log::{debug, info, trace, warn};
|
||||
use tokio::time::timeout;
|
||||
|
||||
use crate::{
|
||||
buffers::{ByteBuf, ByteString},
|
||||
chunk_tracker::ChunkMarkingResult,
|
||||
clone_to_owned::CloneToOwned,
|
||||
peer_binary_protocol::{
|
||||
Handshake, Message, MessageBorrowed, MessageDeserializeError, MessageOwned, Piece, Request,
|
||||
},
|
||||
peer_id::try_decode_peer_id,
|
||||
spawn_utils::{spawn, spawn_blocking},
|
||||
torrent_state::{InflightRequest, PeerHandle, TorrentState},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PeerConnection {
|
||||
state: Arc<TorrentState>,
|
||||
}
|
||||
|
||||
impl PeerConnection {
|
||||
pub fn new(state: Arc<TorrentState>) -> Self {
|
||||
PeerConnection { state }
|
||||
}
|
||||
pub fn into_state(self) -> Arc<TorrentState> {
|
||||
self.state
|
||||
}
|
||||
pub async fn manage_peer(
|
||||
&self,
|
||||
addr: SocketAddr,
|
||||
handle: PeerHandle,
|
||||
// outgoing_chan_tx: tokio::sync::mpsc::Sender<MessageOwned>,
|
||||
mut outgoing_chan: tokio::sync::mpsc::Receiver<MessageOwned>,
|
||||
) -> anyhow::Result<()> {
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let mut conn = tokio::net::TcpStream::connect(addr)
|
||||
.await
|
||||
.context("error connecting")?;
|
||||
let handshake = Handshake::new(self.state.info_hash, self.state.peer_id);
|
||||
conn.write_all(&handshake.serialize())
|
||||
.await
|
||||
.context("error writing handshake")?;
|
||||
let mut read_buf = vec![0u8; 16384 * 2];
|
||||
let read_bytes = conn
|
||||
.read(&mut read_buf)
|
||||
.await
|
||||
.context("error reading handshake")?;
|
||||
if read_bytes == 0 {
|
||||
anyhow::bail!("bad handshake");
|
||||
}
|
||||
let (h, hlen) = Handshake::deserialize(&read_buf[..read_bytes])
|
||||
.map_err(|e| anyhow::anyhow!("error deserializing handshake: {:?}", e))?;
|
||||
|
||||
let mut read_so_far = 0usize;
|
||||
debug!(
|
||||
"connected peer {}: {:?}",
|
||||
addr,
|
||||
try_decode_peer_id(h.peer_id)
|
||||
);
|
||||
if h.info_hash != self.state.info_hash {
|
||||
anyhow::bail!("info hash does not match");
|
||||
}
|
||||
|
||||
self.state.set_peer_live(handle, h);
|
||||
|
||||
if read_bytes > hlen {
|
||||
read_buf.copy_within(hlen..read_bytes, 0);
|
||||
read_so_far = read_bytes - hlen;
|
||||
}
|
||||
|
||||
let (mut read_half, mut write_half) = tokio::io::split(conn);
|
||||
|
||||
let this = self.clone();
|
||||
let writer = async move {
|
||||
let mut buf = Vec::<u8>::new();
|
||||
let keep_alive_interval = Duration::from_secs(120);
|
||||
|
||||
if this.state.stats.have.load(Ordering::Relaxed) > 0 {
|
||||
let len = {
|
||||
let g = this.state.locked.read();
|
||||
let msg = Message::Bitfield(ByteBuf(g.chunks.get_have_pieces().as_raw_slice()));
|
||||
let len = msg.serialize(&mut buf);
|
||||
debug!("sending to {}: {:?}, length={}", handle, &msg, len);
|
||||
len
|
||||
};
|
||||
|
||||
write_half
|
||||
.write_all(&buf[..len])
|
||||
.await
|
||||
.context("error writing bitfield to peer")?;
|
||||
debug!("sent bitfield to {}", handle);
|
||||
}
|
||||
|
||||
loop {
|
||||
let msg = match timeout(keep_alive_interval, outgoing_chan.recv()).await {
|
||||
Ok(Some(msg)) => msg,
|
||||
Ok(None) => {
|
||||
anyhow::bail!("closing writer, channel closed")
|
||||
}
|
||||
Err(_) => MessageOwned::KeepAlive,
|
||||
};
|
||||
|
||||
let uploaded_add = match &msg {
|
||||
Message::Piece(p) => Some(p.block.len()),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let len = msg.serialize(&mut buf);
|
||||
debug!("sending to {}: {:?}, length={}", handle, &msg, len);
|
||||
|
||||
write_half
|
||||
.write_all(&buf[..len])
|
||||
.await
|
||||
.context("error writing the message to peer")?;
|
||||
|
||||
if let Some(uploaded_add) = uploaded_add {
|
||||
this.state
|
||||
.stats
|
||||
.uploaded
|
||||
.fetch_add(uploaded_add as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
// For type inference.
|
||||
#[allow(unreachable_code)]
|
||||
Ok::<_, anyhow::Error>(())
|
||||
};
|
||||
|
||||
let reader = async move {
|
||||
loop {
|
||||
let message = loop {
|
||||
match MessageBorrowed::deserialize(&read_buf[..read_so_far]) {
|
||||
Ok((msg, size)) => {
|
||||
let msg = msg.clone_to_owned();
|
||||
if read_so_far > size {
|
||||
read_buf.copy_within(size..read_so_far, 0);
|
||||
}
|
||||
read_so_far -= size;
|
||||
break msg;
|
||||
}
|
||||
Err(MessageDeserializeError::NotEnoughData(d, _)) => {
|
||||
if read_buf.len() < read_so_far + d {
|
||||
read_buf.reserve(d);
|
||||
read_buf.resize(read_buf.capacity(), 0);
|
||||
}
|
||||
|
||||
let size = read_half
|
||||
.read(&mut read_buf[read_so_far..])
|
||||
.await
|
||||
.context("error reading from peer")?;
|
||||
if size == 0 {
|
||||
anyhow::bail!(
|
||||
"disconnected while reading, read so far: {}",
|
||||
read_so_far
|
||||
)
|
||||
}
|
||||
read_so_far += size;
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
};
|
||||
|
||||
trace!("received from {}: {:?}", handle, &message);
|
||||
|
||||
match message {
|
||||
Message::Request(request) => {
|
||||
self.on_download_request(handle, request)
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!("error handling download request from {}", handle)
|
||||
})?;
|
||||
}
|
||||
Message::Bitfield(b) => self.on_bitfield(handle, b).await?,
|
||||
Message::Choke => self.on_i_am_choked(handle),
|
||||
Message::Unchoke => self.on_i_am_unchoked(handle),
|
||||
Message::Interested => {
|
||||
warn!(
|
||||
"{} is interested, but support for interested messages not implemented",
|
||||
handle
|
||||
)
|
||||
}
|
||||
Message::Piece(piece) => {
|
||||
self.on_received_piece(handle, piece)
|
||||
.context("error in on_received_piece()")?;
|
||||
}
|
||||
Message::KeepAlive => {
|
||||
debug!("keepalive received from {}", handle);
|
||||
}
|
||||
Message::Have(h) => self.on_have(handle, h),
|
||||
Message::NotInterested => {
|
||||
info!("received \"not interested\", but we don't care yet")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For type inference.
|
||||
#[allow(unreachable_code)]
|
||||
Ok::<_, anyhow::Error>(())
|
||||
};
|
||||
|
||||
let r = tokio::select! {
|
||||
r = reader => {r}
|
||||
r = writer => {r}
|
||||
};
|
||||
debug!("{}: either reader or writer are done, exiting", handle);
|
||||
r
|
||||
}
|
||||
|
||||
async 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 state = self.state.clone();
|
||||
let chunk = spawn_blocking(
|
||||
format!(
|
||||
"read_chunk_blocking(peer={}, chunk_info={:?}",
|
||||
peer_handle, &chunk_info
|
||||
),
|
||||
move || state.read_chunk_blocking(peer_handle, chunk_info),
|
||||
)
|
||||
.await??;
|
||||
|
||||
let tx = self
|
||||
.state
|
||||
.locked
|
||||
.read()
|
||||
.peers
|
||||
.clone_tx(peer_handle)
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"peer {} died, dropping chunk that it requested",
|
||||
peer_handle
|
||||
)
|
||||
})?;
|
||||
let message = Message::Piece(Piece::from_vec(
|
||||
chunk_info.piece_index.get(),
|
||||
chunk_info.offset,
|
||||
chunk,
|
||||
));
|
||||
info!("sending to {}: {:?}", peer_handle, &message);
|
||||
Ok::<_, anyhow::Error>(tx.send(message).await?)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
async 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(MessageOwned::Unchoke)
|
||||
.await
|
||||
.context("peer dropped")?;
|
||||
tx.send(MessageOwned::NotInterested)
|
||||
.await
|
||||
.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(MessageOwned::Unchoke)
|
||||
.await
|
||||
.context("peer dropped")?;
|
||||
tx.send(MessageOwned::Interested)
|
||||
.await
|
||||
.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);
|
||||
}
|
||||
|
||||
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.reserve_next_needed_piece(handle) {
|
||||
Some(next) => next,
|
||||
None => {
|
||||
if self.state.get_left_to_download() == 0 {
|
||||
info!("{}: nothing left to download, closing requester", handle);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if let Some(piece) = self.state.try_steal_piece(handle) {
|
||||
info!("{}: stole a piece {}", handle, piece);
|
||||
piece
|
||||
} else {
|
||||
info!("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.outstanding_requests.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(MessageOwned::Request(request))
|
||||
.await
|
||||
.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.outstanding_requests.add_permits(16);
|
||||
}
|
||||
|
||||
fn on_received_piece(
|
||||
&self,
|
||||
handle: PeerHandle,
|
||||
piece: Piece<ByteString>,
|
||||
) -> 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.outstanding_requests.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 should_checksum = 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);
|
||||
true
|
||||
}
|
||||
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) => false,
|
||||
None => {
|
||||
anyhow::bail!(
|
||||
"bogus data received from {}: {:?}, cannot map this to a chunk, dropping peer",
|
||||
handle,
|
||||
piece
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let this = self.clone();
|
||||
|
||||
spawn_blocking(
|
||||
format!(
|
||||
"write_and_check(piece={}, peer={}, block={:?})",
|
||||
piece.index, handle, &piece
|
||||
),
|
||||
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...
|
||||
this.state
|
||||
.write_chunk_blocking(handle, &piece, &chunk_info)?;
|
||||
|
||||
if !should_checksum {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
match this
|
||||
.state
|
||||
.check_piece_blocking(handle, chunk_info.piece_index, &chunk_info)
|
||||
.with_context(|| format!("error checking piece={}", index))?
|
||||
{
|
||||
true => {
|
||||
let piece_len =
|
||||
this.state.lengths.piece_length(chunk_info.piece_index) as u64;
|
||||
this.state
|
||||
.stats
|
||||
.downloaded_and_checked
|
||||
.fetch_add(piece_len, Ordering::Relaxed);
|
||||
this.state
|
||||
.stats
|
||||
.have
|
||||
.fetch_add(piece_len, Ordering::Relaxed);
|
||||
this.state
|
||||
.locked
|
||||
.write()
|
||||
.chunks
|
||||
.mark_piece_downloaded(chunk_info.piece_index);
|
||||
|
||||
debug!(
|
||||
"piece={} successfully downloaded and verified from {}",
|
||||
index, handle
|
||||
);
|
||||
let state_clone = this.state.clone();
|
||||
spawn("transmit haves", async move {
|
||||
state_clone.task_transmit_haves(piece.index).await
|
||||
});
|
||||
}
|
||||
false => {
|
||||
warn!(
|
||||
"checksum for piece={} did not validate, came from {}",
|
||||
index, handle
|
||||
);
|
||||
this.state
|
||||
.locked
|
||||
.write()
|
||||
.chunks
|
||||
.mark_piece_broken(chunk_info.piece_index);
|
||||
}
|
||||
};
|
||||
Ok::<_, anyhow::Error>(())
|
||||
},
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
37
crates/librqbit/src/spawn_utils.rs
Normal file
37
crates/librqbit/src/spawn_utils.rs
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
use std::fmt::Display;
|
||||
|
||||
use log::{debug, error};
|
||||
|
||||
pub fn spawn<N: Display + 'static + Send>(
|
||||
name: N,
|
||||
fut: impl std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
|
||||
) {
|
||||
debug!("starting task \"{}\"", &name);
|
||||
tokio::spawn(async move {
|
||||
match fut.await {
|
||||
Ok(_) => {
|
||||
debug!("task \"{}\" finished", &name);
|
||||
}
|
||||
Err(e) => {
|
||||
error!("error in task \"{}\": {:#}", &name, e)
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub fn spawn_blocking<T: Send + Sync + 'static, N: Display + 'static + Send>(
|
||||
name: N,
|
||||
f: impl FnOnce() -> anyhow::Result<T> + Send + 'static,
|
||||
) -> tokio::task::JoinHandle<anyhow::Result<T>> {
|
||||
debug!("starting blocking task \"{}\"", name);
|
||||
tokio::task::spawn_blocking(move || match f() {
|
||||
Ok(v) => {
|
||||
debug!("blocking task \"{}\" finished", name);
|
||||
Ok(v)
|
||||
}
|
||||
Err(e) => {
|
||||
error!("error in blocking task \"{}\": {:#}", name, &e);
|
||||
Err(e)
|
||||
}
|
||||
})
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
509
crates/librqbit/src/torrent_state.rs
Normal file
509
crates/librqbit/src/torrent_state.rs
Normal file
|
|
@ -0,0 +1,509 @@
|
|||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
fs::File,
|
||||
io::{Read, Seek, SeekFrom, Write},
|
||||
net::SocketAddr,
|
||||
sync::{
|
||||
atomic::{AtomicU64, Ordering},
|
||||
Arc,
|
||||
},
|
||||
};
|
||||
|
||||
use anyhow::Context;
|
||||
use futures::{stream::FuturesUnordered, StreamExt};
|
||||
use log::{debug, warn};
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
use tokio::sync::{mpsc::Sender, Notify, Semaphore};
|
||||
|
||||
use crate::{
|
||||
buffers::ByteString,
|
||||
chunk_tracker::ChunkTracker,
|
||||
file_checking::update_hash_from_file,
|
||||
lengths::{ChunkInfo, Lengths, ValidPieceIndex},
|
||||
peer_binary_protocol::{Handshake, Message, MessageOwned, Piece},
|
||||
torrent_metainfo::TorrentMetaV1Owned,
|
||||
type_aliases::BF,
|
||||
};
|
||||
|
||||
pub type PeerHandle = SocketAddr;
|
||||
|
||||
pub enum PeerState {
|
||||
Connecting(SocketAddr),
|
||||
Live(LivePeerState),
|
||||
}
|
||||
|
||||
#[derive(Debug, Hash, PartialEq, Eq)]
|
||||
pub struct InflightRequest {
|
||||
pub piece: ValidPieceIndex,
|
||||
pub chunk: u32,
|
||||
}
|
||||
|
||||
impl From<&ChunkInfo> for InflightRequest {
|
||||
fn from(c: &ChunkInfo) -> Self {
|
||||
Self {
|
||||
piece: c.piece_index,
|
||||
chunk: c.chunk_index,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct LivePeerState {
|
||||
#[allow(unused)]
|
||||
pub peer_id: [u8; 20],
|
||||
pub i_am_choked: bool,
|
||||
#[allow(unused)]
|
||||
pub peer_choked: bool,
|
||||
#[allow(unused)]
|
||||
pub peer_interested: bool,
|
||||
pub outstanding_requests: Arc<Semaphore>,
|
||||
pub have_notify: Arc<Notify>,
|
||||
pub bitfield: Option<BF>,
|
||||
pub inflight_requests: HashSet<InflightRequest>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct PeerStates {
|
||||
states: HashMap<PeerHandle, PeerState>,
|
||||
seen_peers: HashSet<SocketAddr>,
|
||||
inflight_pieces: HashSet<ValidPieceIndex>,
|
||||
tx: HashMap<PeerHandle, Arc<tokio::sync::mpsc::Sender<MessageOwned>>>,
|
||||
}
|
||||
|
||||
#[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: tokio::sync::mpsc::Sender<MessageOwned>,
|
||||
) -> Option<PeerHandle> {
|
||||
if self.seen_peers.contains(&addr) {
|
||||
return None;
|
||||
}
|
||||
let handle = self.add(addr, tx)?;
|
||||
self.seen_peers.insert(addr);
|
||||
Some(handle)
|
||||
}
|
||||
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: tokio::sync::mpsc::Sender<MessageOwned>,
|
||||
) -> 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> {
|
||||
match self.states.get_mut(&handle) {
|
||||
Some(PeerState::Live(live)) => {
|
||||
let prev = live.i_am_choked;
|
||||
live.i_am_choked = is_choked;
|
||||
return Some(prev);
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
pub fn update_bitfield_from_vec(
|
||||
&mut self,
|
||||
handle: PeerHandle,
|
||||
bitfield: Vec<u8>,
|
||||
) -> Option<Option<BF>> {
|
||||
match self.states.get_mut(&handle) {
|
||||
Some(PeerState::Live(live)) => {
|
||||
let bitfield = BF::from_vec(bitfield);
|
||||
let prev = live.bitfield.take();
|
||||
live.bitfield = Some(bitfield);
|
||||
Some(prev)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
pub fn clone_tx(&self, handle: PeerHandle) -> Option<Arc<Sender<MessageOwned>>> {
|
||||
Some(self.tx.get(&handle)?.clone())
|
||||
}
|
||||
pub fn remove_inflight_piece(&mut self, piece: ValidPieceIndex) -> bool {
|
||||
self.inflight_pieces.remove(&piece)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TorrentStateLocked {
|
||||
pub peers: PeerStates,
|
||||
pub chunks: ChunkTracker,
|
||||
}
|
||||
|
||||
pub struct AtomicStats {
|
||||
pub have: AtomicU64,
|
||||
pub downloaded_and_checked: AtomicU64,
|
||||
pub uploaded: AtomicU64,
|
||||
pub fetched_bytes: AtomicU64,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
impl TorrentState {
|
||||
pub fn read_chunk_blocking(
|
||||
&self,
|
||||
who_sent: PeerHandle,
|
||||
chunk_info: ChunkInfo,
|
||||
) -> anyhow::Result<Vec<u8>> {
|
||||
let mut absolute_offset = self.lengths.chunk_absolute_offset(&chunk_info);
|
||||
let mut result_buf = vec![0u8; chunk_info.size as usize];
|
||||
let mut buf = &mut result_buf[..];
|
||||
|
||||
for (file_idx, file_len) in self.torrent.info.iter_file_lengths().enumerate() {
|
||||
if absolute_offset > file_len {
|
||||
absolute_offset -= file_len;
|
||||
continue;
|
||||
}
|
||||
let file_remaining_len = file_len - absolute_offset;
|
||||
let to_read_in_file = std::cmp::min(file_remaining_len, buf.len() as u64) as usize;
|
||||
|
||||
let mut file_g = self.files[file_idx].lock();
|
||||
debug!(
|
||||
"piece={}, handle={}, file_idx={}, seeking to {}. To read chunk: {:?}",
|
||||
chunk_info.piece_index, who_sent, file_idx, absolute_offset, &chunk_info
|
||||
);
|
||||
file_g
|
||||
.seek(SeekFrom::Start(absolute_offset))
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"error seeking to {}, file id: {}",
|
||||
absolute_offset, file_idx
|
||||
)
|
||||
})?;
|
||||
file_g
|
||||
.read_exact(&mut buf[..to_read_in_file])
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"error reading {} bytes, file_id: {}",
|
||||
file_idx, to_read_in_file
|
||||
)
|
||||
})?;
|
||||
|
||||
buf = &mut buf[to_read_in_file..];
|
||||
|
||||
if buf.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
absolute_offset = 0;
|
||||
}
|
||||
|
||||
return Ok(result_buf);
|
||||
}
|
||||
|
||||
pub fn get_next_needed_piece(&self, peer_handle: PeerHandle) -> Option<ValidPieceIndex> {
|
||||
let g = self.locked.read();
|
||||
let bf = match g.peers.states.get(&peer_handle)? {
|
||||
PeerState::Live(l) => l.bitfield.as_ref()?,
|
||||
_ => return None,
|
||||
};
|
||||
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
|
||||
.states
|
||||
.get(&peer_handle)
|
||||
.and_then(|s| match s {
|
||||
PeerState::Live(l) => Some(l.i_am_choked),
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
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);
|
||||
g.chunks.reserve_needed_piece(n);
|
||||
Some(n)
|
||||
}
|
||||
|
||||
pub fn check_piece_blocking(
|
||||
&self,
|
||||
who_sent: PeerHandle,
|
||||
piece_index: ValidPieceIndex,
|
||||
last_received_chunk: &ChunkInfo,
|
||||
) -> anyhow::Result<bool> {
|
||||
let mut h = sha1::Sha1::new();
|
||||
let piece_length = self.lengths.piece_length(piece_index);
|
||||
let mut absolute_offset = self.lengths.piece_offset(piece_index);
|
||||
let mut buf = vec![0u8; std::cmp::min(65536, piece_length as usize)];
|
||||
|
||||
let mut piece_remaining_bytes = piece_length as usize;
|
||||
|
||||
for (file_idx, (name, file_len)) in
|
||||
self.torrent.info.iter_filenames_and_lengths().enumerate()
|
||||
{
|
||||
if absolute_offset > file_len {
|
||||
absolute_offset -= file_len;
|
||||
continue;
|
||||
}
|
||||
let file_remaining_len = file_len - absolute_offset;
|
||||
|
||||
let to_read_in_file =
|
||||
std::cmp::min(file_remaining_len, piece_remaining_bytes as u64) as usize;
|
||||
let mut file_g = self.files[file_idx].lock();
|
||||
debug!(
|
||||
"piece={}, handle={}, file_idx={}, seeking to {}. Last received chunk: {:?}",
|
||||
piece_index, who_sent, file_idx, absolute_offset, &last_received_chunk
|
||||
);
|
||||
file_g
|
||||
.seek(SeekFrom::Start(absolute_offset))
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"error seeking to {}, file id: {}",
|
||||
absolute_offset, file_idx
|
||||
)
|
||||
})?;
|
||||
update_hash_from_file(&mut file_g, &mut h, &mut buf, to_read_in_file).with_context(
|
||||
|| {
|
||||
format!(
|
||||
"error reading {} bytes, file_id: {} (\"{:?}\")",
|
||||
to_read_in_file, file_idx, name
|
||||
)
|
||||
},
|
||||
)?;
|
||||
|
||||
piece_remaining_bytes -= to_read_in_file;
|
||||
|
||||
if piece_remaining_bytes == 0 {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
absolute_offset = 0;
|
||||
}
|
||||
|
||||
match self.torrent.info.compare_hash(piece_index.get(), &h) {
|
||||
Some(true) => {
|
||||
debug!("piece={} hash matches", piece_index);
|
||||
Ok(true)
|
||||
}
|
||||
Some(false) => {
|
||||
warn!("the piece={} hash does not match", piece_index);
|
||||
Ok(false)
|
||||
}
|
||||
None => {
|
||||
// this is probably a bug?
|
||||
warn!("compare_hash() did not find the piece");
|
||||
anyhow::bail!("compare_hash() did not find the piece");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn am_i_interested_in_peer(&self, handle: PeerHandle) -> bool {
|
||||
self.get_next_needed_piece(handle).is_some()
|
||||
}
|
||||
|
||||
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
|
||||
.iter()
|
||||
.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 {
|
||||
peer_id: h.peer_id,
|
||||
i_am_choked: true,
|
||||
peer_choked: true,
|
||||
peer_interested: false,
|
||||
bitfield: None,
|
||||
have_notify: Arc::new(Notify::new()),
|
||||
outstanding_requests: Arc::new(Semaphore::new(0)),
|
||||
inflight_requests: Default::default(),
|
||||
});
|
||||
}
|
||||
_ => {
|
||||
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 write_chunk_blocking(
|
||||
&self,
|
||||
who_sent: PeerHandle,
|
||||
data: &Piece<ByteString>,
|
||||
chunk_info: &ChunkInfo,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut buf = data.block.as_ref();
|
||||
let mut absolute_offset = self.lengths.chunk_absolute_offset(&chunk_info);
|
||||
|
||||
for (file_idx, (name, file_len)) in
|
||||
self.torrent.info.iter_filenames_and_lengths().enumerate()
|
||||
{
|
||||
if absolute_offset > file_len {
|
||||
absolute_offset -= file_len;
|
||||
continue;
|
||||
}
|
||||
|
||||
let remaining_len = file_len - absolute_offset;
|
||||
let to_write = std::cmp::min(buf.len(), remaining_len as usize);
|
||||
|
||||
let mut file_g = self.files[file_idx].lock();
|
||||
debug!(
|
||||
"piece={}, chunk={:?}, handle={}, begin={}, file={}, writing {} bytes at {}",
|
||||
chunk_info.piece_index,
|
||||
chunk_info,
|
||||
who_sent,
|
||||
chunk_info.offset,
|
||||
file_idx,
|
||||
to_write,
|
||||
absolute_offset
|
||||
);
|
||||
file_g
|
||||
.seek(SeekFrom::Start(absolute_offset))
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"error seeking to {} in file {} (\"{:?}\")",
|
||||
absolute_offset, file_idx, name
|
||||
)
|
||||
})?;
|
||||
file_g
|
||||
.write_all(&buf[..to_write])
|
||||
.with_context(|| format!("error writing to file {} (\"{:?}\")", file_idx, name))?;
|
||||
buf = &buf[to_write..];
|
||||
if buf.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
absolute_offset = 0;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// TODO: this is a task per chunk, not good
|
||||
pub async fn task_transmit_haves(&self, index: u32) -> anyhow::Result<()> {
|
||||
let mut unordered = FuturesUnordered::new();
|
||||
|
||||
for weak in self
|
||||
.locked
|
||||
.read()
|
||||
.peers
|
||||
.tx
|
||||
.values()
|
||||
.map(|v| Arc::downgrade(v))
|
||||
{
|
||||
unordered.push(async move {
|
||||
if let Some(tx) = weak.upgrade() {
|
||||
if tx.send(Message::Have(index)).await.is_err() {
|
||||
// whatever
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
while unordered.next().await.is_some() {}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue