use std::sync::Arc;
use dashmap::{DashMap, DashSet};
use crate::ids::{ConnectionId, SessionId, StreamId};
#[derive(Default)]
pub struct SessionSubscriptions {
inner: DashMap<(ConnectionId, StreamId), Arc<DashSet<ConnectionId>>>,
}
impl SessionSubscriptions {
pub fn new() -> Self {
Self::default()
}
pub fn add(&self, publisher: ConnectionId, strm_id: StreamId, subscriber: ConnectionId) {
let entry = self
.inner
.entry((publisher, strm_id))
.or_insert_with(|| Arc::new(DashSet::new()))
.clone();
entry.insert(subscriber);
}
pub fn remove(&self, publisher: &ConnectionId, strm_id: &StreamId, subscriber: &ConnectionId) -> bool {
let key = (publisher.clone(), strm_id.clone());
let removed = if let Some(entry) = self.inner.get(&key) {
entry.remove(subscriber).is_some()
} else {
false
};
if removed {
if let Some(entry) = self.inner.get(&key) {
if entry.is_empty() {
drop(entry);
self.inner.remove(&key);
}
}
}
removed
}
pub fn subscribers_for(&self, publisher: &ConnectionId, strm_id: &StreamId) -> Vec<ConnectionId> {
match self.inner.get(&(publisher.clone(), strm_id.clone())) {
Some(set) => set.iter().map(|e| e.clone()).collect(),
None => Vec::new(),
}
}
pub fn drop_connection(&self, connid: &ConnectionId) {
let keys: Vec<(ConnectionId, StreamId)> = self
.inner
.iter()
.map(|e| e.key().clone())
.collect();
for key in keys {
if key.0 == *connid {
self.inner.remove(&key);
continue;
}
let collapsed = if let Some(set) = self.inner.get(&key) {
set.remove(connid);
set.is_empty()
} else {
false
};
if collapsed {
self.inner.remove(&key);
}
}
}
pub fn rows(&self) -> Vec<(ConnectionId, StreamId, Vec<ConnectionId>)> {
self.inner
.iter()
.map(|e| {
let (pub_id, strm) = e.key().clone();
let subs = e.value().iter().map(|s| s.clone()).collect();
(pub_id, strm, subs)
})
.collect()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
#[derive(Default)]
pub struct SubscriptionRegistry {
sessions: DashMap<SessionId, Arc<SessionSubscriptions>>,
}
impl SubscriptionRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn for_session(&self, sid: &SessionId) -> Arc<SessionSubscriptions> {
self.sessions
.entry(sid.clone())
.or_insert_with(|| Arc::new(SessionSubscriptions::new()))
.clone()
}
pub fn drop_session(&self, sid: &SessionId) {
self.sessions.remove(sid);
}
pub fn drop_connection(&self, connid: &ConnectionId) {
let sids: Vec<SessionId> = self.sessions.iter().map(|e| e.key().clone()).collect();
for sid in sids {
if let Some(table) = self.sessions.get(&sid) {
let table = Arc::clone(table.value());
drop(self.sessions.get(&sid));
table.drop_connection(connid);
if table.is_empty() {
self.sessions.remove(&sid);
}
}
}
}
}
#[derive(Clone, Debug)]
pub struct PublisherEntry {
pub connection: ConnectionId,
pub participant: String,
pub kind: String,
pub codec: Option<crate::capability::CodecInfo>,
}
#[derive(Default)]
pub struct PublisherRegistry {
inner: DashMap<(SessionId, String), PublisherEntry>,
by_participant: DashMap<(SessionId, String), Vec<String>>,
}
impl PublisherRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&self, sid: SessionId, strm_id: String, entry: PublisherEntry) {
let participant_key = (sid.clone(), entry.participant.clone());
self.inner
.insert((sid.clone(), strm_id.clone()), entry);
self.by_participant
.entry(participant_key)
.and_modify(|v| {
if !v.iter().any(|s| s == &strm_id) {
v.push(strm_id.clone());
}
})
.or_insert_with(|| vec![strm_id]);
}
pub fn publisher(&self, sid: &SessionId, strm_id: &str) -> Option<ConnectionId> {
self.inner
.get(&(sid.clone(), strm_id.to_string()))
.map(|e| e.value().connection.clone())
}
pub fn entry(&self, sid: &SessionId, strm_id: &str) -> Option<PublisherEntry> {
self.inner
.get(&(sid.clone(), strm_id.to_string()))
.map(|e| e.value().clone())
}
pub fn streams_for_participant(&self, sid: &SessionId, participant: &str) -> Vec<String> {
self.by_participant
.get(&(sid.clone(), participant.to_string()))
.map(|e| e.value().clone())
.unwrap_or_default()
}
pub fn drop_publisher(&self, connid: &ConnectionId) {
let to_remove: Vec<(SessionId, String, String)> = self
.inner
.iter()
.filter(|e| &e.value().connection == connid)
.map(|e| {
let (sid, strm) = e.key();
(sid.clone(), strm.clone(), e.value().participant.clone())
})
.collect();
for (sid, strm, participant) in to_remove {
self.inner.remove(&(sid.clone(), strm.clone()));
if let Some(mut entry) = self.by_participant.get_mut(&(sid.clone(), participant.clone())) {
entry.retain(|s| s != &strm);
}
let key = (sid, participant);
let is_empty = self
.by_participant
.get(&key)
.map(|e| e.value().is_empty())
.unwrap_or(false);
if is_empty {
self.by_participant.remove(&key);
}
}
}
pub fn drop_session(&self, sid: &SessionId) {
self.inner.retain(|(s, _), _| s != sid);
self.by_participant.retain(|(s, _), _| s != sid);
}
}