feat: Implement BEP 3-compatible choking algorithm

crates/libtortillas/src/ARCHITECTURE.md

@@ -9,3 +9,11 @@ The library is organized around a small actor hierarchy:
 
 Module facades should export stable public types while keeping actor internals private to the crate.
 Domain types such as torrent state, storage strategy, exported snapshots, tracker model types, and tracker stats live outside actor files so actors can focus on orchestration.
+
+## Choking
+
+`TorrentActor` owns the BEP 3 choking scheduler for its swarm. Active torrents run a rechoke round every 10 seconds, collect peer-local transfer stats from `PeerActor`, and keep at most four interested peers unchoked.
+
+While downloading, regular upload slots are assigned to interested peers with the highest recent download rate. While seeding, regular slots are assigned by recent upload rate. When more interested peers exist than upload slots, one slot is reserved for an optimistic unchoke and rotates every third rechoke round.
+
+`PeerActor` remains responsible for wire-level enforcement: it sends `Choke` and `Unchoke` messages when the torrent scheduler changes state, ignores piece requests from choked peers, and records uploaded bytes when serving piece data.

crates/libtortillas/src/peer/actor.rs

@@ -22,7 +22,7 @@ use tracing::{Span, debug, info, instrument, trace, warn};
 use crate::{
    errors::PeerActorError,
    hashes::InfoHash,
-   peer::Peer,
+   peer::{Peer, PeerId},
    protocol::{stream::PeerRecv, *},
    torrent::{self, TorrentActor},
 };
