futu_backend/conn/
lifecycle.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
4
5use parking_lot::Mutex;
6#[cfg(test)]
7use tokio::sync::mpsc;
8use tokio::sync::watch;
9
10use super::diagnostics::{
11    EndpointFingerprint, InboundProgressSnapshot, PendingFailureFacts, PendingFailureKind,
12    PendingRegistrationIdentity, PendingResponseEntry, trace_pending_failure,
13};
14
15pub(super) type PendingResponses = Arc<Mutex<HashMap<u32, PendingResponseEntry>>>;
16
17pub(super) struct PendingRegistrationGuard {
18    pub(super) pending: PendingResponses,
19    pub(super) serial_no: u32,
20    pub(super) registration_identity: Arc<PendingRegistrationIdentity>,
21}
22
23impl Drop for PendingRegistrationGuard {
24    fn drop(&mut self) {
25        remove_pending_registration_if_current(
26            &self.pending,
27            self.serial_no,
28            &self.registration_identity,
29        );
30    }
31}
32
33fn remove_pending_registration_if_current(
34    pending: &PendingResponses,
35    serial_no: u32,
36    registration_identity: &Arc<PendingRegistrationIdentity>,
37) -> Option<PendingResponseEntry> {
38    let mut pending = pending.lock();
39    let is_current = pending
40        .get(&serial_no)
41        .is_some_and(|entry| Arc::ptr_eq(&entry.registration_identity, registration_identity));
42    if is_current {
43        pending.remove(&serial_no)
44    } else {
45        None
46    }
47}
48
49#[cfg(test)]
50pub(super) type DiagnosticSink = Option<mpsc::Sender<PendingFailureFacts>>;
51
52#[cfg(not(test))]
53#[derive(Clone)]
54pub(super) struct DiagnosticSink;
55
56#[derive(Clone)]
57pub(super) struct ConnectionTerminationFacts {
58    pub(super) kind: PendingFailureKind,
59    pub(super) endpoint_fingerprint: EndpointFingerprint,
60    pub(super) connection_generation: u64,
61    pub(super) progress: InboundProgressSnapshot,
62    pub(super) diagnostic_sink: DiagnosticSink,
63}
64
65pub(super) struct ConnectionLifecycle {
66    state: AtomicU8,
67    #[cfg(any(test, feature = "test-util"))]
68    pub(super) admission_publish_hook: Mutex<Option<LifecyclePauseHook>>,
69    #[cfg(test)]
70    pub(super) termination_drain_hook: Mutex<Option<LifecyclePauseHook>>,
71}
72
73#[cfg(any(test, feature = "test-util"))]
74#[derive(Clone, Default)]
75pub struct LifecyclePauseHook {
76    pub(super) entered: Arc<AtomicBool>,
77    pub(super) release: Arc<AtomicBool>,
78}
79
80#[cfg(any(test, feature = "test-util"))]
81impl LifecyclePauseHook {
82    pub fn entered(&self) -> bool {
83        self.entered.load(Ordering::Acquire)
84    }
85
86    pub fn release(&self) {
87        self.release.store(true, Ordering::Release);
88    }
89}
90
91impl ConnectionLifecycle {
92    const TERMINATED: u8 = 1 << 0;
93    const ADMISSION_ACTIVE: u8 = 1 << 1;
94    const WRITER_ACTIVE: u8 = 1 << 2;
95
96    pub(super) fn new() -> Self {
97        Self {
98            state: AtomicU8::new(0),
99            #[cfg(any(test, feature = "test-util"))]
100            admission_publish_hook: Mutex::new(None),
101            #[cfg(test)]
102            termination_drain_hook: Mutex::new(None),
103        }
104    }
105
106    pub(super) fn is_terminated(&self) -> bool {
107        self.state.load(Ordering::Acquire) & Self::TERMINATED != 0
108    }
109
110    fn claim_termination(&self) -> bool {
111        self.state.fetch_or(Self::TERMINATED, Ordering::AcqRel) & Self::TERMINATED == 0
112    }
113
114    pub(super) fn try_begin_writer(self: &Arc<Self>) -> Option<LifecycleActivityGuard> {
115        self.try_begin(Self::WRITER_ACTIVE)
116    }
117
118    pub(super) fn try_begin_admission(self: &Arc<Self>) -> Option<LifecycleActivityGuard> {
119        self.try_begin(Self::ADMISSION_ACTIVE)
120    }
121
122    #[cfg(any(test, feature = "test-util"))]
123    pub(super) fn pause_admission_publish_for_test(&self) {
124        Self::pause_for_test(self.admission_publish_hook.lock().clone());
125    }
126
127    #[cfg(test)]
128    fn pause_termination_drain_for_test(&self) {
129        Self::pause_for_test(self.termination_drain_hook.lock().clone());
130    }
131
132    #[cfg(any(test, feature = "test-util"))]
133    fn pause_for_test(hook: Option<LifecyclePauseHook>) {
134        if let Some(hook) = hook {
135            hook.entered.store(true, Ordering::Release);
136            while !hook.release.load(Ordering::Acquire) {
137                std::thread::yield_now();
138            }
139        }
140    }
141
142    fn try_begin(self: &Arc<Self>, activity: u8) -> Option<LifecycleActivityGuard> {
143        let mut current = self.state.load(Ordering::Acquire);
144        loop {
145            if current & Self::TERMINATED != 0 {
146                return None;
147            }
148            debug_assert_eq!(current & activity, 0);
149            match self.state.compare_exchange_weak(
150                current,
151                current | activity,
152                Ordering::AcqRel,
153                Ordering::Acquire,
154            ) {
155                Ok(_) => {
156                    return Some(LifecycleActivityGuard {
157                        lifecycle: Arc::clone(self),
158                        activity,
159                    });
160                }
161                Err(observed) => current = observed,
162            }
163        }
164    }
165}
166
167pub(super) struct LifecycleActivityGuard {
168    lifecycle: Arc<ConnectionLifecycle>,
169    activity: u8,
170}
171
172impl Drop for LifecycleActivityGuard {
173    fn drop(&mut self) {
174        self.lifecycle
175            .state
176            .fetch_and(!self.activity, Ordering::AcqRel);
177    }
178}
179
180pub(super) fn mark_disconnected(
181    connected: &Arc<std::sync::atomic::AtomicBool>,
182    connected_tx: &watch::Sender<bool>,
183) {
184    let was_connected = connected.swap(false, Ordering::AcqRel);
185    if was_connected {
186        let _ = connected_tx.send(false);
187    }
188}
189
190#[cfg(test)]
191pub(super) fn no_diagnostic_sink() -> DiagnosticSink {
192    None
193}
194
195#[cfg(not(test))]
196pub(super) fn no_diagnostic_sink() -> DiagnosticSink {
197    DiagnosticSink
198}
199
200pub(super) fn emit_diagnostic(diagnostic_sink: &DiagnosticSink, facts: &PendingFailureFacts) {
201    trace_pending_failure(facts);
202
203    #[cfg(test)]
204    if let Some(tx) = diagnostic_sink
205        && tx.try_send(facts.clone()).is_err()
206    {
207        tracing::debug!("bounded backend diagnostic test sink unavailable");
208    }
209
210    #[cfg(not(test))]
211    let _ = diagnostic_sink;
212}
213
214pub(super) fn pending_failure_facts(
215    kind: PendingFailureKind,
216    cmd_id: u16,
217    serial_no: u32,
218    writer_admitted: bool,
219    endpoint_fingerprint: &EndpointFingerprint,
220    connection_generation: u64,
221    progress: &InboundProgressSnapshot,
222) -> PendingFailureFacts {
223    PendingFailureFacts::new(
224        kind,
225        cmd_id,
226        serial_no,
227        writer_admitted,
228        endpoint_fingerprint.clone(),
229        connection_generation,
230        progress.received_bytes,
231        progress.frame_stage,
232        progress.codec_category,
233    )
234}
235
236fn fail_all_pending(pending: &PendingResponses, base: &ConnectionTerminationFacts) {
237    let entries: Vec<PendingResponseEntry> = pending
238        .lock()
239        .drain()
240        .map(|(serial_no, entry)| {
241            debug_assert_eq!(serial_no, entry.serial_no);
242            entry
243        })
244        .collect();
245
246    fail_pending_entries(entries, base);
247}
248
249fn fail_pending_entries(entries: Vec<PendingResponseEntry>, base: &ConnectionTerminationFacts) {
250    for entry in entries {
251        let facts = pending_failure_facts(
252            base.kind,
253            entry.cmd_id,
254            entry.serial_no,
255            entry.writer_admitted.load(Ordering::Acquire),
256            &base.endpoint_fingerprint,
257            base.connection_generation,
258            &base.progress,
259        );
260        emit_diagnostic(&base.diagnostic_sink, &facts);
261        if entry.tx.send(Err(facts)).is_err() {
262            tracing::debug!(
263                cmd_id = entry.cmd_id,
264                serial_no = entry.serial_no,
265                "pending response receiver dropped before transport failure delivery"
266            );
267        }
268    }
269}
270
271pub(super) fn claim_response_timeout_and_fail_other_pending(
272    termination_lifecycle: &ConnectionLifecycle,
273    connected: &Arc<AtomicBool>,
274    connected_tx: &watch::Sender<bool>,
275    shutdown_tx: &watch::Sender<bool>,
276    pending: &PendingResponses,
277    timed_out_serial_no: u32,
278    timed_out_registration_identity: &Arc<PendingRegistrationIdentity>,
279    base: &ConnectionTerminationFacts,
280) -> bool {
281    // The pending mutex makes exclusion of the triggering request atomic with
282    // claiming this generation. If another terminal observer already won, we
283    // leave our registration intact so that winner can deliver its typed facts.
284    let entries = {
285        let mut pending = pending.lock();
286        if !termination_lifecycle.claim_termination() {
287            return false;
288        }
289        let timed_out_entry_is_current = pending.get(&timed_out_serial_no).is_some_and(|entry| {
290            Arc::ptr_eq(
291                &entry.registration_identity,
292                timed_out_registration_identity,
293            )
294        });
295        if timed_out_entry_is_current {
296            pending.remove(&timed_out_serial_no);
297        }
298        mark_disconnected(connected, connected_tx);
299        pending
300            .drain()
301            .map(|(serial_no, entry)| {
302                debug_assert_eq!(serial_no, entry.serial_no);
303                entry
304            })
305            .collect()
306    };
307
308    let _ = shutdown_tx.send(true);
309    fail_pending_entries(entries, base);
310    true
311}
312
313pub(super) fn claim_termination_and_fail_all_pending(
314    termination_lifecycle: &ConnectionLifecycle,
315    connected: &Arc<AtomicBool>,
316    connected_tx: &watch::Sender<bool>,
317    shutdown_tx: &watch::Sender<bool>,
318    pending: &PendingResponses,
319    base: &ConnectionTerminationFacts,
320) -> bool {
321    if !termination_lifecycle.claim_termination() {
322        return false;
323    }
324    #[cfg(test)]
325    termination_lifecycle.pause_termination_drain_for_test();
326
327    // Wake both transport halves. The writer does not wait for a currently
328    // active sink operation; its activity bit only linearizes whether a frame
329    // had started before this claim. Later queued frames cannot acquire a new
330    // writer activity after TERMINATED is published.
331    let _ = shutdown_tx.send(true);
332
333    // Close the register-vs-terminate race under the same mutex used by
334    // request registration. A request either registered before this closure
335    // and is drained below, or observes connected=false while holding the
336    // mutex and cannot insert a new waiter.
337    {
338        let _pending_guard = pending.lock();
339        mark_disconnected(connected, connected_tx);
340    }
341    fail_all_pending(pending, base);
342    true
343}
futu_backend/conn/lifecycle.rs

futu_backend/conn/
lifecycle.rs
