use crate::ctx::canceller::Canceller; use crate::ctx::reason::Reason; use crate::sql::value::Value; use std::borrow::Cow; use std::collections::HashMap; use std::fmt; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; impl<'a> From for Cow<'a, Value> { fn from(v: Value) -> Cow<'a, Value> { Cow::Owned(v) } } impl<'a> From<&'a Value> for Cow<'a, Value> { fn from(v: &'a Value) -> Cow<'a, Value> { Cow::Borrowed(v) } } pub struct Context<'a> { // An optional parent context. parent: Option<&'a Context<'a>>, // An optional deadline. deadline: Option, // Whether or not this context is cancelled. cancelled: Arc, // A collection of read only values stored in this context. values: Option>>, } impl<'a> Default for Context<'a> { fn default() -> Self { Context::background() } } impl<'a> Context<'a> { // Create an empty background context. pub fn background() -> Self { Context { values: None, parent: None, deadline: None, cancelled: Arc::new(AtomicBool::new(false)), } } // Create a new child from a frozen context. pub fn new(parent: &'a Context) -> Self { Context { values: None, parent: Some(parent), deadline: parent.deadline, cancelled: Arc::new(AtomicBool::new(false)), } } // Add cancellation to the context. The value that is returned will cancel // the context and it's children once called. pub fn add_cancel(&mut self) -> Canceller { let cancelled = self.cancelled.clone(); Canceller::new(cancelled) } // Add a deadline to the context. If the current deadline is sooner than // the provided deadline, this method does nothing. pub fn add_deadline(&mut self, deadline: Instant) { match self.deadline { Some(current) if current < deadline => (), _ => self.deadline = Some(deadline), } } // Add a timeout to the context. If the current timeout is sooner than // the provided timeout, this method does nothing. pub fn add_timeout(&mut self, timeout: Duration) { self.add_deadline(Instant::now() + timeout) } // Add a value to the context. It overwrites any previously set values // with the same key. pub fn add_value(&mut self, key: String, value: V) where V: Into>, { if let Some(ref mut values) = self.values { values.insert(key, value.into()); } else { self.values = Some(HashMap::new()); self.add_value(key, value); } } // Get the deadline for this operation, if any. This is useful for // checking if a long job should be started or not. pub fn deadline(&self) -> Option { self.deadline } // Check if the context is done. If it returns `None` the operation may // proceed, otherwise the operation should be stopped. pub fn done(&self) -> Option { match self.deadline { Some(deadline) if deadline <= Instant::now() => Some(Reason::Timedout), // TODO: see if we can relax the ordering. _ if self.cancelled.load(Ordering::SeqCst) => Some(Reason::Canceled), _ => match self.parent { Some(ctx) => ctx.done(), _ => None, }, } } // Check if the context is ok to continue. pub fn is_ok(&self) -> bool { self.done().is_none() } // Check if the context is not ok to continue. pub fn is_err(&self) -> bool { self.done().is_some() } // Check if the context is not ok to continue. pub fn is_done(&self) -> bool { self.done().is_some() } // Check if the context is not ok to continue, because it timed out. pub fn is_timedout(&self) -> bool { matches!(self.done(), Some(Reason::Timedout)) } // Check if the context is not ok to continue, because it was cancelled. pub fn is_cancelled(&self) -> bool { matches!(self.done(), Some(Reason::Canceled)) } // Check if the status of the context. This will return a Result, with an Ok // if the operation may proceed, and an Error if it should be stopped. pub fn check(&self) -> Result<(), Reason> { match self.done() { Some(reason) => Err(reason), None => Ok(()), } } // Get a value from the context. If no value is stored under the // provided key, then this will return None. pub fn value(&self, key: &str) -> Option<&Value> { match &self.values { Some(v) => match v.get(key) { Some(v) => match v { Cow::Borrowed(v) => Some(*v), Cow::Owned(v) => Some(v), }, None => match self.parent { Some(p) => p.value(key), _ => None, }, }, None => match self.parent { Some(p) => p.value(key), _ => None, }, } } } impl<'a> fmt::Debug for Context<'a> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("Context") .field("parent", &self.parent) .field("deadline", &self.deadline) .field("cancelled", &self.cancelled) .field("values", &self.values.as_ref().map(|_| "values")) .finish() } }