openrtc 2.9.8

//! Keep a physical connection's close observable while ingress is backpressured.

use std::{
    future::{poll_fn, Future},
    pin::pin,
    task::Poll,
};

/// Forward with ordinary backpressure until this exact physical connection closes.
/// `None` cancels and drops the pending handoff; the caller retains close ownership.
pub(crate) async fn forward_or_closed<F: Future, C: Future>(
    forward: F,
    closed: C,
) -> Option<F::Output> {
    let mut forward = pin!(forward);
    let mut closed = pin!(closed);
    poll_fn(|cx| {
        // A ready close owns this physical generation, even if capacity also
        // became available. Never enqueue a pending stream after observing it.
        if closed.as_mut().poll(cx).is_ready() {
            return Poll::Ready(None);
        }
        forward.as_mut().poll(cx).map(Some)
    })
    .await
}

#[cfg(test)]
mod tests {
    use super::forward_or_closed;
    use std::{
        cell::RefCell,
        future::{poll_fn, Future},
        pin::Pin,
        rc::Rc,
        sync::{
            atomic::{AtomicUsize, Ordering},
            Arc,
        },
        task::{Context, Poll, Wake, Waker},
    };

    #[derive(Default)]
    struct WakeCount(AtomicUsize);

    impl Wake for WakeCount {
        fn wake(self: Arc<Self>) {
            self.0.fetch_add(1, Ordering::SeqCst);
        }
    }

    fn poll<F: Future>(future: Pin<&mut F>, wakes: &Arc<WakeCount>) -> Poll<F::Output> {
        future.poll(&mut Context::from_waker(&Waker::from(wakes.clone())))
    }

    #[derive(Clone, Default)]
    struct PhysicalClose(Rc<RefCell<(bool, Option<Waker>)>>);

    impl PhysicalClose {
        fn wait(&self) -> impl Future<Output = ()> {
            let state = self.0.clone();
            poll_fn(move |cx| {
                let mut state = state.borrow_mut();
                if state.0 {
                    Poll::Ready(())
                } else {
                    state.1 = Some(cx.waker().clone());
                    Poll::Pending
                }
            })
        }

        fn close(&self) {
            let waker = {
                let mut state = self.0.borrow_mut();
                state.0 = true;
                state.1.take()
            };
            if let Some(waker) = waker {
                waker.wake();
            }
        }
    }

    #[derive(Debug)]
    struct Stream {
        id: u64,
        drops: Arc<AtomicUsize>,
    }

    impl Drop for Stream {
        fn drop(&mut self) {
            self.drops.fetch_add(1, Ordering::SeqCst);
        }
    }

    // One slot is sufficient to exercise full/available/closed send semantics.
    // This is only the forward future; tests poll the production race primitive.
    #[derive(Default)]
    struct Queue {
        stream: Option<Stream>,
        closed: bool,
        sender: Option<Waker>,
    }

    impl Queue {
        fn send(
            queue: &Rc<RefCell<Self>>,
            stream: Stream,
        ) -> impl Future<Output = Result<(), Stream>> {
            let queue = queue.clone();
            let mut stream = Some(stream);
            poll_fn(move |cx| {
                let mut queue = queue.borrow_mut();
                if queue.closed {
                    Poll::Ready(Err(stream.take().unwrap()))
                } else if queue.stream.is_none() {
                    queue.stream = stream.take();
                    Poll::Ready(Ok(()))
                } else {
                    queue.sender = Some(cx.waker().clone());
                    Poll::Pending
                }
            })
        }

        fn drain(&mut self) -> Stream {
            let stream = self.stream.take().unwrap();
            if let Some(waker) = self.sender.take() {
                waker.wake();
            }
            stream
        }
    }

    fn stream(id: u64) -> (Stream, Arc<AtomicUsize>) {
        let drops = Arc::new(AtomicUsize::new(0));
        (
            Stream {
                id,
                drops: drops.clone(),
            },
            drops,
        )
    }

    fn full_queue() -> Rc<RefCell<Queue>> {
        Rc::new(RefCell::new(Queue {
            stream: Some(stream(1).0),
            ..Queue::default()
        }))
    }

    #[test]
    fn physical_close_cancels_full_handoff_without_consumer_drain() {
        let queue = full_queue();
        let close = PhysicalClose::default();
        let wakes = Arc::new(WakeCount::default());
        let (pending, drops) = stream(2);
        let mut forward = Box::pin(forward_or_closed(
            Queue::send(&queue, pending),
            close.wait(),
        ));
        assert!(poll(forward.as_mut(), &wakes).is_pending());
        close.close();
        assert!(
            matches!(poll(forward.as_mut(), &wakes), Poll::Ready(None)),
            "physical closure must release the owner even while the consumer never drains"
        );
        assert_eq!(wakes.0.load(Ordering::SeqCst), 1);
        assert_eq!(queue.borrow().stream.as_ref().unwrap().id, 1);
        assert_eq!(
            drops.load(Ordering::SeqCst),
            1,
            "drop only the unsent stream"
        );
    }

