use std::any::Any;
use std::future::Future;
use std::sync::Arc;
pub trait Executor: Any + Send + Sync + 'static {
fn as_any(&self) -> &dyn Any;
fn non_transactional(&self) -> Option<Arc<dyn Executor>> {
None
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExecutorScope {
Request,
Job,
}
tokio::task_local! {
static EXECUTOR: Arc<dyn Executor>;
static EXECUTOR_SCOPE: ExecutorScope;
}
pub fn current_executor() -> Option<Arc<dyn Executor>> {
EXECUTOR.try_with(Arc::clone).ok()
}
pub fn current_executor_scope() -> Option<ExecutorScope> {
EXECUTOR_SCOPE.try_with(Clone::clone).ok()
}
pub async fn with_executor<F: Future>(executor: Arc<dyn Executor>, fut: F) -> F::Output {
EXECUTOR.scope(executor, fut).await
}
pub async fn with_request_executor<F: Future>(executor: Arc<dyn Executor>, fut: F) -> F::Output {
EXECUTOR
.scope(executor, EXECUTOR_SCOPE.scope(ExecutorScope::Request, fut))
.await
}
pub async fn with_job_executor<F: Future>(executor: Arc<dyn Executor>, fut: F) -> F::Output {
EXECUTOR
.scope(executor, EXECUTOR_SCOPE.scope(ExecutorScope::Job, fut))
.await
}
#[cfg(test)]
mod tests {
use super::*;
struct StubExecutor;
impl Executor for StubExecutor {
fn as_any(&self) -> &dyn Any {
self
}
}
fn stub() -> Arc<dyn Executor> {
Arc::new(StubExecutor)
}
#[tokio::test]
async fn no_ambient_state_outside_any_scope() {
assert!(current_executor().is_none());
assert!(current_executor_scope().is_none());
}
#[tokio::test]
async fn with_executor_installs_but_does_not_tag() {
with_executor(stub(), async {
assert!(current_executor().is_some());
assert!(current_executor_scope().is_none());
})
.await;
}
#[tokio::test]
async fn with_request_executor_tags_request() {
with_request_executor(stub(), async {
assert_eq!(current_executor_scope(), Some(ExecutorScope::Request));
assert!(current_executor().is_some());
})
.await;
}
#[tokio::test]
async fn with_job_executor_tags_job() {
with_job_executor(stub(), async {
assert_eq!(current_executor_scope(), Some(ExecutorScope::Job));
assert!(current_executor().is_some());
})
.await;
}
#[tokio::test]
async fn scope_unwinds_on_exit() {
with_request_executor(stub(), async {}).await;
assert!(current_executor().is_none());
assert!(current_executor_scope().is_none());
}
#[tokio::test]
async fn nested_scope_shadows_outer() {
with_request_executor(stub(), async {
assert_eq!(current_executor_scope(), Some(ExecutorScope::Request));
with_job_executor(stub(), async {
assert_eq!(current_executor_scope(), Some(ExecutorScope::Job));
})
.await;
assert_eq!(current_executor_scope(), Some(ExecutorScope::Request));
})
.await;
}
#[tokio::test]
async fn downcast_round_trips_the_concrete_type() {
with_request_executor(stub(), async {
let e = current_executor().expect("installed");
assert!(e.as_any().is::<StubExecutor>());
})
.await;
}
}