1use std::{
2 cell::OnceCell,
3 collections::VecDeque,
4 io::{self},
5 net::{SocketAddrV4, TcpListener},
6 path::{Path, PathBuf},
7 sync::mpsc::{Receiver, Sender},
8 time::{Duration, Instant},
9};
10
11use crate::{
12 buf_pool::{Buffer, BufferPool},
13 event_loop::{ConnectionId, EventLoop},
14 file_store::DiskOpType,
15 peer_comm::{extended_protocol::MetadataProgress, peer_protocol::PeerId},
16};
17use crate::{
18 file_store::DiskOp,
19 piece_selector::{DownloadedPiece, Piece, PieceSelector, SUBPIECE_SIZE, Subpiece},
20};
21use crate::{file_store::FileStore, peer_connection::PeerConnection};
22use ahash::HashSetExt;
23use bitvec::{boxed::BitBox, order::Msb0, vec::BitVec};
24use heapless::spsc::Producer;
25use io_uring::IoUring;
26use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
27use slotmap::SlotMap;
28use thiserror::Error;
29
30use crate::peer_connection::DisconnectReason;
31
32use crate::TorrentMetadata;
33
34#[derive(Error, Debug)]
35pub enum Error {
36 #[error("Encountered IO issue: {0}")]
37 Io(#[from] io::Error),
38 #[error("Peer provider disconnected")]
39 PeerProviderDisconnect,
40}
41
42pub const CQE_WAIT_TIME_NS: u32 = 150_000_000;
43
44#[derive(Debug, Clone, Copy, serde::Deserialize, serde::Serialize)]
46#[serde(default)]
47pub struct Config {
48 pub max_connections: usize,
50 pub max_reported_outstanding_requests: u64,
55 pub max_unchoked: u32,
57 pub num_ticks_before_unchoke_recalc: u32,
61 pub num_ticks_before_optimistic_unchoke_recalc: u32,
68 pub seeding_piece_quota: u32,
72 pub cq_size: u32,
74 pub sq_size: u32,
76 pub completion_event_want: usize,
80 pub network_read_buffer_size: usize,
82 pub network_write_buffer_size: usize,
84 pub write_buffer_pool_size: usize,
86 pub read_buffer_pool_size: usize,
88}
89
90impl Default for Config {
91 fn default() -> Self {
92 Self {
93 max_connections: 128,
94 max_reported_outstanding_requests: 512,
95 max_unchoked: 8,
96 num_ticks_before_unchoke_recalc: 15,
97 num_ticks_before_optimistic_unchoke_recalc: 30,
98 seeding_piece_quota: 20,
99 cq_size: 4096,
100 sq_size: 4096,
101 completion_event_want: 32,
102 network_read_buffer_size: (SUBPIECE_SIZE * 2) as usize,
103 read_buffer_pool_size: 512,
104 network_write_buffer_size: (SUBPIECE_SIZE + 4096) as usize,
105 write_buffer_pool_size: 128,
106 }
107 }
108}
109
110pub struct Torrent {
114 our_id: PeerId,
115 state: State,
116}
117
118impl Torrent {
119 pub fn new(our_id: PeerId, state: State) -> Self {
122 Self { our_id, state }
123 }
124
125 pub fn is_complete(&self) -> bool {
131 self.state
132 .torrent_state
133 .as_ref()
134 .is_some_and(|state| state.is_complete)
135 }
136
137 pub fn start(
144 &mut self,
145 event_tx: Producer<TorrentEvent>,
146 command_rc: Receiver<Command>,
147 listener: TcpListener,
148 ) -> Result<(), Error> {
149 let ring: IoUring = IoUring::builder()
151 .setup_single_issuer()
152 .setup_clamp()
153 .setup_cqsize(self.state.config.cq_size)
154 .setup_defer_taskrun()
155 .setup_coop_taskrun()
156 .build(self.state.config.sq_size)
157 .unwrap();
158 let events = SlotMap::with_capacity_and_key(self.state.config.cq_size as usize);
159 let mut event_loop = EventLoop::new(self.our_id, events, &self.state.config);
160 event_loop.run(ring, &mut self.state, event_tx, command_rc, listener)
161 }
162}
163
164#[derive(Debug)]
166pub enum Command {
167 ConnectToPeers(Vec<SocketAddrV4>),
170 Stop,
172 Pause,
174 Resume,
178}
179
180#[derive(Debug)]
182pub struct PeerMetrics {
183 pub download_throughput: u64,
185 pub upload_throughput: u64,
187 pub metadata_progress: Option<MetadataProgress>,
189}
190
191#[derive(Debug)]
193pub enum TorrentEvent {
194 TorrentComplete,
199 MetadataComplete(Box<TorrentMetadata>),
203 Running { port: u16 },
208 Paused,
210 TorrentMetrics {
212 pieces_completed: usize,
214 pieces_allocated: usize,
217 peer_metrics: Vec<PeerMetrics>,
219 },
220}
221
222pub struct InitializedState {
223 pub piece_selector: PieceSelector,
224 pub num_unchoked: u32,
225 pub config: Config,
226 pub ticks_to_recalc_unchoke: u32,
227 pub ticks_to_recalc_optimistic_unchoke: u32,
228 pub downloaded_piece_rc: Receiver<DownloadedPiece>,
229 pub downloaded_piece_tx: Sender<DownloadedPiece>,
230 pub pieces: Vec<Option<Piece>>,
231 pub file_store: FileStore,
232 pub piece_buffer_pool: BufferPool,
233 pub is_complete: bool,
234}
235
236impl InitializedState {
237 pub fn new(root: &Path, metadata: &TorrentMetadata, config: Config) -> io::Result<Self> {
238 let mut pieces = Vec::with_capacity(metadata.pieces.len());
239 for _ in 0..metadata.pieces.len() {
240 pieces.push(None);
241 }
242 let (tx, rc) = std::sync::mpsc::channel();
243 Ok(Self {
244 piece_selector: PieceSelector::new(metadata),
245 num_unchoked: 0,
246 config,
247 ticks_to_recalc_unchoke: config.num_ticks_before_unchoke_recalc,
248 ticks_to_recalc_optimistic_unchoke: config.num_ticks_before_optimistic_unchoke_recalc,
249 downloaded_piece_rc: rc,
250 downloaded_piece_tx: tx,
251 pieces,
252 is_complete: false,
253 piece_buffer_pool: BufferPool::new("pieces", 256, metadata.piece_length as usize),
254 file_store: FileStore::new(root, metadata)?,
255 })
256 }
257
258 #[allow(dead_code)]
259 pub fn num_allocated(&self) -> usize {
260 self.pieces
261 .iter()
262 .filter(|piece| piece.as_ref().is_some_and(|piece| piece.ref_count > 0))
263 .count()
264 }
265
266 #[inline]
267 pub fn num_pieces(&self) -> usize {
268 self.pieces.len()
269 }
270
271 pub(crate) fn complete_piece(
272 &mut self,
273 piece_idx: i32,
274 connections: &mut SlotMap<ConnectionId, PeerConnection>,
275 event_tx: &mut Producer<'_, TorrentEvent>,
276 piece_buffer: Buffer,
277 ) {
278 self.piece_buffer_pool.return_buffer(piece_buffer);
279 self.piece_selector.mark_complete(piece_idx as usize);
280 if !self.is_complete && self.piece_selector.completed_all() {
281 log::info!("Torrent complete!");
282 self.is_complete = true;
283 if event_tx.enqueue(TorrentEvent::TorrentComplete).is_err() {
284 log::error!("Torrent completion event missed");
285 }
286 for (_, peer) in connections.iter_mut() {
289 peer.not_interested();
290 if peer.is_upload_only {
294 peer.pending_disconnect = Some(DisconnectReason::RedundantConnection);
295 }
296 for (_, extension) in peer.extensions.iter_mut() {
298 extension.on_torrent_complete(&mut peer.outgoing_msgs_buffer);
299 }
300 }
301 }
302 for (conn_id, peer) in connections.iter_mut() {
303 if let Some(bitfield) = self.piece_selector.interesting_peer_pieces(conn_id)
304 && !bitfield.any()
305 && peer.is_interesting
306 && peer.queued.is_empty()
307 && peer.inflight.is_empty()
308 {
309 peer.not_interested();
311 if peer.is_upload_only {
315 peer.pending_disconnect = Some(DisconnectReason::RedundantConnection);
316 }
317 }
318 peer.have(piece_idx);
319 }
320 log::debug!("Piece {piece_idx} completed!");
321 }
322
323 pub(crate) fn queue_disk_write_for_downloaded_pieces(
325 &mut self,
326 pending_disk_operations: &mut Vec<DiskOp>,
327 ) {
328 while let Ok(completed_piece) = self.downloaded_piece_rc.try_recv() {
329 if completed_piece.hash_matched {
330 let piece_len = self.piece_selector.piece_len(completed_piece.index as i32);
331 self.file_store.queue_piece_disk_operation(
332 completed_piece.index as i32,
333 completed_piece.buffer,
334 piece_len as usize,
335 DiskOpType::Write,
336 pending_disk_operations,
337 );
338 } else {
339 self.piece_selector
342 .mark_not_downloaded(completed_piece.index);
343 self.piece_buffer_pool.return_buffer(completed_piece.buffer);
344 log::error!("Piece hash didn't match expected hash!");
347 self.piece_selector
348 .mark_not_allocated(completed_piece.index as i32, completed_piece.conn_id);
349 }
350 }
351 }
352
353 pub fn allocate_piece(&mut self, index: i32, conn_id: ConnectionId) -> VecDeque<Subpiece> {
355 log::debug!("Allocating piece: conn_id: {conn_id:?}, index: {index}");
356 self.piece_selector.mark_allocated(index, conn_id);
357 match &mut self.pieces[index as usize] {
358 Some(allocated_piece) => allocated_piece.allocate_remaining_subpieces(),
359 None => {
360 let length = self.piece_selector.piece_len(index);
361 let buffer = self.piece_buffer_pool.get_buffer();
362 let mut piece = Piece::new(index, length, buffer);
363 let subpieces = piece.allocate_remaining_subpieces();
364 self.pieces[index as usize] = Some(piece);
365 subpieces
366 }
367 }
368 }
369
370 pub fn deallocate_piece(&mut self, index: i32, conn_id: ConnectionId) {
372 log::debug!("Deallocating piece: conn_id: {conn_id:?}, index: {index}");
373 self.piece_selector
376 .update_peer_piece_intrest(conn_id, index as usize);
377 if let Some(piece) = self.pieces[index as usize].take_if(|piece| {
381 piece.ref_count = piece.ref_count.saturating_sub(1);
382 piece.ref_count == 0
383 }) {
384 log::debug!("Marked as not allocated: conn_id: {conn_id:?}, index: {index}");
385 self.piece_buffer_pool.return_buffer(piece.into_buffer());
386 self.piece_selector.mark_not_allocated(index, conn_id);
387 }
388 }
389
390 pub fn can_preemtively_unchoke(&self) -> bool {
391 self.num_unchoked < self.config.max_unchoked
392 }
393
394 pub fn recalculate_unchokes(
398 &mut self,
399 connections: &mut SlotMap<ConnectionId, PeerConnection>,
400 ) {
401 struct ComparisonData {
402 is_choking: bool,
403 uploaded_since_unchoked: u64,
404 downloaded_in_last_round: u64,
405 uploaded_in_last_round: u64,
406 last_unchoked: Option<Instant>,
407 }
408 log::info!("Recalculating unchokes");
409 let mut peers = Vec::with_capacity(connections.len());
410 for (id, peer) in connections.iter_mut() {
411 if !peer.peer_interested || peer.pending_disconnect.is_some() {
412 peer.network_stats.reset_round();
413 if !peer.is_choking {
414 if peer.optimistically_unchoked {
415 peer.optimistically_unchoked = false;
416 self.ticks_to_recalc_optimistic_unchoke = 0;
418 }
419 peer.choke(self);
420 }
421
422 continue;
423 }
424 peers.push((
425 id,
426 ComparisonData {
427 is_choking: peer.is_choking,
428 uploaded_since_unchoked: peer.network_stats.upload_since_unchoked,
429 downloaded_in_last_round: peer.network_stats.downloaded_in_last_round,
430 uploaded_in_last_round: peer.network_stats.uploaded_in_last_round,
431 last_unchoked: peer.last_unchoked,
432 },
433 ));
434 }
435 if !self.is_complete {
436 peers.sort_unstable_by(|(_, a), (_, b)| {
438 a.downloaded_in_last_round
439 .cmp(&b.downloaded_in_last_round)
440 .reverse()
441 });
442 } else {
443 let piece_length = self.piece_selector.avg_piece_length();
444 let quota_bytes = (piece_length * self.config.seeding_piece_quota) as u64;
445 peers.sort_unstable_by(|(_, a), (_, b)| {
447 let a_quota_complete = !a.is_choking
448 && a.uploaded_since_unchoked > quota_bytes
449 && a.last_unchoked
450 .is_some_and(|time| time.elapsed() > Duration::from_mins(1));
451 let b_quota_complete = !b.is_choking
452 && b.uploaded_since_unchoked > quota_bytes
453 && b.last_unchoked
454 .is_some_and(|time| time.elapsed() > Duration::from_mins(1));
455 if a_quota_complete != b_quota_complete {
456 a_quota_complete.cmp(&b_quota_complete)
457 } else if a.uploaded_in_last_round != b.uploaded_in_last_round {
458 a.uploaded_in_last_round
459 .cmp(&b.uploaded_in_last_round)
460 .reverse()
461 } else {
462 a.last_unchoked
463 .map_or(Duration::MAX, |time| time.elapsed())
464 .cmp(&b.last_unchoked.map_or(Duration::MAX, |time| time.elapsed()))
465 .reverse()
466 }
467 });
468 }
469 let optimistic_unchoke_slots = std::cmp::max(1, self.config.max_unchoked / 5);
470 let mut remaining_unchoke_slots = self.config.max_unchoked - optimistic_unchoke_slots;
471 for (id, _) in peers {
472 let peer = &mut connections[id];
473 peer.network_stats.reset_round();
474 if remaining_unchoke_slots > 0 {
475 if peer.is_choking {
476 log::debug!(
477 "Peer[{}] now unchoked after recalculating throughputs",
478 peer.peer_id
479 );
480 peer.unchoke(self);
481 }
482 remaining_unchoke_slots -= 1;
483 if peer.optimistically_unchoked {
484 log::trace!(
485 "Peer[{}] previously optimistically unchoked, promoted to normal unchoke",
486 peer.peer_id
487 );
488 peer.optimistically_unchoked = false;
491 self.ticks_to_recalc_optimistic_unchoke = 0;
493 }
494 } else if !peer.is_choking && !peer.optimistically_unchoked {
495 log::debug!(
496 "Peer[{}] no longer unchoked after recalculating throughputs",
497 peer.peer_id
498 );
499 peer.choke(self);
500 }
501 }
502 }
503
504 pub fn recalculate_optimistic_unchokes(
507 &mut self,
508 connections: &mut SlotMap<ConnectionId, PeerConnection>,
509 ) {
510 log::info!("Recalculating optimistic unchokes");
511 let num_opt_unchoked = std::cmp::max(1, self.config.max_unchoked / 5) as usize;
512 let mut previously_opt_unchoked = ahash::HashSet::with_capacity(num_opt_unchoked);
513 let mut candidates = Vec::with_capacity(self.config.max_unchoked as usize);
514 for (id, peer) in connections.iter_mut() {
515 if peer.optimistically_unchoked {
516 previously_opt_unchoked.insert(id);
517 }
518 if peer.pending_disconnect.is_none()
519 && peer.peer_interested
520 && (peer.is_choking || peer.optimistically_unchoked)
521 {
522 candidates.push((
523 id,
524 peer.last_optimistically_unchoked
525 .map_or(u64::MAX, |time| time.elapsed().as_secs()),
526 ));
527 }
528 }
529 candidates.sort_unstable_by(|(_, a), (_, b)| a.cmp(b).reverse());
531 for (id, _) in candidates.iter().take(num_opt_unchoked) {
532 let peer = &mut connections[*id];
533 if peer.optimistically_unchoked {
534 log::debug!("Peer[{}] optmistically unchoked again", peer.peer_id);
535 previously_opt_unchoked.remove(id);
536 } else {
537 peer.optimistically_unchoke(self);
538 }
539 }
540
541 for id in previously_opt_unchoked {
542 let peer = &mut connections[id];
543 peer.choke(self);
544 }
545 }
546}
547
548impl Drop for InitializedState {
549 fn drop(&mut self) {
550 self.piece_buffer_pool.stop_tracking();
551 }
552}
553
554pub struct State {
556 pub(crate) info_hash: [u8; 20],
557 pub(crate) listener_port: Option<u16>,
558 root: PathBuf,
560 torrent_state: Option<InitializedState>,
561 file: OnceCell<Box<TorrentMetadata>>,
562 pub(crate) config: Config,
563}
564
565impl State {
566 pub fn info_hash(&self) -> [u8; 20] {
568 self.info_hash
569 }
570
571 pub fn unstarted(info_hash: [u8; 20], root: PathBuf, config: Config) -> Self {
577 Self {
578 info_hash,
579 root,
580 listener_port: None,
581 torrent_state: None,
582 file: OnceCell::new(),
583 config,
584 }
585 }
586
587 pub fn is_complete(&self) -> bool {
589 self.torrent_state
590 .as_ref()
591 .is_some_and(|state| state.is_complete)
592 }
593
594 pub fn from_metadata_and_root(
606 metadata: TorrentMetadata,
607 root: PathBuf,
608 config: Config,
609 ) -> io::Result<Self> {
610 let mut initialized_state = InitializedState::new(&root, &metadata, config)?;
611 let file_store = &initialized_state.file_store;
612 let completed_pieces: Box<[bool]> = metadata
613 .pieces
614 .as_slice()
615 .par_iter()
616 .enumerate()
617 .map(
618 |(idx, hash)| match file_store.check_piece_hash_sync(idx as i32, hash) {
619 Ok(is_valid) => is_valid,
620 Err(err) => {
621 if err.kind() != io::ErrorKind::NotFound {
622 log::warn!("Error checking piece {idx}: {err}");
623 }
624 false
625 }
626 },
627 )
628 .collect();
629 let completed_pieces: BitVec<u8, Msb0> = completed_pieces.into_iter().collect();
630 let completed_pieces: BitBox<u8, Msb0> = completed_pieces.into_boxed_bitslice();
631 log::trace!("Completed pieces: {completed_pieces}");
632 initialized_state
633 .piece_selector
634 .set_completed_bitfield(completed_pieces);
635 initialized_state.is_complete = initialized_state.piece_selector.completed_all();
636
637 Ok(Self {
638 info_hash: metadata
639 .info_hash_bytes()
640 .try_into()
641 .expect("Invalid info hash"),
642 root,
643 listener_port: None,
644 torrent_state: Some(initialized_state),
645 file: OnceCell::from(Box::new(metadata)),
646 config,
647 })
648 }
649
650 #[cfg(test)]
651 pub fn inprogress(
652 info_hash: [u8; 20],
653 root: PathBuf,
654 metadata: lava_torrent::torrent::v1::Torrent,
655 state: InitializedState,
656 config: Config,
657 ) -> Self {
658 Self {
659 info_hash,
660 root,
661 listener_port: None,
662 torrent_state: Some(state),
663 file: OnceCell::from(Box::new(metadata)),
664 config,
665 }
666 }
667
668 pub fn as_ref(&mut self) -> StateRef<'_> {
669 StateRef {
670 info_hash: self.info_hash,
671 root: &self.root,
672 listener_port: &self.listener_port,
673 torrent: &mut self.torrent_state,
674 full: &self.file,
675 config: &self.config,
676 }
677 }
678}
679
680pub struct StateRef<'state> {
682 info_hash: [u8; 20],
683 root: &'state Path,
684 pub listener_port: &'state Option<u16>,
685 torrent: &'state mut Option<InitializedState>,
686 full: &'state OnceCell<Box<TorrentMetadata>>,
687 pub config: &'state Config,
688}
689
690impl<'e_iter, 'state: 'e_iter> StateRef<'state> {
691 pub fn info_hash(&self) -> &[u8; 20] {
692 &self.info_hash
693 }
694
695 pub fn state(&'e_iter mut self) -> Option<&'e_iter mut InitializedState> {
696 self.torrent.as_mut()
697 }
698
699 #[allow(clippy::borrowed_box)]
700 pub fn metadata(&'e_iter mut self) -> Option<&'state Box<TorrentMetadata>> {
701 self.full.get()
702 }
703
704 #[inline]
705 pub fn is_initialzied(&self) -> bool {
706 self.full.get().is_some()
707 }
708
709 pub fn init(&'e_iter mut self, metadata: TorrentMetadata) -> io::Result<()> {
710 if self.is_initialzied() {
711 return Err(io::Error::other("State initialized twice"));
712 }
713 *self.torrent = Some(InitializedState::new(self.root, &metadata, *self.config)?);
714 self.full
715 .set(Box::new(metadata))
716 .map_err(|_e| io::Error::other("State initialized twice"))?;
717 Ok(())
718 }
719}