soaprs-memory 0.4.0

//! In-memory realtime presence, membership, and delivery adapters.

use std::{
    collections::{BTreeMap, BTreeSet, VecDeque},
    future::Future,
    num::NonZeroUsize,
    pin::Pin,
    sync::{Mutex, RwLock},
    task::{Context, Poll, Waker},
};

use soaprs_auth::{Principal, PrincipalId};
use soaprs_core::{BoxFuture, SoapError, SoapResult};
use soaprs_realtime::{
    BackpressurePolicy, ChannelId, ChannelMembership, ConnectionContext, ConnectionId,
    ConnectionPresence, DeliveryOutcome, JoinOutcome, LeaveOutcome, OutboundDelivery,
    RealtimeDelivery,
};

/// Thread-safe active connection directory.
#[derive(Debug)]
pub struct MemoryConnectionPresence<P> {
    connections: RwLock<BTreeMap<ConnectionId, ConnectionContext<P>>>,
}

impl<P> MemoryConnectionPresence<P> {
    /// Creates an empty connection directory.
    pub const fn new() -> Self {
        Self {
            connections: RwLock::new(BTreeMap::new()),
        }
    }

    fn read(
        &self,
    ) -> SoapResult<std::sync::RwLockReadGuard<'_, BTreeMap<ConnectionId, ConnectionContext<P>>>>
    {
        self.connections
            .read()
            .map_err(|_| SoapError::infrastructure("in-memory presence read lock poisoned"))
    }

    fn write(
        &self,
    ) -> SoapResult<std::sync::RwLockWriteGuard<'_, BTreeMap<ConnectionId, ConnectionContext<P>>>>
    {
        self.connections
            .write()
            .map_err(|_| SoapError::infrastructure("in-memory presence write lock poisoned"))
    }
}

impl<P> Default for MemoryConnectionPresence<P> {
    fn default() -> Self {
        Self::new()
    }
}

impl<P> ConnectionPresence<P> for MemoryConnectionPresence<P>
where
    P: Principal + Clone + Send + Sync,
{
    fn register(&self, context: ConnectionContext<P>) -> BoxFuture<'_, SoapResult<()>> {
        Box::pin(async move {
            let mut connections = self.write()?;
            if connections.contains_key(context.connection_id()) {
                return Err(SoapError::conflict(format!(
                    "realtime connection `{}` is already registered",
                    context.connection_id()
                )));
            }
            connections.insert(context.connection_id().clone(), context);
            Ok(())
        })
    }

    fn remove<'a>(
        &'a self,
        connection_id: &'a ConnectionId,
    ) -> BoxFuture<'a, SoapResult<Option<ConnectionContext<P>>>> {
        Box::pin(async move { Ok(self.write()?.remove(connection_id)) })
    }

    fn connection<'a>(
        &'a self,
        connection_id: &'a ConnectionId,
    ) -> BoxFuture<'a, SoapResult<Option<ConnectionContext<P>>>> {
        Box::pin(async move { Ok(self.read()?.get(connection_id).cloned()) })
    }

    fn connections(&self) -> BoxFuture<'_, SoapResult<Vec<ConnectionContext<P>>>> {
        Box::pin(async move { Ok(self.read()?.values().cloned().collect()) })
    }

    fn principal_connections<'a>(
        &'a self,
        principal_id: &'a PrincipalId,
    ) -> BoxFuture<'a, SoapResult<Vec<ConnectionContext<P>>>> {
        Box::pin(async move {
            Ok(self
                .read()?
                .values()
                .filter(|context| {
                    context.authentication().is_some_and(|authentication| {
                        authentication.principal().principal_id() == principal_id
                    })
                })
                .cloned()
                .collect())
        })
    }
}

#[derive(Debug, Default)]
struct MembershipState {
    channels: BTreeMap<ChannelId, BTreeSet<ConnectionId>>,
    connections: BTreeMap<ConnectionId, BTreeSet<ChannelId>>,
}

/// Thread-safe bidirectional channel membership.
#[derive(Debug, Default)]
pub struct MemoryChannelMembership {
    state: RwLock<MembershipState>,
}

impl MemoryChannelMembership {
    /// Creates empty channel membership.
    pub const fn new() -> Self {
        Self {
            state: RwLock::new(MembershipState {
                channels: BTreeMap::new(),
                connections: BTreeMap::new(),
            }),
        }
    }

