use crate::messaging::dispatch::DispatchData;
use std::{
collections::{HashMap, VecDeque},
net::SocketAddr,
};
pub struct QueueManager {
inner: HashMap<SocketAddr, VecDeque<DispatchData>>,
priority_queue: HashMap<SocketAddr, VecDeque<DispatchData>>,
}
impl QueueManager {
pub fn new() -> Self {
QueueManager {
inner: HashMap::new(),
priority_queue: HashMap::new(),
}
}
pub fn enqueue_data(&mut self, data: DispatchData, dst: SocketAddr) {
self.inner.entry(dst).or_default().push_back(data);
}
pub fn enqueue_priority_data(&mut self, data: DispatchData, dst: SocketAddr) {
self.priority_queue.entry(dst).or_default().push_back(data);
}
pub fn pop_data(&mut self, dst: &SocketAddr) -> Option<DispatchData> {
let mut res = self.priority_queue.get_mut(dst).and_then(|q| q.pop_front());
if self.priority_queue.contains_key(dst) && res.is_none() {
self.priority_queue.remove(dst);
}
if res.is_none() {
res = self.inner.get_mut(dst).and_then(|q| q.pop_front());
if self.inner.contains_key(dst) && res.is_none() {
self.inner.remove(dst);
}
}
res
}
pub fn drop_queue(&mut self, addr: &SocketAddr) {
self.priority_queue.remove(addr);
self.inner.remove(addr);
}
pub fn has_data(&self, dst: &SocketAddr) -> bool {
if self.priority_queue.get(dst).is_some_and(|q| !q.is_empty()) {
return true;
}
self.inner.get(dst).is_some_and(|q| !q.is_empty())
}
}