use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::broadcast;
use crate::security::{TenantContext, tenant_namespaced_name};
const MAX_CHANNEL_BYTES: usize = 128;
const MAX_EVENT_BYTES: usize = 128;
const MAX_PAYLOAD_BYTES: usize = 64 * 1024;
const MIN_CHANNEL_SWEEP_THRESHOLD: usize = 64;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RealtimeMessage {
pub channel: String,
pub event: String,
pub payload: String,
}
pub struct Channel {
pub name: String,
sender: broadcast::Sender<RealtimeMessage>,
}
#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
pub enum RealtimeError {
#[error("Broadcast send error: {0}")]
BroadcastError(String),
#[error("invalid realtime channel: {0}")]
InvalidChannel(String),
#[error("invalid realtime event: {0}")]
InvalidEvent(String),
#[error("realtime payload is too large: {actual} bytes exceeds {maximum}")]
PayloadTooLarge {
actual: usize,
maximum: usize,
},
#[error("invalid realtime presence identity: {0}")]
InvalidPresenceIdentity(String),
}
impl Channel {
pub fn new(name: impl Into<String>, capacity: usize) -> Self {
let (sender, _) = broadcast::channel(capacity.max(1));
Self {
name: name.into(),
sender,
}
}
pub fn broadcast(&self, event: &str, payload: &str) -> Result<usize, RealtimeError> {
let msg = RealtimeMessage {
channel: self.name.clone(),
event: event.to_string(),
payload: payload.to_string(),
};
self.sender
.send(msg)
.map_err(|e| RealtimeError::BroadcastError(e.to_string()))
}
pub fn subscribe(&self) -> broadcast::Receiver<RealtimeMessage> {
self.sender.subscribe()
}
}
#[derive(Clone)]
#[non_exhaustive]
pub struct TenantRealtime {
manager: Arc<BroadcastManager>,
tenant_id: String,
}
impl TenantRealtime {
pub fn from_context(manager: Arc<BroadcastManager>, context: &TenantContext) -> Self {
Self {
manager,
tenant_id: context.tenant_id.clone(),
}
}
pub fn tenant_id(&self) -> &str {
&self.tenant_id
}
pub fn namespaced_channel(&self, logical_channel: &str) -> Result<String, RealtimeError> {
validate_room(logical_channel)?;
Ok(tenant_namespaced_name(&self.tenant_id, logical_channel))
}
pub fn subscribe(
&self,
logical_channel: &str,
) -> Result<broadcast::Receiver<RealtimeMessage>, RealtimeError> {
let channel = self.namespaced_channel(logical_channel)?;
Ok(self.manager.get_or_create(&channel).subscribe())
}
pub fn publish(
&self,
logical_channel: &str,
event: &str,
payload: &str,
) -> Result<usize, RealtimeError> {
let channel = self.namespaced_channel(logical_channel)?;
validate_name(event, MAX_EVENT_BYTES, RealtimeError::InvalidEvent)?;
if payload.len() > MAX_PAYLOAD_BYTES {
return Err(RealtimeError::PayloadTooLarge {
actual: payload.len(),
maximum: MAX_PAYLOAD_BYTES,
});
}
self.manager.publish(&channel, event, payload)
}
}
fn validate_name(
value: &str,
maximum: usize,
error: impl FnOnce(String) -> RealtimeError,
) -> Result<(), RealtimeError> {
if value.is_empty()
|| value.len() > maximum
|| !value.bytes().all(|byte| {
byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b':' | b'/')
})
{
return Err(error(value.to_string()));
}
Ok(())
}
fn validate_room(room: &str) -> Result<(), RealtimeError> {
validate_name(room, MAX_CHANNEL_BYTES, RealtimeError::InvalidChannel)?;
if room
.split('/')
.any(|segment| segment.is_empty() || matches!(segment, "." | ".."))
{
return Err(RealtimeError::InvalidChannel(room.to_string()));
}
Ok(())
}
#[derive(Default)]
pub struct BroadcastManager {
channels: DashMap<String, Arc<Channel>>,
sweep_threshold: AtomicUsize,
}
impl BroadcastManager {
pub fn new() -> Self {
Self::default()
}
pub fn get_or_create(&self, channel_name: &str) -> Arc<Channel> {
let mut created = false;
let channel = self
.channels
.entry(channel_name.to_string())
.or_insert_with(|| {
created = true;
Arc::new(Channel::new(channel_name, 100))
})
.value()
.clone();
if created {
self.sweep_idle_channels_if_due();
}
channel
}
pub fn publish(
&self,
channel_name: &str,
event: &str,
payload: &str,
) -> Result<usize, RealtimeError> {
let Some(channel) = self
.channels
.get(channel_name)
.map(|entry| Arc::clone(entry.value()))
else {
return Err(RealtimeError::BroadcastError(
broadcast::error::SendError(()).to_string(),
));
};
let result = channel.broadcast(event, payload);
if result.is_err() {
drop(channel);
self.channels
.remove_if(channel_name, |_, channel| is_idle_channel(channel));
}
result
}
fn sweep_idle_channels_if_due(&self) {
let threshold = self.sweep_threshold.load(Ordering::Relaxed);
if self.channels.len() >= threshold.max(MIN_CHANNEL_SWEEP_THRESHOLD) {
self.remove_idle_channels();
let next = self.channels.len().saturating_mul(2);
self.sweep_threshold.store(next, Ordering::Relaxed);
}
}
fn remove_idle_channels(&self) {
self.channels.retain(|_, channel| !is_idle_channel(channel));
}
}
fn is_idle_channel(channel: &Arc<Channel>) -> bool {
Arc::strong_count(channel) == 1 && channel.sender.receiver_count() == 0
}
#[derive(Default)]
pub struct PresenceTracker {
online_users: DashMap<String, DashMap<String, u64>>,
}
impl PresenceTracker {
pub fn new() -> Self {
Self::default()
}
pub fn user_joined(&self, room: &str, user_id: &str) {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let room_map = self.online_users.entry(room.to_string()).or_default();
room_map.insert(user_id.to_string(), now);
}
pub fn user_left(&self, room: &str, user_id: &str) {
let now_empty = self.online_users.get(room).is_some_and(|room_map| {
room_map.remove(user_id);
room_map.is_empty()
});
if now_empty {
self.online_users
.remove_if(room, |_, room_map| room_map.is_empty());
}
}
pub fn count_online(&self, room: &str) -> usize {
self.online_users.get(room).map(|m| m.len()).unwrap_or(0)
}
}
#[derive(Clone)]
#[non_exhaustive]
pub struct TenantPresence {
tracker: Arc<PresenceTracker>,
tenant_id: String,
}
impl TenantPresence {
pub fn from_context(tracker: Arc<PresenceTracker>, context: &TenantContext) -> Self {
Self {
tracker,
tenant_id: context.tenant_id.clone(),
}
}
pub fn tenant_id(&self) -> &str {
&self.tenant_id
}
pub fn user_joined(&self, room: &str, user_id: &str) -> Result<(), RealtimeError> {
let room = self.namespaced_room(room)?;
validate_name(
user_id,
MAX_CHANNEL_BYTES,
RealtimeError::InvalidPresenceIdentity,
)?;
self.tracker.user_joined(&room, user_id);
Ok(())
}
pub fn user_left(&self, room: &str, user_id: &str) -> Result<(), RealtimeError> {
let room = self.namespaced_room(room)?;
validate_name(
user_id,
MAX_CHANNEL_BYTES,
RealtimeError::InvalidPresenceIdentity,
)?;
self.tracker.user_left(&room, user_id);
Ok(())
}
pub fn count_online(&self, room: &str) -> Result<usize, RealtimeError> {
Ok(self.tracker.count_online(&self.namespaced_room(room)?))
}
fn namespaced_room(&self, room: &str) -> Result<String, RealtimeError> {
validate_room(room)?;
Ok(tenant_namespaced_name(&self.tenant_id, room))
}
}
#[cfg(test)]
#[path = "realtime_retention_tests.rs"]
mod retention_tests;
#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
use super::*;
use crate::security::TenantMembership;
#[tokio::test]
async fn identical_channels_are_isolated_by_authenticated_tenant_context() {
let membership = TenantMembership::try_new(["school-alpha", "school-beta"])
.expect("valid tenant membership");
let alpha_context = membership.select("school-alpha").expect("alpha membership");
let beta_context = membership.select("school-beta").expect("beta membership");
let manager = Arc::new(BroadcastManager::new());
let alpha = TenantRealtime::from_context(Arc::clone(&manager), &alpha_context);
let beta = TenantRealtime::from_context(manager, &beta_context);
let mut alpha_receiver = alpha.subscribe("course/1").expect("alpha subscription");
let mut beta_receiver = beta.subscribe("course/1").expect("beta subscription");
alpha
.publish("course/1", "lesson.completed", r#"{"lesson_id":7}"#)
.expect("alpha publish");
let alpha_message = alpha_receiver.recv().await.expect("alpha message");
assert_eq!(alpha_message.channel, "tenants:school-alpha:course/1");
assert_eq!(alpha_message.event, "lesson.completed");
assert!(matches!(
beta_receiver.try_recv(),
Err(broadcast::error::TryRecvError::Empty)
));
assert!(matches!(
beta.publish("../school-alpha/course/1", "lesson.completed", "{}"),
Err(RealtimeError::InvalidChannel(_))
));
assert!(matches!(
beta.publish("course/1", "lesson completed", "{}"),
Err(RealtimeError::InvalidEvent(_))
));
assert!(matches!(
beta.publish("course/1", "lesson.completed", &"x".repeat(65_537)),
Err(RealtimeError::PayloadTooLarge { .. })
));
let presence = Arc::new(PresenceTracker::new());
let alpha_presence = TenantPresence::from_context(Arc::clone(&presence), &alpha_context);
let beta_presence = TenantPresence::from_context(presence, &beta_context);
alpha_presence
.user_joined("course/1", "learner-7")
.expect("alpha presence");
assert_eq!(
alpha_presence
.count_online("course/1")
.expect("alpha presence count"),
1
);
assert_eq!(
beta_presence
.count_online("course/1")
.expect("beta presence count"),
0
);
assert!(matches!(
beta_presence.user_joined("course/1", "learner 7"),
Err(RealtimeError::InvalidPresenceIdentity(_))
));
}
#[tokio::test]
async fn colon_in_tenant_or_channel_cannot_alias_another_namespace() {
let membership =
TenantMembership::try_new(["a", "a:b", "acme"]).expect("valid tenant membership");
let a_context = membership.select("a").expect("a membership");
let a_b_context = membership.select("a:b").expect("a:b membership");
let acme_context = membership.select("acme").expect("acme membership");
let manager = Arc::new(BroadcastManager::new());
let a = TenantRealtime::from_context(Arc::clone(&manager), &a_context);
let a_b = TenantRealtime::from_context(Arc::clone(&manager), &a_b_context);
let acme = TenantRealtime::from_context(manager, &acme_context);
assert_eq!(a.namespaced_channel("b:c").expect("a"), "tenants:a:b:c");
assert_eq!(a_b.namespaced_channel("c").expect("a:b"), "tenants:a%3Ab:c");
assert_eq!(
acme.namespaced_channel("course:1").expect("acme"),
"tenants:acme:course:1"
);
let mut a_receiver = a.subscribe("b:c").expect("a subscription");
let mut a_b_receiver = a_b.subscribe("c").expect("a:b subscription");
assert_eq!(
a_b.publish("c", "lesson.completed", "{}")
.expect("a:b publish"),
1
);
assert_eq!(
a_b_receiver.recv().await.expect("a:b message").channel,
"tenants:a%3Ab:c"
);
assert!(matches!(
a_receiver.try_recv(),
Err(broadcast::error::TryRecvError::Empty)
));
let presence = Arc::new(PresenceTracker::new());
let a_presence = TenantPresence::from_context(Arc::clone(&presence), &a_context);
let a_b_presence = TenantPresence::from_context(presence, &a_b_context);
assert_eq!(
a_b_presence.namespaced_room("c").expect("a:b room"),
"tenants:a%3Ab:c"
);
a_presence
.user_joined("b:c", "learner-7")
.expect("a presence");
assert_eq!(a_presence.count_online("b:c").expect("a count"), 1);
assert_eq!(a_b_presence.count_online("c").expect("a:b count"), 0);
}
#[test]
fn zero_capacity_channel_is_panic_free() {
let channel = Channel::new("bounded", 0);
let _receiver = channel.subscribe();
assert_eq!(channel.name, "bounded");
}
}