-
Notifications
You must be signed in to change notification settings - Fork 222
/
fetch_stage.rs
48 lines (41 loc) · 1.32 KB
/
fetch_stage.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
//! The `fetch_stage` batches input from a UDP socket and sends it to a channel.
use service::Service;
use std::net::UdpSocket;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::channel;
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use streamer::{self, PacketReceiver};
pub struct FetchStage {
exit: Arc<AtomicBool>,
thread_hdls: Vec<JoinHandle<()>>,
}
impl FetchStage {
pub fn new(sockets: Vec<UdpSocket>, exit: Arc<AtomicBool>) -> (Self, PacketReceiver) {
let tx_sockets = sockets.into_iter().map(Arc::new).collect();
Self::new_multi_socket(tx_sockets, exit)
}
pub fn new_multi_socket(
sockets: Vec<Arc<UdpSocket>>,
exit: Arc<AtomicBool>,
) -> (Self, PacketReceiver) {
let (sender, receiver) = channel();
let thread_hdls: Vec<_> = sockets
.into_iter()
.map(|socket| streamer::receiver(socket, exit.clone(), sender.clone(), "fetch-stage"))
.collect();
(FetchStage { exit, thread_hdls }, receiver)
}
pub fn close(&self) {
self.exit.store(true, Ordering::Relaxed);
}
}
impl Service for FetchStage {
type JoinReturnType = ();
fn join(self) -> thread::Result<()> {
for thread_hdl in self.thread_hdls {
thread_hdl.join()?;
}
Ok(())
}
}