Merge pull request #121 from artrixdotdev/perf/bitvec-to-arcbitvec

perf: Changed all uses of BitVec to Arc<BitVec>

Author: kurealnum

crates/libtortillas/src/peer/actor.rs

@@ -1,4 +1,8 @@
-use std::{collections::HashMap, time::Instant};
+use std::{
+   collections::HashMap,
+   sync::{Arc, atomic::AtomicU8},
+   time::Instant,
+};
 
 use anyhow::Context;
 use bitvec::vec::BitVec;
@@ -156,18 +160,23 @@ impl PeerActor {
    #[instrument(skip(self), fields(addr = %self.stream, id = %self.peer.id.unwrap()))]
    async fn determine_interest(&mut self) {
       let msg = TorrentRequest::Bitfield;
+
+      let their_bitfield = self.peer.pieces.clone();
       let our_bitfield = match self.supervisor.ask(msg).await.unwrap() {
          TorrentResponse::Bitfield(bitfield) => bitfield,
          _ => unreachable!("Unexpected response from supervisor"),
       };
-      let their_bitfield = self.peer.pieces.clone();
+
+      let is_our_bitfield_empty = our_bitfield.is_empty();
 
       // Find pieces the peer has that we don't have
-      let peer_has_we_dont = their_bitfield & !our_bitfield.clone();
+      let their_bitfield = (*their_bitfield).clone();
+      let our_bitfield = (*our_bitfield).clone();
+      let peer_has_we_dont = their_bitfield & !our_bitfield;
 
       let has_interesting_pieces = peer_has_we_dont.any();
 
-      if our_bitfield.is_empty() || has_interesting_pieces {
+      if is_our_bitfield_empty || has_interesting_pieces {
          self
             .stream
             .send(PeerMessages::Interested)
@@ -318,7 +327,7 @@ impl Message<PeerMessages> for PeerActor {
          }
          PeerMessages::Have(piece_index) => {
             trace!(piece_index, "Peer has a piece");
-            self.peer.pieces.set(piece_index as usize, true);
+            self.peer.pieces.set_aliased(piece_index as usize, true);
          }
          PeerMessages::Request(index, offset, length) => {
             trace!(
@@ -378,7 +387,7 @@ impl Message<PeerMessages> for PeerActor {
 #[allow(dead_code)]
 pub(crate) enum PeerTell {
    NeedPiece(usize, usize, usize),
-   HaveInfoDict(BitVec<u8>),
+   HaveInfoDict(Arc<BitVec<AtomicU8>>),
    // TODO: Add Have(usize) for notifying the peer that we have a new piece
 }
 

crates/libtortillas/src/peer/mod.rs

@@ -8,6 +8,7 @@ use std::{
    fmt::{self, Debug, Display},
    hash::{Hash as InternalHash, Hasher},
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+   sync::{Arc, atomic::AtomicU8},
 };
 
 pub(crate) use actor::*;
@@ -32,7 +33,7 @@ pub struct Peer {
    pub ip: IpAddr,
    pub port: u16,
    pub state: PeerState,
-   pub pieces: BitVec<u8>,
+   pub pieces: Arc<BitVec<AtomicU8>>,
    /// The reserved bytes that the peer sent us in their handshake. This
    /// indicates what extensions the peer supports.
    pub reserved: [u8; 8],
@@ -80,7 +81,7 @@ impl Peer {
          ip,
          port,
          state: PeerState::new(),
-         pieces: BitVec::EMPTY,
+         pieces: Arc::new(BitVec::EMPTY),
          reserved: [0u8; 8],
          peer_supports: PeerSupports::new(),
          id: None,

crates/libtortillas/src/protocol/messages.rs

@@ -1,7 +1,10 @@
 use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
-   sync::Arc,
+   sync::{
+      Arc,
+      atomic::{AtomicU8, Ordering},
+   },
 };
 
 use anyhow::{Error, Result, bail};
@@ -48,7 +51,7 @@ pub enum PeerMessages {
    /// Each boolean in the vector corresponds to a piece in the torrent. A
    /// `true` value indicates that the peer has the piece, while `false`
    /// indicates that it does not. This is typically sent after the handshake.
-   Bitfield(BitVec<u8>) = 5u8,
+   Bitfield(Arc<BitVec<AtomicU8>>) = 5u8,
    /// A request for a specific block of data from a piece.
    /// Length is usually a power of 2.
    ///
@@ -130,7 +133,9 @@ impl PeerMessages {
          PeerMessages::Interested => create_message_with_id(2, &[]),
          PeerMessages::NotInterested => create_message_with_id(3, &[]),
          PeerMessages::Have(index) => create_message_with_id(4, &index.to_be_bytes()),
-         PeerMessages::Bitfield(bits) => create_message_with_id(5, bits.as_raw_slice()),
+         // This code is wildly confusing, but all it does is maps the array of AtomicU8s to an
+         // vector of u8s
+         PeerMessages::Bitfield(bits) => create_message_with_id(5, &bits.as_raw_slice().iter().map(|byte| byte.load(Ordering::Acquire)).collect::<Vec<u8>>()),
          PeerMessages::Request(index, begin, length) // Code is identical
          | PeerMessages::Cancel(index, begin, length) => {
             let id = match self {
@@ -253,7 +258,10 @@ impl PeerMessages {
          }
          5 => {
             trace!(bitfield_len = payload.len(), "Received Bitfield message");
-            Ok(PeerMessages::Bitfield(BitVec::from_slice(&payload)))
+
+            let bitvec: BitVec<AtomicU8> = payload.into_iter().map(AtomicU8::new).collect();
+
+            Ok(PeerMessages::Bitfield(Arc::new(bitvec)))
          }
          6 => {
             if payload.len() != 12 {

crates/libtortillas/src/torrent/mod.rs

@@ -1,4 +1,8 @@
-use std::{fmt, net::SocketAddr, sync::Arc};
+use std::{
+   fmt,
+   net::SocketAddr,
+   sync::{Arc, atomic::AtomicU8},
+};
 
 use bitvec::vec::BitVec;
 use bytes::Bytes;
@@ -30,7 +34,7 @@ pub struct Torrent {
    peers: Arc<DashMap<PeerId, ActorRef<PeerActor>>>,
    trackers: Arc<DashMap<Tracker, ActorRef<TrackerActor>>>,
 
-   bitfield: BitVec<u8>,
+   bitfield: Arc<BitVec<AtomicU8>>,
    id: PeerId,
    info: Option<Info>,
    metainfo: MetaInfo,
@@ -252,7 +256,7 @@ actor_request_response!(
 
    /// Bitfield of the torrent
    Bitfield
-   Bitfield(BitVec<u8>),
+   Bitfield(Arc<BitVec<AtomicU8>>),
 
    /// Current peers of the torrent
    CurrentPeers
@@ -318,11 +322,11 @@ impl Actor for Torrent {
       if info.is_none() {
          debug!("No info dict found in metainfo, you're probably using a magnet uri");
       }
-      let bitfield: BitVec<u8> = if let Some(info) = &info {
+      let bitfield: Arc<BitVec<AtomicU8>> = if let Some(info) = &info {
          debug!("Using bitfield length {}", info.piece_count());
-         BitVec::with_capacity(info.piece_count())
+         Arc::new(BitVec::with_capacity(info.piece_count()))
       } else {
-         BitVec::EMPTY
+         Arc::new(BitVec::EMPTY)
       };
 
       Ok(Self {
@@ -373,7 +377,7 @@ impl Message<TorrentMessage> for Torrent {
                info!("Received valid info dict, starting torrent process...");
                let info: Info =
                   serde_bencode::from_bytes(&bytes).expect("Failed to parse info dict");
-               self.bitfield.resize(info.piece_count(), false);
+               self.bitfield = Arc::new(BitVec::with_capacity(info.piece_count()));
                self.info = Some(info);
                self
                   .broadcast_to_peers(PeerTell::HaveInfoDict(self.bitfield.clone()))