torment/torment-manager/src/tracker_manager.rs

use crate::torrent_manager::TorrentManager;
use crossbeam_channel::{unbounded, Receiver, Sender};
use reqwest::blocking::Client;
use reqwest::header::HeaderMap;
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::RwLock;
use std::thread::JoinHandle;
use torment_core::infohash::v1::U160;
use torment_core::peer_storage::PeerStorageHolder;
use torment_tracker::{
    block_do_announce, Tracker, TrackerAnnounceRequest, TrackerAnnounceResponse,
};
use url::Url;

pub struct TrackerManager {
    id: U160,
    ip: IpAddr,
    port: u16,
    counter: AtomicUsize,
    trackers: HashMap<usize, RwLock<Tracker>>,
    url_index: HashMap<String, usize>,
    peer_storage: PeerStorageHolder,
    tracker_offshore_state: TrackerOffshoreState,
}

struct TrackerOffshoreState {
    threads: Vec<JoinHandle<()>>,
    tracker_request_sender: Sender<(usize, TrackerAnnounceRequest)>,
    tracker_request_receiver: Receiver<(usize, TrackerAnnounceRequest)>,
    tracker_response_sender: Sender<(usize, TrackerAnnounceResponse)>,
    tracker_response_receiver: Receiver<(usize, TrackerAnnounceResponse)>,
}

impl TrackerOffshoreState {
    fn new() -> TrackerOffshoreState {
        let (req_send, req_recv) = unbounded();
        let (resp_send, resp_recv) = unbounded();

        TrackerOffshoreState {
            threads: vec![],
            tracker_request_sender: req_send,
            tracker_request_receiver: req_recv,
            tracker_response_sender: resp_send,
            tracker_response_receiver: resp_recv,
        }
    }
}

pub type TrackerId = usize;

fn tracker_thread(
    input: Receiver<(usize, TrackerAnnounceRequest)>,
    output: Sender<(usize, TrackerAnnounceResponse)>,
) {
    let mut map = HeaderMap::new();
    map.insert(
        reqwest::header::USER_AGENT,
        format!("eater's Torment/{}", env!("CARGO_PKG_VERSION"))
            .parse()
            .unwrap(),
    );
    let client = Client::builder().default_headers(map).build().unwrap();

    for (tracker_id, req) in input {
        output
            .send((tracker_id, block_do_announce(req, &client)))
            .unwrap();
    }
}

impl TrackerManager {
    pub fn new(
        ip: IpAddr,
        port: u16,
        peer_id: U160,
        peer_storage: PeerStorageHolder,
    ) -> TrackerManager {
        TrackerManager {
            id: peer_id,
            ip,
            port,
            counter: Default::default(),
            trackers: Default::default(),
            url_index: Default::default(),
            peer_storage,
            tracker_offshore_state: TrackerOffshoreState::new(),
        }
    }

    pub fn get_tracker_id(&mut self, url: &Url) -> Option<TrackerId> {
        // dht is not a tracker bye
        if url.scheme() == "dht" {
            return None;
        }

        let url_str = url.as_str();
        if let Some(id) = self.url_index.get(url_str) {
            return Some(*id);
        }

        let new_id = self.counter.fetch_add(1, Ordering::AcqRel);
        self.url_index.insert(url_str.to_string(), new_id);
        self.trackers.insert(
            new_id,
            RwLock::new(Tracker::new(url.clone(), self.peer_storage.clone())),
        );

        Some(new_id)
    }

    pub fn announce(&mut self, tracker_id: usize, torrent: &TorrentManager) {
        let id = torrent.info_hash();
        {
            let tracker = &self.trackers[&tracker_id].read().unwrap();
            if !tracker.needs_update(id) {
                return;
            }
        }

        let tracker = &mut self.trackers[&tracker_id].write().unwrap();
        println!(
            "Queueing announce {} on {}",
            torrent.info_hash(),
            tracker.url()
        );

        let req = tracker.create_update_request(
            id,
            self.id,
            self.ip,
            self.port,
            torrent.uploaded(),
            torrent.downloaded(),
            torrent.bytes_left(),
        );

        self.tracker_offshore_state
            .tracker_request_sender
            .send((tracker_id, req))
            .unwrap()
    }

    pub fn house_keeping(&mut self) {
        while self.tracker_offshore_state.threads.len() < 5 {
            let sender = self.tracker_offshore_state.tracker_response_sender.clone();
            let receiver = self.tracker_offshore_state.tracker_request_receiver.clone();

            self.tracker_offshore_state.threads.push(
                std::thread::Builder::new()
                    .name(format!(
                        "tracker/{}",
                        self.tracker_offshore_state.threads.len()
                    ))
                    .spawn(|| tracker_thread(receiver, sender))
                    .unwrap(),
            );
        }

        while let Ok((tracker_id, resp)) = self
            .tracker_offshore_state
            .tracker_response_receiver
            .try_recv()
        {
            self.trackers[&tracker_id].write().unwrap().process(resp);
        }
    }
}