Bep 6 fast extension 13397060133053203187

crates/librqbit/Cargo.toml

@@ -83,13 +83,14 @@ rand.workspace = true
 tracing-subscriber = { workspace = true, features = ["json", "fmt", "env-filter"], optional = true }
 uuid = { workspace = true, features = ["v4"] }
 futures.workspace = true
+futures-util = "0.3"
 url = { workspace = true, features = ["serde"] }
 hex.workspace = true
 backon.workspace = true
 dashmap.workspace = true
 base64.workspace = true
 serde_with.workspace = true
-tokio-util = { workspace = true, features = ["io"] }
+tokio-util = { workspace = true, features = ["io", "codec"] }
 metrics-exporter-prometheus = { workspace = true, optional = true }
 bytes.workspace = true
 rlimit.workspace = true

crates/librqbit/src/chunk_tracker.rs

@@ -294,6 +294,13 @@ impl ChunkTracker {
     }
 
     // return true if the whole piece is marked downloaded
+    pub fn unmark_chunk_downloaded(&mut self, chunk_info: &ChunkInfo) {
+        let chunk_range = self.lengths.chunk_range(chunk_info.piece_index);
+        if let Some(chunk_range) = self.chunk_status.get_mut(chunk_range) {
+            chunk_range.set(chunk_info.chunk_index as usize, false);
+        }
+    }
+
     pub fn mark_chunk_downloaded(
         &mut self,
         piece: &Piece<ByteBuf<'_>>,

crates/librqbit/src/lib.rs

@@ -68,6 +68,7 @@ mod read_buf;
 mod session;
 mod session_persistence;
 pub mod session_stats;
+pub mod socket_rpc;
 pub mod spawn_utils;
 
 pub mod storage;

crates/librqbit/src/piece_tracker.rs

@@ -15,7 +15,7 @@ use std::{
 };
 
 use buffers::ByteBuf;
-use librqbit_core::lengths::ValidPieceIndex;
+use librqbit_core::lengths::{ValidPieceIndex, ChunkInfo};
 use peer_binary_protocol::Piece;
 
 use crate::{
@@ -239,6 +239,10 @@ impl PieceTracker {
         self.chunks.mark_piece_broken_if_not_have(piece);
     }
 
+    pub fn mark_piece_broken_if_not_have(&mut self, piece: ValidPieceIndex) {
+        self.chunks.mark_piece_broken_if_not_have(piece);
+    }
+
     /// Release all pieces owned by a peer (on peer death).
     ///
     /// Moves all pieces owned by the peer from IN_FLIGHT back to QUEUED.
@@ -281,6 +285,10 @@ impl PieceTracker {
 
     // === PASS-THROUGH METHODS ===
 
+    pub fn unmark_chunk_downloaded(&mut self, chunk_info: &ChunkInfo) {
+        self.chunks.unmark_chunk_downloaded(chunk_info)
+    }
+
     /// Mark a chunk as downloaded. Returns the result indicating if the piece is complete.
     pub fn mark_chunk_downloaded(
         &mut self,

crates/librqbit/src/socket_rpc/commands.rs

@@ -0,0 +1,33 @@
+use librqbit_core::hash_id::Id20;
+
+#[derive(Debug, Clone)]
+pub enum WorkerCommand<'a> {
+    AssignTorrent(&'a [u8]), // Infohash or magnet link bytes
+    Pause(Id20),
+    RequestTelemetry,
+}
+
+impl<'a> WorkerCommand<'a> {
+    pub fn parse(buf: &'a [u8]) -> anyhow::Result<Self> {
+        if buf.is_empty() {
+            anyhow::bail!("empty command");
+        }
+        match buf[0] {
+            0x01 => {
+                Ok(WorkerCommand::AssignTorrent(&buf[1..]))
+            }
+            0x02 => {
+                if buf.len() != 21 {
+                    anyhow::bail!("invalid pause command length");
+                }
+                let mut info_hash = [0u8; 20];
+                info_hash.copy_from_slice(&buf[1..21]);
+                Ok(WorkerCommand::Pause(Id20::new(info_hash)))
+            }
+            0x03 => {
+                Ok(WorkerCommand::RequestTelemetry)
+            }
+            _ => anyhow::bail!("unknown command"),
+        }
+    }
+}

crates/librqbit/src/socket_rpc/mod.rs

@@ -0,0 +1,108 @@
+pub mod commands;
+
+use std::sync::Arc;
+use tokio::net::TcpListener;
+use tokio_util::codec::{LengthDelimitedCodec, FramedRead};
+use tokio::io::AsyncWriteExt;
+use futures_util::stream::StreamExt;
+use std::time::Duration;
+
+use crate::Session;
+use self::commands::WorkerCommand;
+
+pub struct SocketRpcServer {
+    session: Arc<Session>,
+    auth_key: [u8; 32],
+}
+
+impl SocketRpcServer {
+    pub fn new(session: Arc<Session>, auth_key: [u8; 32]) -> Self {
+        Self { session, auth_key }
+    }
+
+    pub async fn listen(self, addr: &str) -> anyhow::Result<()> {
+        let listener = TcpListener::bind(addr).await?;
+        let server = Arc::new(self);
+
+        loop {
+            let (mut stream, _addr) = listener.accept().await?;
+            let server_clone = server.clone();
+
+            tokio::spawn(async move {
+                // Enable TCP Keep-Alive
+                {
+                    let sock_ref = socket2::SockRef::from(&stream);
+                    let tcp_keepalive = socket2::TcpKeepalive::new()
+                        .with_time(Duration::from_secs(60))
+                        .with_interval(Duration::from_secs(15));
+                    let _ = sock_ref.set_tcp_keepalive(&tcp_keepalive);
+                }
+
+                // Auth check
+                let mut auth_buf = [0u8; 32];
+                use tokio::io::AsyncReadExt;
+                if let Ok(n) = stream.read_exact(&mut auth_buf).await {
+                    if n != 32 || auth_buf != server_clone.auth_key {
+                        return; // Drop connection
+                    }
+                } else {
+                    return;
+                }
+
+                let (reader, mut writer) = stream.into_split();
+                let mut framed = FramedRead::new(reader, LengthDelimitedCodec::new());
+
+                while let Some(Ok(frame)) = framed.next().await {
+                    let frame_bytes = frame.freeze();
+                    match WorkerCommand::parse(&frame_bytes) {
+                        Ok(cmd) => {
+                            match cmd {
+                                WorkerCommand::AssignTorrent(data) => {
+                                    if let Ok(s) = std::str::from_utf8(data) {
+                                        let session = server_clone.session.clone();
+                                        let url = s.to_owned();
+                                        tokio::spawn(async move {
+                                            let _ = session.add_torrent(crate::AddTorrent::from_url(url), None).await;
+                                        });
+                                    }
+                                }
+                                WorkerCommand::Pause(info_hash) => {
+                                    if let Some(handle) = server_clone.session.get(crate::api::TorrentIdOrHash::Hash(info_hash)) {
+                                        let _ = handle.pause();
+                                    }
+                                }
+                                WorkerCommand::RequestTelemetry => {
+                                    let down_speed = server_clone.session.stats.down_speed_estimator.bps();
+                                    let up_speed = server_clone.session.stats.up_speed_estimator.bps();
+
+                                    let mut resp = Vec::new();
+                                    resp.extend_from_slice(&down_speed.to_be_bytes());
+                                    resp.extend_from_slice(&up_speed.to_be_bytes());
+
+                                    let chunks_data = server_clone.session.with_torrents(|torrents| {
+                                        let mut chunks_data = Vec::new();
+                                        for (_, mt) in torrents {
+                                            if let Ok(bits) = mt.with_chunk_tracker(|ct| ct.get_have_pieces().as_bytes().to_vec()) {
+                                                chunks_data.extend_from_slice(&mt.info_hash().0);
+                                                chunks_data.extend_from_slice(&(bits.len() as u32).to_be_bytes());
+                                                chunks_data.extend_from_slice(&bits);
+                                            }
+                                        }
+                                        chunks_data
+                                    });
+
+                                    resp.extend_from_slice(&(chunks_data.len() as u32).to_be_bytes());
+                                    resp.extend_from_slice(&chunks_data);
+                                    let _ = writer.write_all(&resp).await;
+                                }
+                            }
+                        }
+                        Err(_) => {
+                            return; // Protocol violation, drop connection
+                        }
+                    }
+                }
+            });
+        }
+    }
+}

crates/librqbit/src/torrent_state/live/mod.rs

@@ -1074,6 +1074,46 @@ impl PeerConnectionHandler for &'_ PeerHandler {
             Message::Cancel(_) => {
                 trace!("received \"cancel\", but we don't process it yet")
             }
+            Message::SuggestPiece(piece) => {
+                self.state.peers.with_live_mut(self.addr, "on_suggest", |l| {
+                    l.suggested_pieces.insert(piece);
+                });
+            }
+            Message::AllowedFast(piece) => {
+                self.state.peers.with_live_mut(self.addr, "on_allowed_fast", |l| {
+                    l.allowed_fast.insert(piece);
+                });
+            }
+            Message::RejectRequest(request) => {
+                let piece_index = match self.state.lengths.validate_piece_index(request.index) {
+                    Some(p) => p,
+                    None => return Ok(()),
+                };
+                let chunk_info = match self.state.lengths.chunk_info_from_received_data(
+                    piece_index,
+                    request.begin,
+                    request.length,
+                ) {
+                    Some(d) => d,
+                    None => return Ok(()),
+                };
+
+                self.state.peers.with_live_mut(self.addr, "on_reject_request", |l| {
+                    l.inflight_requests.remove(&chunk_info);
+                });
+
+                // Immediately release chunk and transition from Pending -> Missing
+                let mut g = self.state.lock_write("on_reject_request_release");
+                if let Some(pieces) = g.get_pieces_mut().ok() {
+                    pieces.unmark_chunk_downloaded(&chunk_info);
+                    // Let the ChunkTracker know the piece has broken chunks
+                    pieces.mark_piece_broken_if_not_have(piece_index);
+                }
+
+                // Signal task to trigger acquire_next_piece
+                self.state.new_pieces_notify.notify_waiters();
+                self.unchoke_notify.notify_waiters();
+            }
             Message::Extended(ExtendedMessage::UtMetadata(UtMetadata::Request(
                 metadata_piece_id,
             ))) => {
@@ -1385,27 +1425,71 @@ impl PeerHandler {
             .state
             .peers
             .with_live_mut(self.addr, "acquire_next_piece", |live| {
-                if self.lock_read("i am choked").i_am_choked {
-                    debug!("we are choked, can't acquire piece");
-                    return Ok(None);
-                }
                 let mut g = self.state.lock_write("acquire_next_piece");
-
                 let bf = &live.bitfield;
-                // Extract references to disjoint fields
+                let allowed_fast = &live.allowed_fast;
+                let suggested_pieces = &live.suggested_pieces;
+
                 let TorrentStateLocked {
                     pieces,
                     file_priorities,
                     ..
                 } = &mut **g;
                 let pieces = pieces.as_mut().ok_or(Error::ChunkTrackerEmpty)?;
+
+                let i_am_choked = self.lock_read("i am choked").i_am_choked;
+
+                if i_am_choked {
+                    // check if we can download an allowed_fast piece
+                    let mut found_allowed = None;
+                    for &fast_piece in allowed_fast {
+                        if let Some(vp) = self.state.lengths.validate_piece_index(fast_piece) {
+                            if !pieces.chunks().is_piece_have(vp) {
+                                found_allowed = Some(vp);
+                                break;
+                            }
+                        }
+                    }
+                    if found_allowed.is_none() {
+                        debug!("we are choked, can't acquire piece");
+                        return Ok(None);
+                    }
+                }
+
+                // If not choked (or if we are, but allowed_fast check above would have exited if none available),
+                // check suggestions locally
+                let mut suggestion_to_acquire = None;
+                for &suggested in suggested_pieces {
+                    if let Some(vp) = self.state.lengths.validate_piece_index(suggested) {
+                        if !pieces.chunks().is_piece_have(vp) {
+                            // If choked, we can only acquire allowed_fast pieces
+                            if !i_am_choked || allowed_fast.contains(&suggested) {
+                                suggestion_to_acquire = Some(vp);
+                                break;
+                            }
+                        }
+                    }
+                }
+
+                let peer_has_piece = |p: ValidPieceIndex| {
+                    if i_am_choked {
+                        allowed_fast.contains(&p.get()) && bf.get(p.get() as usize).map(|v| *v) == Some(true)
+                    } else {
+                        bf.get(p.get() as usize).map(|v| *v) == Some(true)
+                    }
+                };
+
+                let priority_pieces = suggestion_to_acquire.into_iter().chain(
+                    self.state.streams.iter_next_pieces(&self.state.lengths)
+                );
+
                 let result = pieces.acquire_piece(AcquireRequest {
                     peer: self.addr,
                     peer_avg_time: self.counters.average_piece_download_time(),
-                    priority_pieces: self.state.streams.iter_next_pieces(&self.state.lengths),
+                    priority_pieces,
                     file_priorities,
                     file_infos: &self.state.metadata.file_infos,
-                    peer_has_piece: |p| bf.get(p.get() as usize).map(|v| *v) == Some(true),
+                    peer_has_piece,
                     can_steal: |p| {
                         self.state.per_piece_locks[p.get_usize()]
                             .try_write()

crates/librqbit/src/torrent_state/live/peer/mod.rs

@@ -256,6 +256,8 @@ pub(crate) struct LivePeerState {
 
     // When the peer sends us data this is used to track if we asked for it.
     pub inflight_requests: HashSet<InflightRequest>,
+    pub suggested_pieces: HashSet<u32>,
+    pub allowed_fast: HashSet<u32>,
 
     // The main channel to send requests to peer.
     pub tx: PeerTx,
@@ -275,6 +277,8 @@ impl LivePeerState {
             peer_interested: initial_interested,
             bitfield: BF::default(),
             inflight_requests: Default::default(),
+            suggested_pieces: Default::default(),
+            allowed_fast: Default::default(),
             tx,
             connection_kind,
         }