    fn read(&self) -> SoapResult<std::sync::RwLockReadGuard<'_, MembershipState>> {
        self.state
            .read()
            .map_err(|_| SoapError::infrastructure("in-memory membership read lock poisoned"))
    }

    fn write(&self) -> SoapResult<std::sync::RwLockWriteGuard<'_, MembershipState>> {
        self.state
            .write()
            .map_err(|_| SoapError::infrastructure("in-memory membership write lock poisoned"))
    }
}

impl ChannelMembership for MemoryChannelMembership {
    fn join<'a>(
        &'a self,
        connection_id: &'a ConnectionId,
        channel_id: &'a ChannelId,
    ) -> BoxFuture<'a, SoapResult<JoinOutcome>> {
        Box::pin(async move {
            let mut state = self.write()?;
            let joined = state
                .channels
                .entry(channel_id.clone())
                .or_default()
                .insert(connection_id.clone());
            state
                .connections
                .entry(connection_id.clone())
                .or_default()
                .insert(channel_id.clone());
            Ok(if joined {
                JoinOutcome::Joined
            } else {
                JoinOutcome::AlreadyMember
            })
        })
    }

    fn leave<'a>(
        &'a self,
        connection_id: &'a ConnectionId,
        channel_id: &'a ChannelId,
    ) -> BoxFuture<'a, SoapResult<LeaveOutcome>> {
        Box::pin(async move {
            let mut state = self.write()?;
            let removed = state
                .channels
                .get_mut(channel_id)
                .is_some_and(|members| members.remove(connection_id));
            if state
                .channels
                .get(channel_id)
                .is_some_and(BTreeSet::is_empty)
            {
                state.channels.remove(channel_id);
            }
            if let Some(channels) = state.connections.get_mut(connection_id) {
                channels.remove(channel_id);
                if channels.is_empty() {
                    state.connections.remove(connection_id);
                }
            }
            Ok(if removed {
                LeaveOutcome::Left
            } else {
                LeaveOutcome::NotMember
            })
        })
    }

    fn members<'a>(
        &'a self,
        channel_id: &'a ChannelId,
    ) -> BoxFuture<'a, SoapResult<Vec<ConnectionId>>> {
        Box::pin(async move {
            Ok(self
                .read()?
                .channels
                .get(channel_id)
                .map(|members| members.iter().cloned().collect())
                .unwrap_or_default())
        })
    }

    fn channels<'a>(
        &'a self,
        connection_id: &'a ConnectionId,
    ) -> BoxFuture<'a, SoapResult<Vec<ChannelId>>> {
        Box::pin(async move {
            Ok(self
                .read()?
                .connections
                .get(connection_id)
                .map(|channels| channels.iter().cloned().collect())
                .unwrap_or_default())
        })
    }

    fn remove_connection<'a>(
        &'a self,
        connection_id: &'a ConnectionId,
    ) -> BoxFuture<'a, SoapResult<Vec<ChannelId>>> {
        Box::pin(async move {
            let mut state = self.write()?;
            let channels = state
                .connections
                .remove(connection_id)
                .unwrap_or_default()
                .into_iter()
                .collect::<Vec<_>>();
            for channel_id in &channels {
                if let Some(members) = state.channels.get_mut(channel_id) {
                    members.remove(connection_id);
                    if members.is_empty() {
                        state.channels.remove(channel_id);
                    }
                }
            }
            Ok(channels)
        })
    }
}

struct DeliveryState<M> {
    queue: VecDeque<OutboundDelivery<M>>,
    waiters: Vec<Waker>,
}

/// Runtime-independent bounded or unbounded outbound delivery recorder.
///
/// Bounded instances expose the contract semantics of wait, reject, and drop
/// backpressure. Draining the queue wakes pending waiters.
pub struct MemoryRealtimeDelivery<M> {
    capacity: Option<NonZeroUsize>,
    state: Mutex<DeliveryState<M>>,
}

impl<M> MemoryRealtimeDelivery<M> {
    /// Creates an unbounded recorder that accepts every delivery.
    pub const fn new() -> Self {
        Self {
            capacity: None,
            state: Mutex::new(DeliveryState {
                queue: VecDeque::new(),
                waiters: Vec::new(),
            }),
        }
    }

    /// Creates a recorder with explicit non-zero queue capacity.
    pub const fn bounded(capacity: NonZeroUsize) -> Self {
        Self {
            capacity: Some(capacity),
            state: Mutex::new(DeliveryState {
                queue: VecDeque::new(),
                waiters: Vec::new(),
            }),
        }
    }

