use serde_json::Value;
use tokio::task_local;
task_local! {
pub static INTERRUPT_CONTEXT: std::cell::RefCell<Option<InterruptContext>>;
}
#[derive(Clone, Debug)]
pub struct InterruptContext {
pub interrupt_value: Option<Value>,
pub resume_values: Vec<Value>,
pub current_index: usize,
}
impl InterruptContext {
pub fn new() -> Self {
Self {
interrupt_value: None,
resume_values: Vec::new(),
current_index: 0,
}
}
pub fn with_resume_value(value: Value) -> Self {
Self {
interrupt_value: None,
resume_values: vec![value],
current_index: 0,
}
}
pub fn with_resume_values(values: Vec<Value>) -> Self {
Self {
interrupt_value: None,
resume_values: values,
current_index: 0,
}
}
pub fn has_interrupt(&self) -> bool {
self.interrupt_value.is_some()
}
pub fn interrupt_value(&self) -> Option<&Value> {
self.interrupt_value.as_ref()
}
pub fn reset(&mut self) {
self.interrupt_value = None;
self.current_index = 0;
}
}
impl Default for InterruptContext {
fn default() -> Self {
Self::new()
}
}
pub async fn set_interrupt_context<F, R>(context: InterruptContext, f: F) -> R
where
F: std::future::Future<Output = R>,
{
INTERRUPT_CONTEXT
.scope(std::cell::RefCell::new(Some(context)), f)
.await
}
pub fn get_interrupt_value() -> Option<Value> {
INTERRUPT_CONTEXT
.try_with(|ctx| {
ctx.borrow()
.as_ref()
.and_then(|c| c.interrupt_value.clone())
})
.ok()
.flatten()
}
pub fn has_interrupt() -> bool {
INTERRUPT_CONTEXT
.try_with(|ctx| {
ctx.borrow()
.as_ref()
.map(|c| c.has_interrupt())
.unwrap_or(false)
})
.ok()
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_interrupt_context() {
let ctx = InterruptContext::with_resume_value(serde_json::json!("resume"));
set_interrupt_context(ctx, async {
let value = get_interrupt_value();
assert_eq!(value, None);
})
.await;
}
}