Skip to main content

compio_actor\cluster/
current.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use pin_project_lite::pin_project;
8
9use super::Cluster;
10
11scoped_tls::scoped_thread_local!(static CURRENT_CLUSTER: Cluster);
12
13impl Cluster {
14    /// Returns the cluster running the current actor.
15    ///
16    /// # Panics
17    ///
18    /// Panics when called outside an actor managed by a cluster.
19    pub fn current() -> Self {
20        Cluster::try_current().expect("not running in an actor cluster")
21    }
22
23    /// Try to get the cluster running the current actor.
24    pub fn try_current() -> Option<Self> {
25        CURRENT_CLUSTER
26            .is_set()
27            .then(|| CURRENT_CLUSTER.with(Clone::clone))
28    }
29
30    /// Drive the future in the scope of `self`.
31    pub(super) fn drive<F: Future>(self, future: F) -> Scope<F> {
32        Scope {
33            cluster: self,
34            future,
35        }
36    }
37}
38
39pin_project! {
40    pub(super) struct Scope<F> {
41        cluster: Cluster,
42        #[pin]
43        future: F,
44    }
45}
46
47impl<F: Future> Future for Scope<F> {
48    type Output = F::Output;
49
50    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
51        let this = self.project();
52        CURRENT_CLUSTER.set(this.cluster, || this.future.poll(cx))
53    }
54}