fusen_register/
lib.rs

1#![warn(missing_docs)]
2//! Cancellation-safe service registration and discovery contracts for fusen-rs.
3
4use fusen_contract::{ServiceRegistration, ServiceSelector};
5use futures_util::FutureExt;
6use std::{
7    future::Future,
8    panic::{AssertUnwindSafe, catch_unwind},
9    pin::Pin,
10    sync::{
11        Arc, Mutex,
12        atomic::{AtomicBool, AtomicUsize, Ordering},
13    },
14};
15use tokio::{runtime::Handle, sync::watch};
16
17use crate::{
18    directory::Directory,
19    error::{RegistryError, RegistryErrorKind, RegistryOperation},
20};
21
22/// Latest-wins service directories and provider publication handles.
23pub mod directory;
24/// Classified registry failures.
25pub mod error;
26/// Owned, sendable future returned by registry lifecycle APIs.
27pub type RegistryFuture<T> =
28    Pin<Box<dyn Future<Output = Result<T, RegistryError>> + Send + 'static>>;
29
30/// Parameters for preparing one provider registration.
31#[derive(Clone, Debug)]
32pub struct RegistrationRequest {
33    registration: Arc<ServiceRegistration>,
34}
35
36impl RegistrationRequest {
37    /// Creates a registration request.
38    pub fn new(registration: Arc<ServiceRegistration>) -> Self {
39        Self { registration }
40    }
41
42    /// Returns the immutable provider registration.
43    pub fn registration(&self) -> &Arc<ServiceRegistration> {
44        &self.registration
45    }
46
47    /// Consumes this request into the immutable provider registration.
48    pub fn into_registration(self) -> Arc<ServiceRegistration> {
49        self.registration
50    }
51}
52
53/// Parameters for preparing one discovery subscription.
54#[derive(Clone, Debug)]
55pub struct SubscriptionRequest {
56    selector: ServiceSelector,
57}
58
59impl SubscriptionRequest {
60    /// Creates a subscription request.
61    pub fn new(selector: ServiceSelector) -> Self {
62        Self { selector }
63    }
64
65    /// Returns the service selector.
66    pub const fn selector(&self) -> &ServiceSelector {
67        &self.selector
68    }
69
70    /// Consumes this request into the service selector.
71    pub fn into_selector(self) -> ServiceSelector {
72        self.selector
73    }
74}
75
76/// Pluggable provider that prepares registration and subscription ownership before activation.
77///
78/// Implementations must construct handles without starting remote side effects. The runtime stores
79/// each handle before calling its `activate` method, so cancellation always has a cleanup owner.
80pub trait Registry: Send + Sync + 'static {
81    /// Prepares one service registration without publishing it yet.
82    fn prepare_registration(
83        &self,
84        request: RegistrationRequest,
85    ) -> Result<RegistrationHandle, RegistryError>;
86
87    /// Prepares one service subscription without installing it yet.
88    fn prepare_subscription(
89        &self,
90        request: SubscriptionRequest,
91    ) -> Result<SubscriptionHandle, RegistryError>;
92}
93
94impl<T> Registry for Arc<T>
95where
96    T: Registry + ?Sized,
97{
98    fn prepare_registration(
99        &self,
100        request: RegistrationRequest,
101    ) -> Result<RegistrationHandle, RegistryError> {
102        (**self).prepare_registration(request)
103    }
104
105    fn prepare_subscription(
106        &self,
107        request: SubscriptionRequest,
108    ) -> Result<SubscriptionHandle, RegistryError> {
109        (**self).prepare_subscription(request)
110    }
111}
112
113/// Safe constructors for provider-owned registry lifecycles.
114pub mod provider {
115    use super::*;
116
117    /// Creates a registration handle from activation and cleanup operations.
118    pub fn registration<A, C, CF>(activate: A, close: C) -> RegistrationHandle
119    where
120        A: Future<Output = Result<(), RegistryError>> + Send + 'static,
121        C: FnOnce() -> CF + Send + 'static,
122        CF: Future<Output = Result<(), RegistryError>> + Send + 'static,
123    {
124        super::prepare_registration(activate, close)
125    }
126
127    /// Creates a subscription handle and stable directory from provider operations.
128    pub fn subscription<A, C, CF>(directory: Directory, activate: A, close: C) -> SubscriptionHandle
129    where
130        A: Future<Output = Result<(), RegistryError>> + Send + 'static,
131        C: FnOnce() -> CF + Send + 'static,
132        CF: Future<Output = Result<(), RegistryError>> + Send + 'static,
133    {
134        super::prepare_subscription(directory, activate, close)
135    }
136}
137
138/// Creates a registration handle from provider-owned activation and cleanup operations.
139///
140/// Neither future is polled before the first call to [`RegistrationHandle::activate`]. Cleanup is
141/// constructed at most once and only after activation has reached a terminal result.
142fn prepare_registration<A, C, CF>(activate: A, close: C) -> RegistrationHandle
143where
144    A: Future<Output = Result<(), RegistryError>> + Send + 'static,
145    C: FnOnce() -> CF + Send + 'static,
146    CF: Future<Output = Result<(), RegistryError>> + Send + 'static,
147{
148    RegistrationHandle {
149        lifecycle: Arc::new(Lifecycle::new(
150            RegistryOperation::ActivateRegistration,
151            RegistryOperation::CloseRegistration,
152            Box::pin(activate),
153            Box::new(move || Box::pin(close())),
154        )),
155    }
156}
157
158/// Creates a subscription handle from a provider directory, activation, and cleanup operations.
159///
160/// Neither future is polled before the first call to [`SubscriptionHandle::activate`]. Cleanup is
161/// constructed at most once and only after activation has reached a terminal result.
162fn prepare_subscription<A, C, CF>(directory: Directory, activate: A, close: C) -> SubscriptionHandle
163where
164    A: Future<Output = Result<(), RegistryError>> + Send + 'static,
165    C: FnOnce() -> CF + Send + 'static,
166    CF: Future<Output = Result<(), RegistryError>> + Send + 'static,
167{
168    SubscriptionHandle {
169        lifecycle: Arc::new(Lifecycle::new(
170            RegistryOperation::ActivateSubscription,
171            RegistryOperation::CloseSubscription,
172            Box::pin(activate),
173            Box::new(move || Box::pin(close())),
174        )),
175        directory,
176    }
177}
178
179/// Prepared ownership of one provider registration.
180///
181/// Clones share activation and cleanup terminal results. Dropping the last clone only requests
182/// cleanup; provider work remains owned by the worker started during activation. Cancelling every
183/// pending activation waiter also requests cleanup, so a late provider success is compensated.
184#[derive(Clone)]
185pub struct RegistrationHandle {
186    lifecycle: Arc<Lifecycle>,
187}
188
189impl RegistrationHandle {
190    /// Starts provider activation once and shares its terminal result with every caller.
191    pub fn activate(&self) -> RegistryFuture<()> {
192        let lifecycle = self.lifecycle.clone();
193        let mut waiter = ActivationWaiter::new(lifecycle.clone());
194        Box::pin(async move {
195            let result = match lifecycle.ensure_started() {
196                Ok(()) => lifecycle.wait_activation().await,
197                Err(error) => Err(error),
198            };
199            waiter.complete();
200            result
201        })
202    }
203
204    /// Requests cleanup without waiting for provider completion.
205    pub fn request_close(&self) {
206        self.lifecycle.request_close();
207    }
208
209    /// Requests cleanup and shares its terminal result with every caller.
210    pub fn close(&self) -> RegistryFuture<()> {
211        let lifecycle = self.lifecycle.clone();
212        Box::pin(async move {
213            lifecycle.request_close();
214            lifecycle.wait_close().await
215        })
216    }
217}
218
219impl std::fmt::Debug for RegistrationHandle {
220    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        formatter
222            .debug_struct("RegistrationHandle")
223            .finish_non_exhaustive()
224    }
225}
226
227/// Prepared ownership of one provider subscription and its stable directory.
228///
229/// Clones share activation and cleanup terminal results. Dropping the last clone only requests
230/// cleanup; provider work remains owned by the worker started during activation. Cancelling every
231/// pending activation waiter also requests cleanup, so a late provider success is compensated.
232#[derive(Clone)]
233pub struct SubscriptionHandle {
234    lifecycle: Arc<Lifecycle>,
235    directory: Directory,
236}
237
238impl SubscriptionHandle {
239    /// Starts provider activation once and returns the shared directory after successful setup.
240    pub fn activate(&self) -> RegistryFuture<Directory> {
241        let lifecycle = self.lifecycle.clone();
242        let directory = self.directory.clone();
243        let mut waiter = ActivationWaiter::new(lifecycle.clone());
244        Box::pin(async move {
245            let result = match lifecycle.ensure_started() {
246                Ok(()) => lifecycle.wait_activation().await.map(|()| directory),
247                Err(error) => Err(error),
248            };
249            waiter.complete();
250            result
251        })
252    }
253
254    /// Requests cleanup without waiting for provider completion.
255    pub fn request_close(&self) {
256        self.lifecycle.request_close();
257    }
258
259    /// Requests cleanup and shares its terminal result with every caller.
260    pub fn close(&self) -> RegistryFuture<()> {
261        let lifecycle = self.lifecycle.clone();
262        Box::pin(async move {
263            lifecycle.request_close();
264            lifecycle.wait_close().await
265        })
266    }
267}
268
269impl std::fmt::Debug for SubscriptionHandle {
270    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
271        formatter
272            .debug_struct("SubscriptionHandle")
273            .field("directory", &self.directory)
274            .finish_non_exhaustive()
275    }
276}
277
278type CloseFactory = Box<dyn FnOnce() -> RegistryFuture<()> + Send + 'static>;
279type SharedResult = Option<Result<(), RegistryError>>;
280
281struct PreparedLifecycle {
282    activate: RegistryFuture<()>,
283    close: CloseFactory,
284}
285
286enum StartState {
287    Prepared(Option<PreparedLifecycle>),
288    Started,
289    Finished,
290}
291
292struct Lifecycle {
293    activation_operation: RegistryOperation,
294    close_operation: RegistryOperation,
295    start: Mutex<StartState>,
296    activation_waiters: AtomicUsize,
297    activation_observed: AtomicBool,
298    close_requested: AtomicBool,
299    close_request: watch::Sender<bool>,
300    activation_result: watch::Sender<SharedResult>,
301    close_result: watch::Sender<SharedResult>,
302}
303
304impl Lifecycle {
305    fn new(
306        activation_operation: RegistryOperation,
307        close_operation: RegistryOperation,
308        activate: RegistryFuture<()>,
309        close: CloseFactory,
310    ) -> Self {
311        let (close_request, _) = watch::channel(false);
312        let (activation_result, _) = watch::channel(None);
313        let (close_result, _) = watch::channel(None);
314        Self {
315            activation_operation,
316            close_operation,
317            start: Mutex::new(StartState::Prepared(Some(PreparedLifecycle {
318                activate,
319                close,
320            }))),
321            activation_waiters: AtomicUsize::new(0),
322            activation_observed: AtomicBool::new(false),
323            close_requested: AtomicBool::new(false),
324            close_request,
325            activation_result,
326            close_result,
327        }
328    }
329
330    fn ensure_started(&self) -> Result<(), RegistryError> {
331        if self.activation_result.borrow().is_some() {
332            return Ok(());
333        }
334        let runtime = match Handle::try_current() {
335            Ok(runtime) => runtime,
336            Err(error) => {
337                let error = RegistryError::new(
338                    self.activation_operation,
339                    RegistryErrorKind::Internal,
340                    error,
341                );
342                self.finish_before_start(Err(error.clone()));
343                return Err(error);
344            }
345        };
346        let mut start = self.start.lock().unwrap_or_else(|error| error.into_inner());
347        match &mut *start {
348            StartState::Started => return Ok(()),
349            StartState::Finished => return Ok(()),
350            StartState::Prepared(_) => {}
351        }
352        if self.close_requested.load(Ordering::Acquire) {
353            let prepared = match &mut *start {
354                StartState::Prepared(prepared) => prepared.take(),
355                StartState::Started | StartState::Finished => None,
356            };
357            drop(prepared);
358            *start = StartState::Finished;
359            self.publish_pre_activation_close();
360            return Ok(());
361        }
362        let prepared = match &mut *start {
363            StartState::Prepared(prepared) => prepared
364                .take()
365                .expect("prepared lifecycle is present before activation"),
366            StartState::Started | StartState::Finished => unreachable!(),
367        };
368        let worker = LifecycleWorker {
369            activation_operation: self.activation_operation,
370            close_operation: self.close_operation,
371            activate: Some(prepared.activate),
372            close: Some(prepared.close),
373            close_request: self.close_request.subscribe(),
374            activation_result: self.activation_result.clone(),
375            close_result: self.close_result.clone(),
376            activation_published: false,
377            close_published: false,
378        };
379        runtime.spawn(worker.run());
380        *start = StartState::Started;
381        Ok(())
382    }
383
384    fn request_close(&self) {
385        if !self.close_requested.swap(true, Ordering::AcqRel) {
386            self.close_request.send_replace(true);
387        }
388        let mut start = self.start.lock().unwrap_or_else(|error| error.into_inner());
389        let prepared = match &mut *start {
390            StartState::Prepared(prepared) => prepared.take(),
391            StartState::Started | StartState::Finished => return,
392        };
393        drop(prepared);
394        *start = StartState::Finished;
395        self.publish_pre_activation_close();
396    }
397
398    fn finish_before_start(&self, activation: Result<(), RegistryError>) {
399        let mut start = self.start.lock().unwrap_or_else(|error| error.into_inner());
400        let prepared = match &mut *start {
401            StartState::Prepared(prepared) => prepared.take(),
402            StartState::Started | StartState::Finished => return,
403        };
404        drop(prepared);
405        *start = StartState::Finished;
406        self.activation_result.send_replace(Some(activation));
407        self.close_result.send_replace(Some(Ok(())));
408    }
409
410    fn publish_pre_activation_close(&self) {
411        self.activation_result
412            .send_replace(Some(Err(RegistryError::message(
413                self.activation_operation,
414                RegistryErrorKind::Cancelled,
415                "registry handle closed before activation",
416            ))));
417        self.close_result.send_replace(Some(Ok(())));
418    }
419
420    async fn wait_activation(&self) -> Result<(), RegistryError> {
421        wait_for_result(
422            self.activation_result.subscribe(),
423            self.activation_operation,
424            "activation worker ended without a result",
425        )
426        .await
427    }
428
429    async fn wait_close(&self) -> Result<(), RegistryError> {
430        wait_for_result(
431            self.close_result.subscribe(),
432            self.close_operation,
433            "cleanup worker ended without a result",
434        )
435        .await
436    }
437}
438
439struct ActivationWaiter {
440    lifecycle: Arc<Lifecycle>,
441    registered: bool,
442    completed: bool,
443}
444
445impl ActivationWaiter {
446    fn new(lifecycle: Arc<Lifecycle>) -> Self {
447        lifecycle.activation_waiters.fetch_add(1, Ordering::AcqRel);
448        Self {
449            lifecycle,
450            registered: true,
451            completed: false,
452        }
453    }
454
455    fn complete(&mut self) {
456        self.completed = true;
457        self.lifecycle
458            .activation_observed
459            .store(true, Ordering::Release);
460        self.release();
461    }
462
463    fn release(&mut self) {
464        if !self.registered {
465            return;
466        }
467        self.registered = false;
468        let previous = self
469            .lifecycle
470            .activation_waiters
471            .fetch_sub(1, Ordering::AcqRel);
472        debug_assert!(previous > 0);
473        if previous == 1
474            && !self.completed
475            && !self.lifecycle.activation_observed.load(Ordering::Acquire)
476        {
477            self.lifecycle.request_close();
478        }
479    }
480}
481
482impl Drop for ActivationWaiter {
483    fn drop(&mut self) {
484        self.release();
485    }
486}
487
488impl Drop for Lifecycle {
489    fn drop(&mut self) {
490        self.request_close();
491    }
492}
493
494struct LifecycleWorker {
495    activation_operation: RegistryOperation,
496    close_operation: RegistryOperation,
497    activate: Option<RegistryFuture<()>>,
498    close: Option<CloseFactory>,
499    close_request: watch::Receiver<bool>,
500    activation_result: watch::Sender<SharedResult>,
501    close_result: watch::Sender<SharedResult>,
502    activation_published: bool,
503    close_published: bool,
504}
505
506impl LifecycleWorker {
507    async fn run(mut self) {
508        let activate = self
509            .activate
510            .take()
511            .expect("activation future is present until the worker starts");
512        let activation = match AssertUnwindSafe(activate).catch_unwind().await {
513            Ok(result) => result,
514            Err(_) => Err(RegistryError::message(
515                self.activation_operation,
516                RegistryErrorKind::Internal,
517                "registry provider activation panicked",
518            )),
519        };
520        let activation = if *self.close_request.borrow() && activation.is_ok() {
521            Err(RegistryError::message(
522                self.activation_operation,
523                RegistryErrorKind::Cancelled,
524                "registry handle closed while activation was pending",
525            ))
526        } else {
527            activation
528        };
529        self.activation_result.send_replace(Some(activation));
530        self.activation_published = true;
531
532        if !*self.close_request.borrow() {
533            let _ = self.close_request.changed().await;
534        }
535        let close = self
536            .close
537            .take()
538            .expect("cleanup factory is present until close is requested");
539        let close = match catch_unwind(AssertUnwindSafe(close)) {
540            Ok(close) => match AssertUnwindSafe(close).catch_unwind().await {
541                Ok(result) => result,
542                Err(_) => Err(RegistryError::message(
543                    self.close_operation,
544                    RegistryErrorKind::Internal,
545                    "registry provider cleanup panicked",
546                )),
547            },
548            Err(_) => Err(RegistryError::message(
549                self.close_operation,
550                RegistryErrorKind::Internal,
551                "registry provider cleanup factory panicked",
552            )),
553        };
554        self.close_result.send_replace(Some(close));
555        self.close_published = true;
556    }
557}
558
559impl Drop for LifecycleWorker {
560    fn drop(&mut self) {
561        if !self.activation_published {
562            self.activation_result
563                .send_replace(Some(Err(RegistryError::message(
564                    self.activation_operation,
565                    RegistryErrorKind::Internal,
566                    "registry activation worker was aborted",
567                ))));
568        }
569        if !self.close_published {
570            self.close_result
571                .send_replace(Some(Err(RegistryError::message(
572                    self.close_operation,
573                    RegistryErrorKind::CleanupAborted,
574                    "registry cleanup worker was aborted",
575                ))));
576        }
577    }
578}
579
580async fn wait_for_result(
581    mut result: watch::Receiver<SharedResult>,
582    operation: RegistryOperation,
583    ended_message: &'static str,
584) -> Result<(), RegistryError> {
585    loop {
586        if let Some(result) = result.borrow().clone() {
587            return result;
588        }
589        result.changed().await.map_err(|_| {
590            RegistryError::message(operation, RegistryErrorKind::Internal, ended_message)
591        })?;
592    }
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598    use crate::directory::directory;
599    use std::{
600        sync::atomic::{AtomicUsize, Ordering},
601        task::{Context, Poll, Waker},
602    };
603    use tokio::sync::{Notify, oneshot};
604
605    fn close_error(message: &str) -> RegistryError {
606        RegistryError::message(
607            RegistryOperation::CloseRegistration,
608            RegistryErrorKind::Internal,
609            message,
610        )
611    }
612
613    #[tokio::test]
614    async fn prepared_handle_has_no_side_effect_before_activation() {
615        let activations = Arc::new(AtomicUsize::new(0));
616        let cleanups = Arc::new(AtomicUsize::new(0));
617        let handle = prepare_registration(
618            {
619                let activations = activations.clone();
620                async move {
621                    activations.fetch_add(1, Ordering::SeqCst);
622                    Ok(())
623                }
624            },
625            {
626                let cleanups = cleanups.clone();
627                move || async move {
628                    cleanups.fetch_add(1, Ordering::SeqCst);
629                    Ok(())
630                }
631            },
632        );
633
634        tokio::task::yield_now().await;
635        assert_eq!(activations.load(Ordering::SeqCst), 0);
636        handle.close().await.unwrap();
637        assert_eq!(activations.load(Ordering::SeqCst), 0);
638        assert_eq!(cleanups.load(Ordering::SeqCst), 0);
639        assert_eq!(
640            handle.activate().await.unwrap_err().kind(),
641            RegistryErrorKind::Cancelled
642        );
643    }
644
645    #[tokio::test]
646    async fn cancelling_last_activation_waiter_requests_late_success_cleanup() {
647        let started = Arc::new(Notify::new());
648        let (release_sender, release_receiver) = oneshot::channel();
649        let cleanups = Arc::new(AtomicUsize::new(0));
650        let cleanup_completed = Arc::new(Notify::new());
651        let handle = prepare_registration(
652            {
653                let started = started.clone();
654                async move {
655                    started.notify_one();
656                    let _ = release_receiver.await;
657                    Ok(())
658                }
659            },
660            {
661                let cleanups = cleanups.clone();
662                let cleanup_completed = cleanup_completed.clone();
663                move || async move {
664                    cleanups.fetch_add(1, Ordering::SeqCst);
665                    cleanup_completed.notify_one();
666                    Ok(())
667                }
668            },
669        );
670        let waiter = tokio::spawn({
671            let handle = handle.clone();
672            async move { handle.activate().await }
673        });
674        started.notified().await;
675        waiter.abort();
676        assert!(waiter.await.unwrap_err().is_cancelled());
677
678        release_sender.send(()).unwrap();
679        cleanup_completed.notified().await;
680        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
681        handle.close().await.unwrap();
682        assert_eq!(
683            handle.activate().await.unwrap_err().kind(),
684            RegistryErrorKind::Cancelled
685        );
686    }
687
688    #[tokio::test]
689    async fn cancelling_one_of_two_activation_waiters_keeps_the_shared_activation_alive() {
690        let started = Arc::new(Notify::new());
691        let (release_sender, release_receiver) = oneshot::channel();
692        let cleanups = Arc::new(AtomicUsize::new(0));
693        let handle = prepare_registration(
694            {
695                let started = started.clone();
696                async move {
697                    started.notify_one();
698                    let _ = release_receiver.await;
699                    Ok(())
700                }
701            },
702            {
703                let cleanups = cleanups.clone();
704                move || async move {
705                    cleanups.fetch_add(1, Ordering::SeqCst);
706                    Ok(())
707                }
708            },
709        );
710        let cancelled_waiter = handle.activate();
711        let surviving_waiter = handle.activate();
712        let cancelled = tokio::spawn(cancelled_waiter);
713        let surviving = tokio::spawn(surviving_waiter);
714        started.notified().await;
715
716        cancelled.abort();
717        assert!(cancelled.await.unwrap_err().is_cancelled());
718        release_sender.send(()).unwrap();
719        surviving.await.unwrap().unwrap();
720        assert_eq!(cleanups.load(Ordering::SeqCst), 0);
721
722        handle.close().await.unwrap();
723        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
724    }
725
726    #[tokio::test]
727    async fn concurrent_activation_and_close_share_one_provider_operation() {
728        let activations = Arc::new(AtomicUsize::new(0));
729        let cleanups = Arc::new(AtomicUsize::new(0));
730        let close_started = Arc::new(Notify::new());
731        let close_release = Arc::new(Notify::new());
732        let handle = prepare_registration(
733            {
734                let activations = activations.clone();
735                async move {
736                    activations.fetch_add(1, Ordering::SeqCst);
737                    Ok(())
738                }
739            },
740            {
741                let cleanups = cleanups.clone();
742                let close_started = close_started.clone();
743                let close_release = close_release.clone();
744                move || async move {
745                    cleanups.fetch_add(1, Ordering::SeqCst);
746                    close_started.notify_one();
747                    close_release.notified().await;
748                    Err(close_error("expected cleanup failure"))
749                }
750            },
751        );
752
753        let first_activation = tokio::spawn({
754            let handle = handle.clone();
755            async move { handle.activate().await }
756        });
757        let second_activation = tokio::spawn({
758            let handle = handle.clone();
759            async move { handle.activate().await }
760        });
761        first_activation.await.unwrap().unwrap();
762        second_activation.await.unwrap().unwrap();
763        assert_eq!(activations.load(Ordering::SeqCst), 1);
764
765        let first_close = tokio::spawn({
766            let handle = handle.clone();
767            async move { handle.close().await }
768        });
769        let second_close = tokio::spawn({
770            let handle = handle.clone();
771            async move { handle.close().await }
772        });
773        close_started.notified().await;
774        close_release.notify_waiters();
775        let first = first_close.await.unwrap().unwrap_err();
776        let second = second_close.await.unwrap().unwrap_err();
777        assert_eq!(first.to_string(), second.to_string());
778        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
779    }
780
781    #[tokio::test]
782    async fn activation_error_is_preserved_and_cleanup_runs_once() {
783        let cleanups = Arc::new(AtomicUsize::new(0));
784        let activation_error = RegistryError::message(
785            RegistryOperation::ActivateRegistration,
786            RegistryErrorKind::Unavailable,
787            "expected activation failure",
788        );
789        let handle = prepare_registration(
790            {
791                let activation_error = activation_error.clone();
792                async move { Err(activation_error) }
793            },
794            {
795                let cleanups = cleanups.clone();
796                move || async move {
797                    cleanups.fetch_add(1, Ordering::SeqCst);
798                    Ok(())
799                }
800            },
801        );
802
803        let error = handle.activate().await.unwrap_err();
804        assert_eq!(error.kind(), RegistryErrorKind::Unavailable);
805        handle.close().await.unwrap();
806        handle.close().await.unwrap();
807        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
808    }
809
810    #[tokio::test]
811    async fn activation_panic_is_isolated_and_cleanup_still_runs_once() {
812        let cleanups = Arc::new(AtomicUsize::new(0));
813        let handle = prepare_registration(
814            async {
815                panic!("expected provider activation panic");
816            },
817            {
818                let cleanups = cleanups.clone();
819                move || async move {
820                    cleanups.fetch_add(1, Ordering::SeqCst);
821                    Ok(())
822                }
823            },
824        );
825
826        let error = handle.activate().await.unwrap_err();
827        assert_eq!(error.kind(), RegistryErrorKind::Internal);
828        handle.close().await.unwrap();
829        handle.close().await.unwrap();
830        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
831    }
832
833    #[tokio::test]
834    async fn cleanup_panic_is_isolated_and_shared() {
835        let cleanups = Arc::new(AtomicUsize::new(0));
836        let handle = prepare_registration(async { Ok(()) }, {
837            let cleanups = cleanups.clone();
838            move || async move {
839                cleanups.fetch_add(1, Ordering::SeqCst);
840                panic!("expected provider cleanup panic");
841            }
842        });
843        handle.activate().await.unwrap();
844
845        let first = handle.close().await.unwrap_err();
846        let second = handle.close().await.unwrap_err();
847        assert_eq!(first.kind(), RegistryErrorKind::Internal);
848        assert_eq!(first.to_string(), second.to_string());
849        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
850    }
851
852    #[tokio::test]
853    async fn subscription_activation_returns_the_prepared_directory() {
854        let (publisher, directory) = directory();
855        let expected = directory.clone();
856        let handle = prepare_subscription(
857            directory,
858            async move {
859                publisher.publish_ready(Vec::new())?;
860                Ok(())
861            },
862            || async { Ok(()) },
863        );
864
865        let active = handle.activate().await.unwrap();
866        assert_eq!(active.snapshot().revision(), expected.snapshot().revision());
867        assert_eq!(active.snapshot().state(), expected.snapshot().state());
868        handle.close().await.unwrap();
869    }
870
871    #[tokio::test]
872    async fn last_handle_drop_only_requests_background_close() {
873        let completed = Arc::new(Notify::new());
874        let cleanups = Arc::new(AtomicUsize::new(0));
875        let handle = prepare_registration(async { Ok(()) }, {
876            let completed = completed.clone();
877            let cleanups = cleanups.clone();
878            move || async move {
879                cleanups.fetch_add(1, Ordering::SeqCst);
880                completed.notify_one();
881                Ok(())
882            }
883        });
884        handle.activate().await.unwrap();
885        drop(handle);
886
887        completed.notified().await;
888        assert_eq!(cleanups.load(Ordering::SeqCst), 1);
889    }
890
891    #[test]
892    fn close_before_activation_is_ready_without_a_runtime() {
893        let handle = prepare_registration(async { Ok(()) }, || async { Ok(()) });
894        let mut future = handle.close();
895        let waker = Waker::noop();
896        let mut context = Context::from_waker(waker);
897        assert!(matches!(
898            Pin::as_mut(&mut future).poll(&mut context),
899            Poll::Ready(Ok(()))
900        ));
901    }
902}
fusen_register/lib.rs

fusen_register/
lib.rs