@@ -32,6 +32,34 @@ const MAX_PENDING_MESSAGES: usize = 8;
 const PEER_KEEPALIVE_TIMEOUT: u64 = 120;
 const PEER_DISCONNECT_TIMEOUT: u64 = 240;
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(crate) struct PeerStats {
+   pub(crate) id: PeerId,
+   pub(crate) interested: bool,
+   pub(crate) choked: bool,
+   pub(crate) download_rate: usize,
+   pub(crate) upload_rate: usize,
+   pub(crate) bytes_downloaded: usize,
+   pub(crate) bytes_uploaded: usize,
+}
+
+#[derive(Clone, Copy, Debug)]
+struct RateSample {
+   at: Instant,
+   bytes_downloaded: usize,
+   bytes_uploaded: usize,
+}
+
+impl RateSample {
+   fn new(peer: &Peer) -> Self {
+      Self {
+         at: Instant::now(),
+         bytes_downloaded: peer.bytes_downloaded(),
+         bytes_uploaded: peer.bytes_uploaded(),
+      }
+   }
+}
+
 /// The actor that handles all communications with a given peer.
 pub(crate) struct PeerActor {
    /// The peers state and statistics
@@ -43,6 +71,7 @@ pub(crate) struct PeerActor {
 
    pending_block_requests: HashSet<(usize, usize, usize)>,
    pending_message_requests: VecDeque<PeerMessages>,
+   last_rate_sample: RateSample,
 }
 
 impl PeerActor {
@@ -306,6 +335,41 @@ impl PeerActor {
 
       self.stream.send(msg).await
    }
+
+   fn snapshot_stats(&mut self) -> Option<PeerStats> {
+      let id = self.peer.id?;
+      let now = Instant::now();
+      let bytes_downloaded = self.peer.bytes_downloaded();
+      let bytes_uploaded = self.peer.bytes_uploaded();
+      let elapsed_secs = now
+         .duration_since(self.last_rate_sample.at)
+         .as_secs()
+         .max(1) as usize;
+
+      let download_rate = bytes_downloaded.saturating_sub(self.last_rate_sample.bytes_downloaded)
+         / 1024
+         / elapsed_secs;
+      let upload_rate =
+         bytes_uploaded.saturating_sub(self.last_rate_sample.bytes_uploaded) / 1024 / elapsed_secs;
+
+      self.peer.set_download_rate(download_rate);
+      self.peer.set_upload_rate(upload_rate);
+      self.last_rate_sample = RateSample {
+         at: now,
+         bytes_downloaded,
+         bytes_uploaded,
+      };
+
+      Some(PeerStats {
+         id,
+         interested: self.peer.interested(),
+         choked: self.peer.choked(),
+         download_rate,
+         upload_rate,
+         bytes_downloaded,
+         bytes_uploaded,
+      })
+   }
 }
 
 impl Actor for PeerActor {
@@ -340,6 +404,7 @@ impl Actor for PeerActor {
          .map_err(|e| PeerActorError::SupervisorCommunicationFailed(e.to_string()))?;
 
       Ok(Self {
+         last_rate_sample: RateSample::new(&peer),
          peer,
          stream,
          supervisor,
@@ -431,17 +496,15 @@ impl Message<PeerMessages> for PeerActor {
                data_len = data.len(),
                "Received piece data"
             );
-            self.peer.increment_bytes_downloaded(data.len());
             let key = (index as usize, offset as usize, data.len());
             // Only send the piece to the supervisor if they request it or it hasn't been
             // cancelled
-            if self.pending_block_requests.contains(&key) {
-               self.pending_block_requests.remove(&key);
-
+            if self.pending_block_requests.remove(&key) {
                let Some(peer_id) = self.peer.id else {
                   warn!("Received piece from peer without id; ignoring");
                   return;
                };
+               self.peer.increment_bytes_downloaded(data.len());
                let supervisor_msg = torrent::events::IncomingPiece {
                   peer_id,
                   index: index as usize,
@@ -508,6 +571,14 @@ impl Message<PeerMessages> for PeerActor {
             }
          }
          PeerMessages::Request(index, offset, length) => {
+            if self.peer.choked() {
+               trace!(
+                  piece_index = index,
+                  offset, length, "Ignoring request from choked peer"
+               );
+               return;
+            }
+
             debug!(
                piece_index = index,
                offset, length, "Peer requested piece data"
@@ -525,11 +596,13 @@ impl Message<PeerMessages> for PeerActor {
 
             match data {
                Some(data) => {
+                  let uploaded_bytes = data.len();
                   self
                      .stream
                      .send(PeerMessages::Piece(index as u32, offset as u32, data))
                      .await
                      .expect("Failed to send piece");
+                  self.peer.increment_bytes_uploaded(uploaded_bytes);
                }
                None => {
                   warn!(
@@ -682,6 +755,33 @@ pub(crate) mod commands {
          }
       }
 
+      #[message(derive(Clone, Debug))]
+      #[instrument(skip(self), fields(peer_addr = %self.stream, peer_id = %self.peer.id.unwrap()))]
+      pub(crate) async fn set_choked(&mut self, choked: bool) {
+         if self.peer.choked() == choked {
+            return;
+         }
+
+         let message = if choked {
+            PeerMessages::Choke
+         } else {
+            PeerMessages::Unchoke
+         };
+
+         if let Err(err) = self.stream.send(message).await {
+            warn!(?err, choked, "Failed to update peer choke state");
+            return;
+         }
+
+         self.peer.set_choked(choked);
+         trace!(choked, "Updated peer choke state");
+      }
+
+      #[message]
+      pub(crate) fn stats(&mut self) -> Option<PeerStats> {
+         self.snapshot_stats()
+      }
+
       #[message(derive(Clone, Debug))]
       pub(crate) async fn have_info_dict(&mut self, bitfield: Arc<BitVec<AtomicU8>>) {
          self

crates/libtortillas/src/torrent/actor.rs

@@ -18,12 +18,15 @@ use kameo::{
    mailbox::{MailboxReceiver, Signal},
    supervision::{RestartPolicy, SupervisionStrategy},
 };
-use kameo_actors::scheduler::Scheduler;
+use kameo_actors::scheduler::{Scheduler, SetTimeout};
 use librqbit_utp::UtpSocketUdp;
-use tokio::sync::oneshot;
+use tokio::{sync::oneshot, task::AbortHandle};
 use tracing::{debug, error, info, instrument, trace, warn};
 
-use super::util;
+use super::{
+   choking::{ChokingScheduler, RECHOKE_INTERVAL},
+   util,
+};
 use crate::{
    errors::TorrentError,
    hashes::InfoHash,
@@ -102,6 +105,9 @@ pub(crate) struct TorrentActor {
    pub state: TorrentState,
    /// Scheduler for managing piece and block requests
    pub(super) piece_scheduler: PieceScheduler,
+   /// Scheduler for BEP 3 upload choking decisions.
+   pub(super) choking_scheduler: ChokingScheduler,
+   pub(super) next_rechoke: Option<AbortHandle>,
 
    pub(super) start_time: Option<Instant>,
    /// The number of peers we need to have before we start downloading, defaults
@@ -235,6 +241,45 @@ impl TorrentActor {
          state = ?self.state,
          "Started torrenting process"
       );
+
+      self.rechoke_peers().await;
+      self.schedule_next_rechoke().await;
+   }
+
+   pub(super) async fn schedule_next_rechoke(&mut self) {
+      if !matches!(
+         self.state,
+         TorrentState::Downloading | TorrentState::Seeding
+      ) {
+         return;
+      }
+
+      if let Some(next_rechoke) = self.next_rechoke.take() {
+         next_rechoke.abort();
+      }
+
+      match self
+         .scheduler
+         .ask(SetTimeout::new(
+            self.actor_ref.downgrade(),
+            RECHOKE_INTERVAL,
+            super::commands::Rechoke,
+         ))
+         .await
+      {
+         Ok(next_rechoke) => self.next_rechoke = Some(next_rechoke),
+         Err(err) => {
+            warn!(?err, "Failed to schedule next rechoke; using local timeout");
+            let actor_ref = self.actor_ref.clone();
+            let fallback = tokio::spawn(async move {
+               tokio::time::sleep(RECHOKE_INTERVAL).await;
+               if let Err(err) = actor_ref.tell(super::commands::Rechoke).await {
+                  warn!(?err, "Failed to run fallback rechoke");
+               }
+            });
+            self.next_rechoke = Some(fallback.abort_handle());
+         }
+      }
    }
 
    pub(super) fn send_ready_hooks(&mut self) {
@@ -482,6 +527,8 @@ impl Actor for TorrentActor {
          piece_store,
          state: TorrentState::default(),
          piece_scheduler: PieceScheduler::new(piece_count),
+         choking_scheduler: ChokingScheduler::default(),
+         next_rechoke: None,
          start_time: None,
          sufficient_peers: sufficient_peers.unwrap_or(6),
          autostart: autostart.unwrap_or(true),
@@ -511,6 +558,9 @@ impl Actor for TorrentActor {
       for tracker in self.trackers.values() {
          tracker.kill();
       }
+      if let Some(next_rechoke) = self.next_rechoke.take() {
+         next_rechoke.abort();
+      }
       self.piece_store.kill();
       self.scheduler.kill();
 
@@ -532,14 +582,20 @@ mod tests {
    use std::{path::PathBuf, time::Duration};
 
    use librqbit_utp::UtpSocket;
-   use tokio::{fs, time::sleep};
+   use tokio::{
+      fs,
+      time::{sleep, timeout},
+   };
    use tracing::trace;
 
    use super::*;
    use crate::{
       metainfo::MetaInfo,
       testing,
-      torrent::{BLOCK_SIZE, Torrent, TorrentExport, commands::HasInfoDict},
+      torrent::{
+         BLOCK_SIZE, Torrent, TorrentExport,
+         commands::{ExportState, HasInfoDict},
+      },
    };
 
    #[tokio::test(flavor = "multi_thread")]
@@ -672,27 +728,38 @@ mod tests {
 
       assert!(torrent.poll_ready().await.is_ok());
 
-      loop {
-         let mut entries = fs::read_dir(&piece_path).await.unwrap();
-         let mut found_piece = false;
-         while let Some(entry) = entries.next_entry().await.unwrap() {
-            let path = entry.path();
-            if let Some(ext) = path.extension()
-               && ext == "piece"
-            {
-               let metadata = entry.metadata().await.unwrap();
-               if metadata.len() == info_dict.piece_length {
-                  found_piece = true;
-                  break;
+      let wrote_piece_block = timeout(Duration::from_secs(60), async {
+         loop {
+            let export = actor.ask(ExportState).await.unwrap();
+            let has_persisted_progress = export.bitfield.count_ones() > 0
+               || export
+                  .block_map
+                  .iter()
+                  .any(|entry| entry.value().count_ones() > 0);
+
+            if has_persisted_progress {
+               let mut entries = fs::read_dir(&piece_path).await.unwrap();
+               while let Some(entry) = entries.next_entry().await.unwrap() {
+                  let path = entry.path();
+                  if let Some(ext) = path.extension()
+                     && ext == "piece"
+                     && entry.metadata().await.unwrap().len() > 0
+                  {
+                     return true;
+                  }
                }
             }
+
+            sleep(Duration::from_millis(200)).await;
          }
-         if found_piece {
-            break;
-         }
+      })
+      .await
+      .unwrap_or(false);
 
-         sleep(Duration::from_millis(200)).await;
-      }
+      assert!(
+         wrote_piece_block,
+         "timed out waiting for piece storage write"
+      );
 
       actor.stop_gracefully().await.unwrap();
       clear_piece_files(&piece_path).await;
@@ -770,6 +837,8 @@ mod tests {
          )),
          state: TorrentState::Inactive,
          piece_scheduler,
+         choking_scheduler: ChokingScheduler::default(),
+         next_rechoke: None,
          start_time: None,
          sufficient_peers: 6,
          autostart: false,