ruchei/multicast/
buffered.rs

1use std::{
2    collections::VecDeque,
3    convert::Infallible,
4    ops::Index,
5    pin::Pin,
6    sync::Arc,
7    task::{Context, Poll, Wake},
8};
9
10use extend_pinned::ExtendPinned;
11use futures_util::{
12    Sink, SinkExt, Stream, TryStream, TryStreamExt, ready, stream::FusedStream, task::AtomicWaker,
13};
14use pin_project::pin_project;
15use ruchei_collections::{
16    as_linked_slab::{AsLinkedSlab, SlabKey},
17    linked_slab::LinkedSlab,
18};
19use ruchei_connection::{ConnectionWaker, Ready};
20use ruchei_extend::{Extending, ExtendingExt};
21
22use crate::connection_item::ConnectionItem;
23
24const OP_WAKE_NEXT: usize = 0;
25const OP_WAKE_READY: usize = 1;
26const OP_WAKE_FLUSH: usize = 1;
27const OP_WAKE_CLOSE: usize = 2;
28/// `start`ed, haven't yet reached the `flush_target`
29const OP_IS_S_PRE_F: usize = 3;
30/// `start`ed, already reached the `flush_target`
31const OP_IS_S_POST_F: usize = 4;
32/// `OP_IS_S_PRE_F` and `sent == items.len()`
33const OP_IS_FLUSHING: usize = 5;
34/// ordered by `sent`
35const OP_SENT_COUNT: usize = 6;
36/// first representative of `OP_SENT_COUNT` per `sent`
37const OP_SENT_FIRST: usize = 7;
38const OP_COUNT: usize = 8;
39
40#[derive(Debug)]
41pub(crate) struct Connection<S> {
42    pub(crate) stream: S,
43    pub(crate) next: Arc<ConnectionWaker>,
44    pub(crate) ready: Arc<ConnectionWaker>,
45    pub(crate) flush: Arc<ConnectionWaker>,
46    pub(crate) close: Arc<ConnectionWaker>,
47    sent: usize,
48    flushed: usize,
49}
50
51#[derive(Debug, Default)]
52struct NextFlush {
53    next: AtomicWaker,
54    flush: AtomicWaker,
55}
56
57impl Wake for NextFlush {
58    fn wake(self: Arc<Self>) {
59        self.next.wake();
60        self.flush.wake();
61    }
62}
63
64#[derive(Debug)]
65struct Item<T> {
66    item: T,
67    first: Option<SlabKey>,
68}
69
70#[derive(Debug)]
71struct Items<T> {
72    items: VecDeque<Item<T>>,
73    offset: usize,
74}
75
76impl<T> Default for Items<T> {
77    fn default() -> Self {
78        Self {
79            items: Default::default(),
80            offset: Default::default(),
81        }
82    }
83}
84
85impl<T> Items<T> {
86    #[must_use]
87    fn len(&self) -> usize {
88        self.items.len() + self.offset
89    }
90
91    fn push(&mut self, item: T, first: Option<SlabKey>) {
92        self.items.push_back(Item { item, first });
93    }
94}
95
96impl<T> Index<usize> for Items<T> {
97    type Output = T;
98
99    fn index(&self, index: usize) -> &Self::Output {
100        &self.items[index.checked_sub(self.offset).expect("early index")].item
101    }
102}
103
104#[pin_project]
105#[derive(Debug)]
106pub struct Multicast<S, T, E = <S as TryStream>::Error> {
107    connections: LinkedSlab<Connection<S>, OP_COUNT>,
108    #[pin]
109    next: Ready,
110    #[pin]
111    ready: Ready,
112    #[pin]
113    flush: Ready,
114    #[pin]
115    close: Ready,
116    items: Items<T>,
117    first_sent_all: Option<SlabKey>,
118    flush_target: usize,
119    next_flush: Arc<NextFlush>,
120    closed: VecDeque<(S, Option<E>)>,
121}
122
123impl<S, T, E> Default for Multicast<S, T, E> {
124    fn default() -> Self {
125        Self {
126            connections: Default::default(),
127            next: Default::default(),
128            ready: Default::default(),
129            flush: Default::default(),
130            close: Default::default(),
131            items: Default::default(),
132            first_sent_all: Default::default(),
133            flush_target: Default::default(),
134            next_flush: Default::default(),
135            closed: Default::default(),
136        }
137    }
138}
139
140impl<S: Unpin + Sink<T, Error = E>, T: Clone, E> Multicast<S, T, E> {
141    #[must_use]
142    fn first_for(self: Pin<&mut Self>, sent: usize) -> &mut Option<SlabKey> {
143        let this = self.project();
144        if sent == this.items.len() {
145            this.first_sent_all
146        } else {
147            &mut this.items.items[sent.checked_sub(this.items.offset).expect("early index")].first
148        }
149    }
150
151    #[must_use]
152    fn uncount_first(
153        mut self: Pin<&mut Self>,
154        key: SlabKey,
155        sent: usize,
156    ) -> (Option<SlabKey>, Option<SlabKey>) {
157        let mut this = self.as_mut().project();
158        assert!(this.connections.link_contains::<OP_SENT_FIRST>(key));
159        assert_eq!(this.connections[key].sent, sent);
160        assert_eq!(self.as_mut().first_for(sent).take(), Some(key));
161        this = self.as_mut().project();
162        let (long_prev, long_next) = this.connections.link_of::<OP_SENT_FIRST>(Some(key));
163        assert!(this.connections.link_pop_at::<OP_SENT_FIRST>(key));
164        let (_, short_next) = this.connections.link_of::<OP_SENT_COUNT>(Some(key));
165        assert!(this.connections.link_pop_at::<OP_SENT_COUNT>(key));
166        if let Some(short_next) = short_next
167            && this.connections[short_next].sent == sent
168        {
169            assert!(!this.connections.link_contains::<OP_SENT_FIRST>(short_next));
170            this.connections
171                .link_insert::<OP_SENT_FIRST>(long_prev, short_next, long_next);
172            *self.first_for(sent) = Some(short_next);
173            (Some(short_next), long_next)
174        } else {
175            (long_prev, long_next)
176        }
177    }
178
179    #[must_use]
180    fn uncount_non_first(
181        mut self: Pin<&mut Self>,
182        key: SlabKey,
183        sent: usize,
184    ) -> (SlabKey, Option<SlabKey>) {
185        let mut this = self.as_mut().project();
186        assert!(!this.connections.link_contains::<OP_SENT_FIRST>(key));
187        assert_eq!(this.connections[key].sent, sent);
188        let first = self
189            .as_mut()
190            .first_for(sent)
191            .as_ref()
192            .copied()
193            .expect("first not found");
194        assert_ne!(first, key);
195        this = self.project();
196        let (_, long_next) = this.connections.link_of::<OP_SENT_FIRST>(Some(first));
197        assert!(this.connections.link_pop_at::<OP_SENT_COUNT>(key));
198        (first, long_next)
199    }
200
201    #[must_use]
202    fn uncount(
203        mut self: Pin<&mut Self>,
204        key: SlabKey,
205        sent: usize,
206    ) -> (Option<SlabKey>, Option<SlabKey>) {
207        let mut this = self.as_mut().project();
208        assert_eq!(this.connections[key].sent, sent);
209        let (prev, next) = if this.connections.link_contains::<OP_SENT_FIRST>(key) {
210            self.as_mut().uncount_first(key, sent)
211        } else {
212            let (prev, next) = self.as_mut().uncount_non_first(key, sent);
213            (Some(prev), next)
214        };
215        this = self.as_mut().project();
216        if let Some(prev) = prev {
217            assert!(this.connections.link_contains::<OP_SENT_FIRST>(prev));
218            let prev_sent = this.connections[prev].sent;
219            assert!(prev_sent <= sent);
220            assert_eq!(*self.as_mut().first_for(prev_sent), Some(prev));
221            this = self.as_mut().project();
222        }
223        if let Some(next) = next {
224            assert!(this.connections.link_contains::<OP_SENT_FIRST>(next));
225            let next_sent = this.connections[next].sent;
226            assert!(sent < next_sent);
227            assert_eq!(*self.as_mut().first_for(next_sent), Some(next));
228            this = self.as_mut().project();
229        }
230        assert!(!this.connections.link_contains::<OP_SENT_FIRST>(key));
231        assert!(!this.connections.link_contains::<OP_SENT_COUNT>(key));
232        (prev, next)
233    }
234
235    fn count(
236        mut self: Pin<&mut Self>,
237        prev: Option<SlabKey>,
238        next: Option<SlabKey>,
239        key: SlabKey,
240        sent: usize,
241    ) {
242        let mut this = self.as_mut().project();
243        assert_eq!(this.connections[key].sent, sent);
244        assert!(!this.connections.link_contains::<OP_SENT_FIRST>(key));
245        assert!(!this.connections.link_contains::<OP_SENT_COUNT>(key));
246        if let Some(prev) = prev {
247            assert!(this.connections.link_contains::<OP_SENT_FIRST>(prev));
248            let prev_sent = this.connections[prev].sent;
249            assert!(prev_sent < sent);
250            assert_eq!(*self.as_mut().first_for(prev_sent), Some(prev));
251            this = self.as_mut().project();
252        }
253        if let Some(next) = next {
254            assert!(this.connections.link_contains::<OP_SENT_FIRST>(next));
255            let next_sent = this.connections[next].sent;
256            assert!(sent <= next_sent);
257            assert_eq!(*self.as_mut().first_for(next_sent), Some(next));
258            this = self.as_mut().project();
259        }
260        let (_, prev_next) = this.connections.link_of::<OP_SENT_FIRST>(prev);
261        assert_eq!(prev_next, next);
262        let (next_prev, long_next) = this.connections.link_of::<OP_SENT_FIRST>(next);
263        assert_eq!(next_prev, prev);
264        if let Some(next) = next
265            && sent == this.connections[next].sent
266        {
267            let (short_prev, _) = this.connections.link_of::<OP_SENT_COUNT>(long_next);
268            let short_prev = short_prev.expect("should at least be what next is");
269            assert_eq!(this.connections[short_prev].sent, sent);
270            this.connections
271                .link_insert::<OP_SENT_COUNT>(Some(short_prev), key, long_next);
272        } else {
273            this.connections
274                .link_insert::<OP_SENT_FIRST>(prev, key, next);
275            let (short_prev, _) = this.connections.link_of::<OP_SENT_COUNT>(next);
276            match (prev, short_prev) {
277                (None, None) => {}
278                (Some(prev), Some(short_prev)) => {
279                    let sent = this.connections[short_prev].sent;
280                    assert_eq!(this.connections[prev].sent, sent);
281                    assert_eq!(*self.as_mut().first_for(sent), Some(prev));
282                    this = self.as_mut().project();
283                }
284                _ => panic!("inconsistent state"),
285            }
286            this.connections
287                .link_insert::<OP_SENT_COUNT>(short_prev, key, next);
288            *self.first_for(sent) = Some(key);
289        }
290    }
291
292    fn increment_sent(mut self: Pin<&mut Self>, key: SlabKey, sent: usize) {
293        let (prev, next) = self.as_mut().uncount(key, sent);
294        let this = self.as_mut().project();
295        this.connections[key].sent += 1;
296        let sent = this.connections[key].sent;
297        self.count(prev, next, key, sent);
298    }
299
300    fn remove(mut self: Pin<&mut Self>, key: SlabKey, error: Option<E>) {
301        let mut this = self.as_mut().project();
302        if this.connections.link_contains::<OP_SENT_FIRST>(key) {
303            let sent = this.connections[key].sent;
304            let _ = self.as_mut().uncount_first(key, sent);
305            this = self.project();
306        }
307        let connection = this.connections.remove(key);
308        connection.next.wake();
309        connection.ready.wake();
310        connection.flush.wake();
311        connection.close.wake();
312        this.closed.push_back((connection.stream, error));
313        this.next.wake();
314    }
315
316    pub fn push(self: Pin<&mut Self>, stream: S) {
317        let this = self.project();
318        let key = this.connections.vacant_key();
319        let next = this.next.downgrade();
320        let ready = this.ready.downgrade();
321        let flush = this.flush.downgrade();
322        let close = this.close.downgrade();
323        let sent = this.items.len();
324        let connection = Connection {
325            stream,
326            next: ConnectionWaker::new(key, next),
327            ready: ConnectionWaker::new(key, ready),
328            flush: ConnectionWaker::new(key, flush),
329            close: ConnectionWaker::new(key, close),
330            sent,
331            flushed: sent,
332        };
333        this.connections.insert_at(key, connection);
334        assert!(this.connections.link_push_back::<OP_WAKE_NEXT>(key));
335        assert!(this.connections.link_push_back::<OP_WAKE_READY>(key));
336        assert!(this.connections.link_push_back::<OP_WAKE_CLOSE>(key));
337        this.next.wake();
338        this.ready.wake();
339        this.close.wake();
340        assert!(this.connections.link_push_back::<OP_SENT_COUNT>(key));
341        if this.first_sent_all.is_none() {
342            assert!(this.connections.link_push_back::<OP_SENT_FIRST>(key));
343            *this.first_sent_all = Some(key);
344        }
345    }
346
347    fn start_flush_one(self: Pin<&mut Self>, key: SlabKey) {
348        let this = self.project();
349        assert!(this.connections[key].sent == this.items.len());
350        assert!(this.connections.link_contains::<OP_IS_S_PRE_F>(key));
351        assert!(this.connections[key].flushed < *this.flush_target);
352        assert!(this.connections.link_push_back::<OP_IS_FLUSHING>(key));
353        this.flush.downgrade().insert(key);
354    }
355
356    /// wait until `sent` reaches `items.len()`
357    fn poll_send_one(
358        mut self: Pin<&mut Self>,
359        key: SlabKey,
360        cx: &mut Context<'_>,
361    ) -> Poll<Result<(), S::Error>> {
362        let mut this = self.as_mut().project();
363        assert!(this.connections[key].sent < this.items.len());
364        while this.connections[key].sent < this.items.len() {
365            ready!(this.connections[key].stream.poll_ready_unpin(cx))?;
366            let sent = this.connections[key].sent;
367            let item = this.items[sent].clone();
368            this.connections[key].stream.start_send_unpin(item)?;
369            if !this.connections.link_contains::<OP_IS_S_PRE_F>(key) {
370                if this.connections[key].flushed < *this.flush_target {
371                    this.connections.link_push_back::<OP_IS_S_PRE_F>(key);
372                } else {
373                    this.connections.link_push_back::<OP_IS_S_POST_F>(key);
374                }
375            }
376            self.as_mut().increment_sent(key, sent);
377            this = self.as_mut().project();
378        }
379        if this.connections.link_contains::<OP_IS_S_PRE_F>(key) {
380            self.as_mut().start_flush_one(key);
381        }
382        Poll::Ready(Ok(()))
383    }
384
385    /// wait until `flushed` reaches `flush_target`
386    fn poll_flush_one(
387        self: Pin<&mut Self>,
388        key: SlabKey,
389        cx: &mut Context<'_>,
390    ) -> Poll<Result<(), S::Error>> {
391        let this = self.project();
392        assert!(this.connections.link_contains::<OP_IS_FLUSHING>(key));
393        assert!(this.connections[key].sent == this.items.len());
394        assert!(this.connections.link_contains::<OP_IS_S_PRE_F>(key));
395        ready!(this.connections[key].stream.poll_flush_unpin(cx))?;
396        this.connections[key].flushed = this.connections[key].sent;
397        assert!(this.connections.link_pop_at::<OP_IS_FLUSHING>(key));
398        assert!(this.connections.link_pop_at::<OP_IS_S_PRE_F>(key));
399        Poll::Ready(Ok(()))
400    }
401
402    fn poll_send_all(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
403        let mut this = self.as_mut().project();
404        this.ready.register(cx);
405        while let Some(key) = this.ready.as_mut().next::<OP_WAKE_READY>(this.connections) {
406            if this.connections[key].sent < this.items.len()
407                && let Some(connection) = this.connections.get_mut(key)
408                && let Poll::Ready(Err(e)) = connection
409                    .ready
410                    .clone()
411                    .poll(cx, |cx| self.as_mut().poll_send_one(key, cx))
412            {
413                self.as_mut().remove(key, Some(e));
414            }
415            this = self.as_mut().project();
416        }
417        while let Some(item) = this.items.items.front()
418            && item.first.is_none()
419        {
420            this.items.items.pop_front();
421            this.items.offset += 1;
422        }
423        if this.items.items.is_empty() {
424            Poll::Ready(())
425        } else {
426            Poll::Pending
427        }
428    }
429
430    fn poll_flush_all(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
431        let mut this = self.as_mut().project();
432        this.flush.register(cx);
433        while let Some(key) = this.flush.as_mut().next::<OP_WAKE_FLUSH>(this.connections) {
434            if this.connections.link_contains::<OP_IS_FLUSHING>(key)
435                && let Some(connection) = this.connections.get_mut(key)
436                && let Poll::Ready(Err(e)) = connection
437                    .flush
438                    .clone()
439                    .poll(cx, |cx| self.as_mut().poll_flush_one(key, cx))
440            {
441                self.as_mut().remove(key, Some(e));
442            }
443            this = self.as_mut().project();
444        }
445        if this.connections.link_empty::<OP_IS_FLUSHING>() {
446            Poll::Ready(())
447        } else {
448            Poll::Pending
449        }
450    }
451
452    fn poll_send_flush(mut self: Pin<&mut Self>) -> Poll<()> {
453        let waker = self.next_flush.clone().into();
454        let cx = &mut Context::from_waker(&waker);
455        let sent = self.as_mut().poll_send_all(cx);
456        ready!(self.as_mut().poll_flush_all(cx));
457        sent
458    }
459
460    fn start_flush(mut self: Pin<&mut Self>) {
461        let mut this = self.as_mut().project();
462        assert!(*this.flush_target < this.items.len());
463        *this.flush_target = this.items.len();
464        while let Some(key) = this.connections.link_pop_front::<OP_IS_S_POST_F>() {
465            assert!(this.connections[key].flushed < *this.flush_target);
466            this.connections.link_push_back::<OP_IS_S_PRE_F>(key);
467            if this.connections[key].sent == this.items.len() {
468                self.as_mut().start_flush_one(key);
469                this = self.as_mut().project();
470            }
471        }
472    }
473}
474
475impl<S: Unpin + TryStream<Error = E> + Sink<T, Error = E>, T: Clone, E> Stream
476    for Multicast<S, T, E>
477{
478    type Item = ConnectionItem<S>;
479
480    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
481        self.next_flush.next.register(cx.waker());
482        let _ = self.as_mut().poll_send_flush();
483        let mut this = self.as_mut().project();
484        if let Some((stream, error)) = this.closed.pop_front() {
485            return Poll::Ready(Some(ConnectionItem::Closed(stream, error)));
486        }
487        while let Some(key) = this.next.as_mut().next::<OP_WAKE_NEXT>(this.connections) {
488            if let Some(connection) = this.connections.get_mut(key)
489                && let Poll::Ready(o) = connection
490                    .next
491                    .poll(cx, |cx| connection.stream.try_poll_next_unpin(cx))
492            {
493                match o {
494                    Some(Ok(item)) => {
495                        this.next.downgrade().insert(key);
496                        return Poll::Ready(Some(ConnectionItem::Item(item)));
497                    }
498                    Some(Err(e)) => {
499                        self.as_mut().remove(key, Some(e));
500                    }
501                    None => {
502                        self.as_mut().remove(key, None);
503                    }
504                }
505            }
506            this = self.as_mut().project();
507        }
508        if this.connections.is_empty() {
509            Poll::Ready(None)
510        } else {
511            Poll::Pending
512        }
513    }
514}
515
516impl<S: Unpin + TryStream<Error = E> + Sink<T, Error = E>, T: Clone, E> FusedStream
517    for Multicast<S, T, E>
518{
519    fn is_terminated(&self) -> bool {
520        self.closed.is_empty() && self.connections.is_empty()
521    }
522}
523
524impl<S: Unpin + Sink<T, Error = E>, T: Clone, E> Sink<T> for Multicast<S, T, E> {
525    type Error = Infallible;
526
527    fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
528        Poll::Ready(Ok(()))
529    }
530
531    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
532        let this = self.project();
533        let mut key = this.first_sent_all.as_ref().copied();
534        while let Some(k) = key {
535            this.connections.link_pop_at::<OP_IS_FLUSHING>(k);
536            this.ready.downgrade().insert(k);
537            (_, key) = this.connections.link_of::<OP_SENT_COUNT>(key);
538        }
539        this.items.push(item, this.first_sent_all.take());
540        Ok(())
541    }
542
543    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
544        self.next_flush.flush.register(cx.waker());
545        if self.flush_target < self.items.len() {
546            self.as_mut().start_flush();
547        }
548        ready!(self.poll_send_flush());
549        Poll::Ready(Ok(()))
550    }
551
552    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
553        ready!(self.as_mut().poll_send_all(cx));
554        let mut this = self.as_mut().project();
555        this.close.register(cx);
556        while let Some(key) = this.close.as_mut().next::<OP_WAKE_CLOSE>(this.connections) {
557            if let Some(connection) = this.connections.get_mut(key)
558                && let Poll::Ready(r) = connection
559                    .close
560                    .poll(cx, |cx| connection.stream.poll_close_unpin(cx))
561            {
562                match r {
563                    Ok(()) => {
564                        self.as_mut().remove(key, None);
565                    }
566                    Err(e) => {
567                        self.as_mut().remove(key, Some(e));
568                    }
569                }
570            }
571            this = self.as_mut().project();
572        }
573        if this.connections.is_empty() {
574            Poll::Ready(Ok(()))
575        } else {
576            Poll::Pending
577        }
578    }
579}
580
581impl<S: Unpin + Sink<T, Error = E>, T: Clone, E> ExtendPinned<S> for Multicast<S, T, E> {
582    fn extend_pinned<I: IntoIterator<Item = S>>(mut self: Pin<&mut Self>, iter: I) {
583        for stream in iter {
584            self.as_mut().push(stream);
585        }
586    }
587}
588
589pub type MulticastExtending<T, R> = Extending<Multicast<<R as MulticastBuffered<T>>::S, T>, R>;
590
591pub trait MulticastBuffered<T: Clone>: Sized + Stream<Item = Self::S> {
592    /// Single [`Stream`]/[`Sink`].
593    type S: Unpin + TryStream<Error = Self::E> + Sink<T, Error = Self::E>;
594    /// Error.
595    type E;
596
597    #[must_use]
598    fn multicast_buffered(self) -> MulticastExtending<T, Self> {
599        self.extending_default()
600    }
601}
602
603impl<S: Unpin + TryStream<Error = E> + Sink<T, Error = E>, T: Clone, E, R: Stream<Item = S>>
604    MulticastBuffered<T> for R
605{
606    type S = S;
607    type E = E;
608}
ruchei/multicast/buffered.rs

ruchei/multicast/
buffered.rs
