Skip to main content

vortex_bittorrent/
torrent.rs

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/// Configuration settings for a given torrent
45#[derive(Debug, Clone, Copy, serde::Deserialize, serde::Serialize)]
46#[serde(default)]
47pub struct Config {
48    /// The max number of total connections that can be open at a time for the torrent
49    pub max_connections: usize,
50    /// The max number of outstanding requests this vortex should report to other peers via the
51    /// `reqq` extension. This currently does not impact the behavoir of vortex but it should impact
52    /// the amount of load connected peers send to us. Vortex does respect connected peers reported
53    /// `reqq` value.
54    pub max_reported_outstanding_requests: u64,
55    /// The maximal amount of allowed unchoked peers at any given time
56    pub max_unchoked: u32,
57    /// Controls how frequently "which peers should be unchoked" is calculated.
58    /// After this number of ticks, vortex will look over all peers and redistribute
59    /// which ones are unchoked and not.
60    pub num_ticks_before_unchoke_recalc: u32,
61    /// Controls the longest possible interval "which peers should be optimistically unchoked" is calculated.
62    /// After this number of ticks, vortex will look over all peers and redistribute
63    /// which ones are optimistically unchoked and not. Note that this may happen more frequently
64    /// than this number suggests due to the normal unchoke distribution "promoting" optimistically
65    /// unchoked peers to normally unchoked peers. In that case this recaluclation will happen
66    /// immedieately afterwards.
67    pub num_ticks_before_optimistic_unchoke_recalc: u32,
68    /// This determines the minimal target of pieces we want to upload to a peer before
69    /// moving on to another peer when the "round-robin" unchoking strategy is used. The
70    /// "round-robin" strategy is currently only used when seeding.
71    pub seeding_piece_quota: u32,
72    /// Controls the size of the io_uring completion queue
73    pub cq_size: u32,
74    /// Controls the size of the io_uring submission queue
75    pub sq_size: u32,
76    /// The event loop will wait for at least these amont of completion events
77    /// before it starts processing them. If the target isn't reached it will wait for
78    /// at most [ `CQE_WAIT_TIME_NS` ] nanoseconds before processing the ones currently in the completion queue.
79    pub completion_event_want: usize,
80    /// The size of the Read buffers used for network operations. Defaults to SUBPIECE_SIZE * 2
81    pub network_read_buffer_size: usize,
82    /// The size of the Write buffers used for network operations. Defaults to SUBPIECE_SIZE + 4Kb
83    pub network_write_buffer_size: usize,
84    /// The size of the Write network buffer pools
85    pub write_buffer_pool_size: usize,
86    /// The size of the Reaad network buffer pools
87    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
110/// A single torrent containing the current torrent state
111/// and our peer id. This object is responsible for starting
112/// the I/O event loop.
113pub struct Torrent {
114    our_id: PeerId,
115    state: State,
116}
117
118impl Torrent {
119    /// Create a new torrent from the given state. Use the given peer id
120    /// when communicating with other peers.
121    pub fn new(our_id: PeerId, state: State) -> Self {
122        Self { our_id, state }
123    }
124
125    /// Returns if the torrent is completed or not.
126    /// This is mostly useful to check status BEFORE
127    /// starting the torrent. Use [ `TorrentEvent::TorrentComplete` ]
128    /// to detect if a torrent have been completed download if you start
129    /// an incomplete torrent.
130    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    /// Start the I/O event loop and the torrent itself. This will start
138    /// a new io_uring instance and new buffer pools
139    /// will be allocated and registered to the io uring instance.
140    /// This is expected to run in a separate thread and interaction with
141    /// the event loop happens via the mpsc command queue. Torrent events will be sent to the
142    /// consumer of the `event_tx` spsc channel.
143    pub fn start(
144        &mut self,
145        event_tx: Producer<TorrentEvent>,
146        command_rc: Receiver<Command>,
147        listener: TcpListener,
148    ) -> Result<(), Error> {
149        // check ulimit
150        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/// Commands that can be sent to the torrent event loop through the command channel
165#[derive(Debug)]
166pub enum Command {
167    /// Connect to peers at the given address.
168    /// Already connected peers will be filtered out
169    ConnectToPeers(Vec<SocketAddrV4>),
170    /// Stop the event loop gracefully, disconnecting peers and freeing resources
171    Stop,
172    /// Pause the torrent gracefully, disconnecting any peers but keeping resources in memory.
173    Pause,
174    /// Resumes a paused torrent, only relevant when the torrent is in the "Paused" state.
175    /// Use the `TorrentEvent::Paused` event to determine when the torrent successfully entered
176    /// the paused state.
177    Resume,
178}
179
180/// Metrics for a given peer
181#[derive(Debug)]
182pub struct PeerMetrics {
183    /// The numer of bytes per second we are currently downloading from the peer
184    pub download_throughput: u64,
185    /// The number of bytes per second we are currently uploading to the peer
186    pub upload_throughput: u64,
187    /// Progress for downloading metadata if supported
188    pub metadata_progress: Option<MetadataProgress>,
189}
190
191/// Events from the torrent
192#[derive(Debug)]
193pub enum TorrentEvent {
194    /// The torrent finished downloading. Note you will NOT receive
195    /// this event if the torrent already is completed when you are
196    /// starting the torrent. Use [ `Torrent::is_complete` ] to check if it's already completed
197    /// before starting the torrent.
198    TorrentComplete,
199    /// The metadata has finished downloading from the connected peers.
200    /// Note you will NOT receive this event if the metadata already is
201    /// completed when starting the torrent.
202    MetadataComplete(Box<TorrentMetadata>),
203    /// The listener for incoming connection has finished set up
204    /// on the provided port and the event loop has started.
205    /// This will be emitted on first startup and every time
206    /// the torrent is resumed after a previous pause.
207    Running { port: u16 },
208    /// The torrent was successfully paused.
209    Paused,
210    /// Over all metrics for the torrent, sent every tick.
211    TorrentMetrics {
212        /// How many pieces have currently been completed
213        pieces_completed: usize,
214        /// How many pieces have been allocated by peers for
215        /// download
216        pieces_allocated: usize,
217        /// Peer metrics for all currently connected peers
218        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            // We are no longer interestead in any of the
287            // peers
288            for (_, peer) in connections.iter_mut() {
289                peer.not_interested();
290                // If the peer is upload only and
291                // we are upload only there is no reason
292                // to stay connected
293                if peer.is_upload_only {
294                    peer.pending_disconnect = Some(DisconnectReason::RedundantConnection);
295                }
296                // Notify all extensions that the torrent completed
297                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                // We are no longer interestead in this peer
310                peer.not_interested();
311                // if it's upload only we can close the
312                // connection since it will never download from
313                // us
314                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    // TODO: Put this in the event loop directly instead when that is easier to test
324    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                // Only need to mark this as not downloaded when it fails
340                // since otherwise it will be marked as completed and this is moot
341                self.piece_selector
342                    .mark_not_downloaded(completed_piece.index);
343                self.piece_buffer_pool.return_buffer(completed_piece.buffer);
344                // deallocate piece peer peer
345                // TODO: disconnect
346                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    // Allocates a piece and increments the piece ref count
354    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    // Deallocate a piece
371    pub fn deallocate_piece(&mut self, index: i32, conn_id: ConnectionId) {
372        log::debug!("Deallocating piece: conn_id: {conn_id:?}, index: {index}");
373        // Mark the piece as interesting again so it can be picked again
374        // if necessary
375        self.piece_selector
376            .update_peer_piece_intrest(conn_id, index as usize);
377        // The piece might have been mid hashing when a timeout is received
378        // (two separate peer) which causes to be completed whilst another peer
379        // tried to download it. It's fine
380        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    /// Determine the most suited peers to be unchoked based on the selected unchoking strategy.
395    /// Some unchoke slots are left for optimistic unchokes. The strategy currently is only
396    /// affected if the torrent is completed or not.
397    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                        // reset so another peer can be optimistically unchoked
417                        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            // Sort based of downloaded_in_last_round
437            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            // The torrent is completed. Do round robin sorting like in libtorrent
446            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                    // no longer optimistic, the peer is "promoted"
489                    // to a normal unchoke slot
490                    peer.optimistically_unchoked = false;
491                    // reset so another peer can be optimistically unchoked
492                    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    // Give some lucky winners unchokes to test if they will have better throughput than the
505    // currently unchoked peers
506    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        // Sort in the order of peers that have waited the longest
530        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
554/// Current state of the torrent
555pub struct State {
556    pub(crate) info_hash: [u8; 20],
557    pub(crate) listener_port: Option<u16>,
558    // TODO: Consider checking this is accessible at construction
559    root: PathBuf,
560    torrent_state: Option<InitializedState>,
561    file: OnceCell<Box<TorrentMetadata>>,
562    pub(crate) config: Config,
563}
564
565impl State {
566    /// The torrent info hash
567    pub fn info_hash(&self) -> [u8; 20] {
568        self.info_hash
569    }
570
571    /// Use this constructor if the torrent is unstarted.
572    /// `info_hash` is the info hash of the torrent that should be downloaded.
573    /// `root` is the directory where the torrent will be downloaded into.
574    /// Vortex will create this directory if it doesn't already exist.
575    /// `config` is the vortex config that should be used
576    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    /// Returns if the torrent has been completed
588    pub fn is_complete(&self) -> bool {
589        self.torrent_state
590            .as_ref()
591            .is_some_and(|state| state.is_complete)
592    }
593
594    /// Use this constructor if you have access to the torrent metadata
595    /// and/or if the torrent has already started.
596    ///
597    /// `metadata` is the metadata associated with the torrent.
598    /// `root` is the directory where potentially already started torrent files
599    /// are expected to be found. If the folder doesn't exist vortex will create it.
600    ///
601    /// `config` is the vortex config that should be used.
602    ///
603    /// NOTE: This will go through all files in `root` and hash their pieces (in parallel) to determine torrent progress
604    /// which may be slow on large torrents.
605    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
680// is this even needed?
681pub 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}