    /// Returns the current accepted queue length.
    pub fn len(&self) -> SoapResult<usize> {
        Ok(self.lock()?.queue.len())
    }

    /// Reports whether no accepted deliveries are queued.
    pub fn is_empty(&self) -> SoapResult<bool> {
        self.len().map(|length| length == 0)
    }

    /// Removes every accepted delivery in enqueue order and wakes waiters.
    pub fn drain(&self) -> SoapResult<Vec<OutboundDelivery<M>>> {
        let (deliveries, waiters) = {
            let mut state = self.lock()?;
            let deliveries = state.queue.drain(..).collect();
            let waiters = std::mem::take(&mut state.waiters);
            (deliveries, waiters)
        };
        for waiter in waiters {
            waiter.wake();
        }
        Ok(deliveries)
    }

    fn lock(&self) -> SoapResult<std::sync::MutexGuard<'_, DeliveryState<M>>> {
        self.state
            .lock()
            .map_err(|_| SoapError::infrastructure("in-memory delivery lock poisoned"))
    }

    fn has_capacity(&self, length: usize) -> bool {
        self.capacity.is_none_or(|capacity| length < capacity.get())
    }
}

impl<M> Default for MemoryRealtimeDelivery<M> {
    fn default() -> Self {
        Self::new()
    }
}

impl<M> RealtimeDelivery<M> for MemoryRealtimeDelivery<M>
where
    M: Send,
{
    fn deliver(&self, delivery: OutboundDelivery<M>) -> BoxFuture<'_, SoapResult<DeliveryOutcome>> {
        Box::pin(MemoryDeliveryFuture {
            delivery: Mutex::new(Some(delivery)),
            sink: self,
        })
    }
}

struct MemoryDeliveryFuture<'a, M> {
    delivery: Mutex<Option<OutboundDelivery<M>>>,
    sink: &'a MemoryRealtimeDelivery<M>,
}

