handler: use responses.extend (from drain)

This commit is contained in:
Joakim Frostegård 2020-04-05 03:17:14 +02:00
parent 06533331f2
commit 43286685df

View file

@ -3,7 +3,7 @@ use std::sync::atomic::Ordering;
use std::time::Instant; use std::time::Instant;
use std::vec::Drain; use std::vec::Drain;
use rand::{self, SeedableRng, rngs::SmallRng, thread_rng}; use rand::{Rng, SeedableRng, rngs::SmallRng, thread_rng};
use rand::seq::IteratorRandom; use rand::seq::IteratorRandom;
use bittorrent_udp::types::*; use bittorrent_udp::types::*;
@ -17,9 +17,10 @@ pub fn handle_connect_requests(
requests: Drain<(ConnectRequest, SocketAddr)>, requests: Drain<(ConnectRequest, SocketAddr)>,
){ ){
let now = Time(Instant::now()); let now = Time(Instant::now());
let mut rng = thread_rng();
for (request, src) in requests { responses.extend(requests.map(|(request, src)| {
let connection_id = ConnectionId(rand::random()); let connection_id = ConnectionId(rng.gen());
let key = ConnectionKey { let key = ConnectionKey {
connection_id, connection_id,
@ -28,13 +29,15 @@ pub fn handle_connect_requests(
state.connections.insert(key, now); state.connections.insert(key, now);
responses.push((Response::Connect( let response = Response::Connect(
ConnectResponse { ConnectResponse {
connection_id, connection_id,
transaction_id: request.transaction_id, transaction_id: request.transaction_id,
} }
), src)); );
}
(response, src)
}));
} }
@ -43,14 +46,14 @@ pub fn handle_announce_requests(
responses: &mut Vec<(Response, SocketAddr)>, responses: &mut Vec<(Response, SocketAddr)>,
requests: Drain<(AnnounceRequest, SocketAddr)>, requests: Drain<(AnnounceRequest, SocketAddr)>,
){ ){
for (request, src) in requests { responses.extend(requests.filter_map(|(request, src)| {
let connection_key = ConnectionKey { let connection_key = ConnectionKey {
connection_id: request.connection_id, connection_id: request.connection_id,
socket_addr: src, socket_addr: src,
}; };
if !state.connections.contains_key(&connection_key){ if !state.connections.contains_key(&connection_key){
continue; return None;
} }
let mut torrent_data = state.torrents let mut torrent_data = state.torrents
@ -105,8 +108,8 @@ pub fn handle_announce_requests(
peers: response_peers peers: response_peers
}); });
responses.push((response, src)); Some((response, src))
} }));
} }
@ -117,7 +120,7 @@ pub fn handle_scrape_requests(
){ ){
let empty_stats = create_torrent_scrape_statistics(0, 0); let empty_stats = create_torrent_scrape_statistics(0, 0);
for (request, src) in requests { responses.extend(requests.map(|(request, src)| {
let mut stats: Vec<TorrentScrapeStatistics> = Vec::with_capacity(256); let mut stats: Vec<TorrentScrapeStatistics> = Vec::with_capacity(256);
for info_hash in request.info_hashes.iter() { for info_hash in request.info_hashes.iter() {
@ -136,8 +139,8 @@ pub fn handle_scrape_requests(
torrent_stats: stats, torrent_stats: stats,
}); });
responses.push((response, src)); (response, src)
} }));
} }