use crate::zeromq::codec::{CodecError, TrySend, *};
use crate::zeromq::endpoint::Endpoint;
use crate::zeromq::error::ZmqResult;
use crate::zeromq::fair_queue::{FairQueue, QueueInner};
use crate::zeromq::message::*;
use crate::zeromq::transport::AcceptStopHandle;
use crate::zeromq::util::PeerIdentity;
use crate::zeromq::{CaptureSocket, SocketOptions};
use crate::zeromq::{
MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketRecv, SocketSend, SocketType,
ZmqError,
};
use async_trait::async_trait;
use dashmap::DashMap;
use futures_channel::mpsc;
use futures_util::StreamExt;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::pin::Pin;
use std::sync::Arc;
pub(crate) struct XPubSubscriber {
pub(crate) subscriptions: Vec<Vec<u8>>,
pub(crate) send_queue: Pin<Box<ZmqFramedWrite>>,
}
pub(crate) struct XPubSocketBackend {
subscribers: DashMap<PeerIdentity, XPubSubscriber>,
fair_queue_inner: Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>,
socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
socket_options: SocketOptions,
}
impl XPubSocketBackend {
fn message_received(&self, peer_id: &PeerIdentity, message: &ZmqMessage) {
let data: Vec<u8> = match message.get(0) {
Some(frame) if message.len() == 1 => frame.to_vec(),
_ => return,
};
if data.is_empty() {
return;
}
match data[0] {
1 => {
if let Some(mut sub) = self.subscribers.get_mut(peer_id) {
sub.subscriptions.push(Vec::from(&data[1..]));
}
}
0 => {
let needle = Vec::from(&data[1..]);
if let Some(mut sub) = self.subscribers.get_mut(peer_id) {
if let Some(idx) = sub.subscriptions.iter().position(|s| s == &needle) {
sub.subscriptions.remove(idx);
}
}
}
_ => {}
}
}
}
impl SocketBackend for XPubSocketBackend {
fn socket_type(&self) -> SocketType {
SocketType::XPUB
}
fn socket_options(&self) -> &SocketOptions {
&self.socket_options
}
fn shutdown(&self) {
self.subscribers.clear();
}
fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
&self.socket_monitor
}
}
#[async_trait]
impl MultiPeerBackend for XPubSocketBackend {
async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
let (recv_queue, send_queue) = io.into_parts();
self.subscribers.insert(
peer_id.clone(),
XPubSubscriber {
subscriptions: vec![],
send_queue: Box::pin(send_queue),
},
);
self.fair_queue_inner
.lock()
.insert(peer_id.clone(), recv_queue);
}
fn peer_disconnected(&self, peer_id: &PeerIdentity) {
self.subscribers.remove(peer_id);
self.fair_queue_inner.lock().remove(peer_id);
}
}
pub struct XPubSocket {
pub(crate) backend: Arc<XPubSocketBackend>,
fair_queue: FairQueue<ZmqFramedRead, PeerIdentity>,
binds: HashMap<Endpoint, AcceptStopHandle>,
}
impl Drop for XPubSocket {
fn drop(&mut self) {
self.backend.shutdown();
}
}
#[async_trait]
impl SocketSend for XPubSocket {
async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
let first = match message.get(0) {
Some(f) => f,
None => return Ok(()),
};
let mut dead_peers = Vec::new();
for mut subscriber in self.backend.subscribers.iter_mut() {
for sub_filter in &subscriber.subscriptions {
if sub_filter.len() <= first.len()
&& sub_filter.as_slice() == &first[0..sub_filter.len()]
{
let res = subscriber
.send_queue
.as_mut()
.try_send(Message::Message(message.clone()));
match res {
Ok(()) => {}
Err(ZmqError::Codec(CodecError::Io(e)))
if e.kind() == ErrorKind::BrokenPipe =>
{
dead_peers.push(subscriber.key().clone());
}
Err(ZmqError::BufferFull(_)) => {}
Err(e) => {
log::warn!("xpub send error: {e}");
dead_peers.push(subscriber.key().clone());
}
}
break;
}
}
}
for peer in dead_peers {
self.backend.peer_disconnected(&peer);
}
Ok(())
}
}
#[async_trait]
impl SocketRecv for XPubSocket {
async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
loop {
match self.fair_queue.next().await {
Some((peer_id, Ok(Message::Message(message)))) => {
self.backend.message_received(&peer_id, &message);
return Ok(message);
}
Some((_peer_id, Ok(_))) => {}
Some((peer_id, Err(_))) => {
self.backend.peer_disconnected(&peer_id);
}
None => return Err(ZmqError::NoMessage),
}
}
}
}
impl CaptureSocket for XPubSocket {}
#[async_trait]
impl Socket for XPubSocket {
fn with_options(options: SocketOptions) -> Self {
let fair_queue = FairQueue::new(true);
Self {
backend: Arc::new(XPubSocketBackend {
subscribers: DashMap::new(),
fair_queue_inner: fair_queue.inner(),
socket_monitor: Mutex::new(None),
socket_options: options,
}),
fair_queue,
binds: HashMap::new(),
}
}
fn backend(&self) -> Arc<dyn MultiPeerBackend> {
self.backend.clone()
}
fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
&mut self.binds
}
fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
let (sender, receiver) = mpsc::channel(1024);
self.backend.socket_monitor.lock().replace(sender);
receiver
}
}