rx_rust/operators/connectable/
ref_count.rs

1use crate::delegate_disposal;
2use crate::disposable::{Disposable, chain_disposal::ChainDisposal};
3use crate::observable::{Observable, Subscription};
4use crate::observer::Observer;
5use crate::operators::connectable::connectable_controller::{
6    ConnectableController, Connected, Disconnected,
7};
8use crate::subject::Subject;
9use crate::subject::subject_observable::SubjectObservable;
10use crate::utils::mutable::{Mutable, MutableHelper};
11use crate::utils::on_panic::OnPanic;
12use crate::utils::types::{MaybeSend, Shared};
13use educe::Educe;
14use std::num::NonZeroUsize;
15
16#[derive(Educe)]
17#[educe(Debug)]
18enum State<OE, S, D>
19where
20    D: Disposable,
21{
22    Disconnected {
23        controller: ConnectableController<OE, S, Disconnected>,
24    },
25    ConnectingOrDisconnecting {
26        /// Use `usize` instead of `NonZeroUsize`. If it's 0, there are no subscribers and should
27        /// disconnect. If it's not 0, there is at least one subscriber and should connect.
28        subscribers: usize,
29    },
30    Connected {
31        subscribers: NonZeroUsize,
32        controller: ConnectableController<OE, S, Connected<D>>,
33    },
34}
35
36/// Makes a [`ConnectableController`] behave like an ordinary `Observable` that automatically connects
37/// on the first subscription and disconnects when the last subscription is disposed.
38/// See <https://reactivex.io/documentation/operators/refcount.html>
39///
40/// # Examples
41/// ```rust
42/// use rx_rust::{
43///     observable::ObservableExt,
44///     observer::Termination,
45///     operators::{
46///         connectable::{connectable_controller::ConnectableController, ref_count::RefCount},
47///         creating::from_iter::FromIter,
48///     },
49///     subject::publish_subject::PublishSubject,
50/// };
51/// use std::{convert::Infallible, sync::{Arc, Mutex}};
52///
53/// let values = Arc::new(Mutex::new(Vec::new()));
54/// let terminations = Arc::new(Mutex::new(Vec::new()));
55///
56/// let subject: PublishSubject<'_, i32, Infallible> = PublishSubject::default();
57/// let controller = ConnectableController::new(FromIter::new(vec![1, 2]), subject);
58/// let observable = controller.ref_count();
59/// let values_observer = Arc::clone(&values);
60/// let terminations_observer = Arc::clone(&terminations);
61///
62/// let subscription = observable.clone().subscribe_with_callback(
63///     move |value| values_observer.lock().unwrap().push(value),
64///     move |termination| terminations_observer
65///         .lock()
66///         .unwrap()
67///         .push(termination),
68/// );
69///
70/// drop(subscription);
71///
72/// assert_eq!(&*values.lock().unwrap(), &[1, 2]);
73/// assert_eq!(
74///     &*terminations.lock().unwrap(),
75///     &[Termination::Completed]
76/// );
77/// ```
78#[derive(Educe)]
79#[educe(Debug, Clone)]
80pub struct RefCount<'or, T, E, OE, S>
81where
82    OE: Observable<'or, T, E>,
83{
84    observable: SubjectObservable<S>,
85    state: Shared<Mutable<State<OE, S, OE::D>>>,
86}
87
88impl<'or, T, E, OE, S> RefCount<'or, T, E, OE, S>
89where
90    OE: Observable<'or, T, E>,
91    S: Clone,
92{
93    pub fn new(controller: ConnectableController<OE, S, Disconnected>) -> Self {
94        Self {
95            observable: controller.observable(),
96            state: Shared::new(Mutable::new(State::Disconnected { controller })),
97        }
98    }
99}
100
101delegate_disposal!(
102    Disposal<'or, T, E, OE, S>,
103    ChainDisposal<S::D, RefCountDisposal<'or, T, E, OE, S>>,
104    where OE: Observable<'or, T, E>,
105        S: Subject<'or, T, E>
106);
107
108impl<'or, T, E, OE, S> Observable<'or, T, E> for RefCount<'or, T, E, OE, S>
109where
110    OE: Observable<'or, T, E> + Clone,
111    S: Subject<'or, T, E> + Clone + MaybeSend + 'or,
112{
113    type D = Disposal<'or, T, E, OE, S>;
114
115    fn subscribe(self, observer: impl Observer<T, E> + MaybeSend + 'or) -> Subscription<Self::D> {
116        let controller = self.state.with_mut(|current| match &mut *current {
117            state @ State::Disconnected { .. } => {
118                let State::Disconnected { controller } =
119                    std::mem::replace(state, State::ConnectingOrDisconnecting { subscribers: 1 })
120                else {
121                    unreachable!()
122                };
123                Some(controller)
124            }
125            State::ConnectingOrDisconnecting { subscribers } => {
126                *subscribers = subscribers
127                    .checked_add(1)
128                    .expect("RefCount subscriber count overflowed");
129                None
130            }
131            State::Connected { subscribers, .. } => {
132                *subscribers = subscribers
133                    .checked_add(1)
134                    .expect("RefCount subscriber count overflowed");
135                None
136            }
137        });
138        // Subscribing notifies the observer inline when the subject has already terminated, so it
139        // can unwind while this subscriber is counted in and the controller is out of the state.
140        // Disarming the guard is what ends that scope: the connection that follows is not covered,
141        // and cannot be, since the controller is consumed by then.
142        let state = self.state.clone();
143        let guard = OnPanic::new(controller, move |controller| {
144            restore_subscriber(&state, controller);
145        });
146        let sub = self.observable.subscribe(observer);
147        let controller = guard.disarm();
148        if let Some(controller) = controller {
149            handle_connecting_or_disconnecting(self.state.clone(), Purpose::Connect(controller));
150        }
151        sub.then(RefCountDisposal { state: self.state }).map_into()
152    }
153}
154
155/// Puts back what [`RefCount::subscribe`] took, when subscribing the observer unwinds.
156///
157/// The subscriber is counted in — and the controller of a first subscription is taken out of the
158/// state — before the observer is subscribed to the subject. Without this, a panic from there
159/// would leave the count inflated forever, so the last subscription would no longer disconnect,
160/// and would lose the controller with the state stuck in
161/// [`ConnectingOrDisconnecting`](State::ConnectingOrDisconnecting).
162fn restore_subscriber<'or, T, E, OE, S>(
163    state: &Shared<Mutable<State<OE, S, OE::D>>>,
164    controller: Option<ConnectableController<OE, S, Disconnected>>,
165) where
166    OE: Observable<'or, T, E> + Clone,
167    S: Observer<T, E> + Clone + MaybeSend + 'or,
168{
169    let Some(controller) = controller else {
170        // Only the count was changed, so removing this subscriber is what disposing the
171        // subscription it never got would have done.
172        RefCountDisposal {
173            state: state.clone(),
174        }
175        .dispose();
176        return;
177    };
178    // The controller is held here, so nothing else can leave the connecting state, and the count
179    // can only have grown: the state is the one this subscription wrote.
180    let controller = state.with_mut(|current| {
181        let State::ConnectingOrDisconnecting { subscribers } = current else {
182            unreachable!("the controller of the connecting state is held by this guard")
183        };
184        let remaining = subscribers
185            .checked_sub(1)
186            .expect("RefCount subscriber count underflowed");
187        if remaining == 0 {
188            // Nobody is left to connect for, so the state goes back to what it was.
189            *current = State::Disconnected { controller };
190            None
191        } else {
192            *subscribers = remaining;
193            Some(controller)
194        }
195    });
196    if let Some(controller) = controller {
197        // Subscribers that arrived meanwhile are still waiting for the connection this
198        // subscription was going to make, so it is made here, as the returning path would.
199        handle_connecting_or_disconnecting(state.clone(), Purpose::Connect(controller));
200    }
201}
202
203struct RefCountDisposal<'or, T, E, OE, S>
204where
205    OE: Observable<'or, T, E>,
206{
207    state: Shared<Mutable<State<OE, S, OE::D>>>,
208}
209
210impl<'or, T, E, OE, S> Disposable for RefCountDisposal<'or, T, E, OE, S>
211where
212    OE: Observable<'or, T, E> + Clone,
213    S: Observer<T, E> + Clone + MaybeSend + 'or,
214{
215    fn dispose(self) {
216        let controller = self.state.with_mut(|current| match &mut *current {
217            State::Disconnected { .. } => unreachable!(),
218            State::ConnectingOrDisconnecting { subscribers } => {
219                *subscribers = subscribers
220                    .checked_sub(1)
221                    .expect("RefCount subscriber count underflowed");
222                None
223            }
224            State::Connected { subscribers, .. } if subscribers.get() > 1 => {
225                *subscribers = NonZeroUsize::new(subscribers.get() - 1)
226                    .expect("decremented subscriber count is non-zero");
227                None
228            }
229            State::Connected { .. } => {
230                let State::Connected { controller, .. } =
231                    std::mem::replace(current, State::ConnectingOrDisconnecting { subscribers: 0 })
232                else {
233                    unreachable!()
234                };
235                Some(controller)
236            }
237        });
238        if let Some(controller) = controller {
239            handle_connecting_or_disconnecting(self.state.clone(), Purpose::Disconnect(controller));
240        }
241    }
242}
243
244enum Purpose<'or, T, E, OE, S>
245where
246    OE: Observable<'or, T, E>,
247{
248    Connect(ConnectableController<OE, S, Disconnected>),
249    Disconnect(ConnectableController<OE, S, Connected<OE::D>>),
250}
251
252fn handle_connecting_or_disconnecting<'or, T, E, OE, S>(
253    state: Shared<Mutable<State<OE, S, OE::D>>>,
254    mut purpose: Purpose<'or, T, E, OE, S>,
255) where
256    OE: Observable<'or, T, E> + Clone,
257    S: Observer<T, E> + Clone + MaybeSend + 'or,
258{
259    loop {
260        let next = match purpose {
261            Purpose::Connect(controller) => {
262                let controller = controller.connect();
263                state.with_mut(|current| match &mut *current {
264                    State::Disconnected { .. } => unreachable!(),
265                    State::ConnectingOrDisconnecting { subscribers } => {
266                        if *subscribers == 0 {
267                            // It's unsubscribed, so we should disconnect
268                            Some(Purpose::Disconnect(controller))
269                        } else {
270                            *current = State::Connected {
271                                subscribers: NonZeroUsize::new(*subscribers).unwrap(),
272                                controller,
273                            };
274                            None
275                        }
276                    }
277                    State::Connected { .. } => unreachable!(),
278                })
279            }
280            Purpose::Disconnect(controller) => {
281                let controller = controller.disconnect();
282                state.with_mut(|current| match &mut *current {
283                    State::Disconnected { .. } => unreachable!(),
284                    State::ConnectingOrDisconnecting { subscribers } => {
285                        if *subscribers == 0 {
286                            *current = State::Disconnected { controller };
287                            None
288                        } else {
289                            // It's not unsubscribed, so we should reconnect
290                            Some(Purpose::Connect(controller))
291                        }
292                    }
293                    State::Connected { .. } => unreachable!(),
294                })
295            }
296        };
297        match next {
298            Some(next) => purpose = next,
299            None => break,
300        }
301    }
302}
rx_rust/operators/connectable/ref_count.rs

rx_rust/operators/connectable/
ref_count.rs