impl<M> Future for MemoryDeliveryFuture<'_, M>
where
    M: Send,
{
    type Output = SoapResult<DeliveryOutcome>;

    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        let future = self.as_ref().get_ref();
        let mut state = match self.sink.lock() {
            Ok(state) => state,
            Err(error) => return Poll::Ready(Err(error)),
        };
        if self.sink.has_capacity(state.queue.len()) {
            let mut delivery = match future.delivery.lock() {
                Ok(delivery) => delivery,
                Err(_) => {
                    return Poll::Ready(Err(SoapError::infrastructure(
                        "in-memory delivery future lock poisoned",
                    )));
                }
            };
            let Some(delivery) = delivery.take() else {
                return Poll::Ready(Err(SoapError::infrastructure(
                    "in-memory delivery future was polled after completion",
                )));
            };
            state.queue.push_back(delivery);
            return Poll::Ready(Ok(DeliveryOutcome::Accepted));
        }

        let backpressure = match future.delivery.lock() {
            Ok(delivery) => delivery.as_ref().map(|delivery| delivery.backpressure),
            Err(_) => {
                return Poll::Ready(Err(SoapError::infrastructure(
                    "in-memory delivery future lock poisoned",
                )));
            }
        };
        let Some(backpressure) = backpressure else {
            return Poll::Ready(Err(SoapError::infrastructure(
                "in-memory delivery future lost its message",
            )));
        };
        match backpressure {
            BackpressurePolicy::Wait => {
                if !state
                    .waiters
                    .iter()
                    .any(|waiter| waiter.will_wake(context.waker()))
                {
                    state.waiters.push(context.waker().clone());
                }
                Poll::Pending
            }
            BackpressurePolicy::Reject => Poll::Ready(Err(SoapError::unavailable(
                "realtime outbound delivery queue is full",
            ))),
            BackpressurePolicy::DropNewest => match future.delivery.lock() {
                Ok(mut delivery) => {
                    *delivery = None;
                    Poll::Ready(Ok(DeliveryOutcome::Dropped))
                }
                Err(_) => Poll::Ready(Err(SoapError::infrastructure(
                    "in-memory delivery future lock poisoned",
                ))),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{
        num::NonZeroUsize,
        pin::pin,
        task::{Context, Poll, Waker},
        time::SystemTime,
    };

    use soaprs_auth::{Authentication, Principal, StandardPrincipal};
    use soaprs_contract_tests::{
        block_on, verify_channel_membership_contract, verify_connection_presence_contract,
    };
    use soaprs_core::{MessageEnvelope, MessageMetadata, SoapErrorKind};
    use soaprs_realtime::{
        BackpressurePolicy, ChannelId, ConnectionContext, ConnectionId, DeliveryOutcome,
        OutboundDelivery, RealtimeDelivery,
    };

    use super::{MemoryChannelMembership, MemoryConnectionPresence, MemoryRealtimeDelivery};

    #[test]
    fn presence_and_membership_pass_shared_contracts() {
        let Some(first_id) = ConnectionId::new("connection-b").ok() else {
            panic!("valid connection id");
        };
        let Some(second_id) = ConnectionId::new("connection-a").ok() else {
            panic!("valid connection id");
        };
        let principal = || {
            StandardPrincipal::new("player-1")
                .and_then(|principal| principal.role("player"))
                .and_then(|principal| Authentication::new("session", principal))
        };
        let (Some(first_auth), Some(second_auth)) = (principal().ok(), principal().ok()) else {
            panic!("valid authentication");
        };
        let first = ConnectionContext::new(first_id.clone(), SystemTime::UNIX_EPOCH)
            .authenticated(first_auth);
        let second = ConnectionContext::new(second_id.clone(), SystemTime::UNIX_EPOCH)
            .authenticated(second_auth);
        let principal_id = first
            .authentication()
            .map(|authentication| authentication.principal().principal_id().clone());
        let Some(principal_id) = principal_id else {
            panic!("authenticated fixture");
        };

        let presence = MemoryConnectionPresence::new();
        assert!(
            block_on(verify_connection_presence_contract(
                &presence,
                first,
                second,
                &principal_id,
            ))
            .is_ok()
        );

        let Some(first_channel) = ChannelId::new("channel-b").ok() else {
            panic!("valid channel id");
        };
        let Some(second_channel) = ChannelId::new("channel-a").ok() else {
            panic!("valid channel id");
        };
        let membership = MemoryChannelMembership::new();
        assert!(
            block_on(verify_channel_membership_contract(
                &membership,
                &first_id,
                &second_id,
                &first_channel,
                &second_channel,
            ))
            .is_ok()
        );
    }

    #[test]
    fn bounded_delivery_exposes_wait_reject_and_drop_backpressure() {
        let capacity = NonZeroUsize::new(1).unwrap_or(NonZeroUsize::MIN);
        let delivery = MemoryRealtimeDelivery::bounded(capacity);
        let Some(connection_id) = ConnectionId::new("connection-1").ok() else {
            panic!("valid connection id");
        };
        let message = |id: &'static str| {
            MessageEnvelope::new(
                id.to_owned(),
                MessageMetadata::new(id, SystemTime::UNIX_EPOCH),
            )
        };

        let first = OutboundDelivery::direct(connection_id.clone(), message("first"));
        assert_eq!(
            block_on(delivery.deliver(first)).ok(),
            Some(DeliveryOutcome::Accepted)
        );

        let dropped = OutboundDelivery::direct(connection_id.clone(), message("dropped"))
            .backpressure(BackpressurePolicy::DropNewest);
        assert_eq!(
            block_on(delivery.deliver(dropped)).ok(),
            Some(DeliveryOutcome::Dropped)
        );

        let rejected = OutboundDelivery::direct(connection_id.clone(), message("rejected"))
            .backpressure(BackpressurePolicy::Reject);
        assert_eq!(
            block_on(delivery.deliver(rejected))
                .as_ref()
                .map_err(soaprs_core::SoapError::kind),
            Err(SoapErrorKind::Unavailable)
        );

        let waiting = OutboundDelivery::direct(connection_id, message("waiting"));
        let mut waiting = pin!(delivery.deliver(waiting));
        let mut context = Context::from_waker(Waker::noop());
        assert!(matches!(waiting.as_mut().poll(&mut context), Poll::Pending));

        let first_batch = delivery.drain();
        assert_eq!(
            first_batch
                .ok()
                .and_then(|batch| batch.first().map(|item| item.envelope.message.clone())),
            Some("first".to_owned())
        );
        assert!(matches!(
            waiting.as_mut().poll(&mut context),
            Poll::Ready(Ok(DeliveryOutcome::Accepted))
        ));
        assert_eq!(delivery.len().ok(), Some(1));
    }
}