    #[test]
    fn live_full_handoff_keeps_backpressure_and_delivers_after_drain() {
        let queue = full_queue();
        let close = PhysicalClose::default();
        let wakes = Arc::new(WakeCount::default());
        let (pending, drops) = stream(2);
        let mut forward = Box::pin(forward_or_closed(
            Queue::send(&queue, pending),
            close.wait(),
        ));
        assert!(poll(forward.as_mut(), &wakes).is_pending());
        assert!(poll(forward.as_mut(), &wakes).is_pending());
        assert_eq!(drops.load(Ordering::SeqCst), 0);
        assert_eq!(queue.borrow_mut().drain().id, 1);
        assert_eq!(wakes.0.load(Ordering::SeqCst), 1);
        assert!(matches!(
            poll(forward.as_mut(), &wakes),
            Poll::Ready(Some(Ok(())))
        ));
        assert_eq!(queue.borrow().stream.as_ref().unwrap().id, 2);
        assert_eq!(drops.load(Ordering::SeqCst), 0);
    }

    #[test]
    fn physical_close_wins_when_capacity_also_becomes_available() {
        let queue = full_queue();
        let close = PhysicalClose::default();
        let wakes = Arc::new(WakeCount::default());
        let (pending, drops) = stream(2);
        let mut forward = Box::pin(forward_or_closed(
            Queue::send(&queue, pending),
            close.wait(),
        ));
        assert!(poll(forward.as_mut(), &wakes).is_pending());
        close.close();
        drop(queue.borrow_mut().drain());
        assert!(matches!(poll(forward.as_mut(), &wakes), Poll::Ready(None)));
        assert!(
            queue.borrow().stream.is_none(),
            "do not enqueue after physical close"
        );
        assert_eq!(drops.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn closing_old_physical_connection_does_not_cancel_replacement_or_sibling() {
        let old = PhysicalClose::default();
        let replacement = PhysicalClose::default();
        let sibling = PhysicalClose::default();
        let wakes = Arc::new(WakeCount::default());
        let old_queue = full_queue();
        let replacement_queue = full_queue();
        let sibling_queue = full_queue();
        let mut old_forward = Box::pin(forward_or_closed(
            Queue::send(&old_queue, stream(2).0),
            old.wait(),
        ));
        let mut replacement_forward = Box::pin(forward_or_closed(
            Queue::send(&replacement_queue, stream(3).0),
            replacement.wait(),
        ));
        let mut sibling_forward = Box::pin(forward_or_closed(
            Queue::send(&sibling_queue, stream(4).0),
            sibling.wait(),
        ));
        assert!(poll(old_forward.as_mut(), &wakes).is_pending());
        assert!(poll(replacement_forward.as_mut(), &wakes).is_pending());
        assert!(poll(sibling_forward.as_mut(), &wakes).is_pending());
        old.close();
        assert!(matches!(
            poll(old_forward.as_mut(), &wakes),
            Poll::Ready(None)
        ));
        assert!(poll(replacement_forward.as_mut(), &wakes).is_pending());
        assert!(poll(sibling_forward.as_mut(), &wakes).is_pending());
        drop(replacement_queue.borrow_mut().drain());
        drop(sibling_queue.borrow_mut().drain());
        assert!(matches!(
            poll(replacement_forward.as_mut(), &wakes),
            Poll::Ready(Some(Ok(())))
        ));
        assert!(matches!(
            poll(sibling_forward.as_mut(), &wakes),
            Poll::Ready(Some(Ok(())))
        ));
    }

    #[test]
    fn live_receiver_error_remains_visible_to_caller() {
        let queue = Rc::new(RefCell::new(Queue {
            closed: true,
            ..Queue::default()
        }));
        let close = PhysicalClose::default();
        let wakes = Arc::new(WakeCount::default());
        let mut forward = Box::pin(forward_or_closed(
            Queue::send(&queue, stream(2).0),
            close.wait(),
        ));
        match poll(forward.as_mut(), &wakes) {
            Poll::Ready(Some(Err(stream))) => assert_eq!(stream.id, 2),
            result => panic!("receiver error was lost: {result:?}"),
        }
    }
}