procsem/

lib.rs

//! A simple process semaphore.
//!
//! (Note: The word _process_ should be read as a sequence of operations,
//! rather than an operating system process).
//!
//! The [`ProcSem`] is intended to allow mutual exclusion of a chain of
//! operations that may span over several threads/tasks.
//!
//! This is much like a `Mutex`, but it differs in that it holds no generic
//! parameter and the `ProcCtx` (the equivalent of `Mutex`'s `MutexGuard`) is
//! `Send`, because it is explicitly meant to be passed around between
//! threads/tasks.  It supports blocking, nonblocking and async lock
//! acquisition.
//!
//! # Progress reporting
//! When a `ProcSem` is created, it can optionally be handed an object that
//! implements [`StateReporter`].  If a `ProcSem` has a `StateReporter`
//! associated with it, then an acquired `ProcCtx` can use
//! [`ProcCtx::action()`], [`ProcCtx::progress()`] to pass progress information
//! to the `StateReporter` implementor.
//!
//! Once a `ProcCtx` is about to terminate, [`ProcCtx::end()`] can be used to
//! signal a final message that can be retrieved by the `ProcSem` object.

use std::{
  future::Future,
  num::NonZeroUsize,
  pin::Pin,
  sync::{
    atomic::{AtomicUsize, Ordering},
    Arc
  },
  task::{Context, Poll, Waker}
};

use parking_lot::{Condvar, Mutex};

use indexmap::IndexMap;


/// Interface used to implement objects that are reponsible for sending process
/// reports.
#[allow(unused_variables)]
pub trait StateReporter: Send + Sync {
  /// Called when a [`ProcCtx`] has been acquired.
  fn begin(&self) {}

  /// Called when [`ProcCtx::action()`] has been called to set a current
  /// action.
  fn action(&self, s: &str) {}

  /// The [`ProcCtx`] has progress has been made.
  fn progress(&self, percent: Option<u8>) {}

  /// Called when the [`ProcCtx`] is being dropped.
  ///
  /// This method will be called while the process ownership context still
  /// exists.  If the `ProcCtx` is used as a resource lock, the application
  /// can rely on mutual exclusive access to the resource within this callback.
  fn end(&self, msg: Option<String>) {}
}


/// Shared data that must be protected behind mutually exclusive access.
#[derive(Default)]
struct Inner {
  /// Flag denoting whether the semaphore is currently owned or not.
  ///
  /// If this is `true` it means there's a `ProcCtx` object.  If this is
  /// `false` there's no `ProcCtx` object for this semaphore.
  owned: bool,

  /// Collection of async task wakers waiting to acquire this lock.
  wakers: IndexMap<usize, Waker>,

  /// The last final message.
  endmsg: Option<String>
}

#[derive(Default)]
struct Shared {
  signal: Condvar,
  inner: Mutex<Inner>,
  idgen: AtomicUsize,
  reporter: Option<Box<dyn StateReporter>>
}


/// Representation of a process that may be active or inactive.
#[repr(transparent)]
#[derive(Default)]
pub struct ProcSem(Arc<Shared>);

impl ProcSem {
  /// Create a new process semaphore object.
  #[must_use]
  pub fn new() -> Self {
    Self::default()
  }

  /// Create a new process semaphore with a reporter object.
  pub fn with_reporter(reporter: impl StateReporter + 'static) -> Self {
    Self(Arc::new(Shared {
      reporter: Some(Box::new(reporter)),
      ..Default::default()
    }))
  }

  /// Return a `Future` that will yield a [`ProcCtx`] when possible.
  ///
  /// On successful acquire the internal endmsg will be cleared.
  #[must_use]
  pub fn acquire_async(&self) -> AcquireFuture {
    AcquireFuture {
      shared: Arc::clone(&self.0),
      id: None
    }
  }

  /// Acquire a [`ProcCtx`], blocking until successful.
  ///
  /// On successful acquire the internal endmsg will be cleared.
  #[must_use]
  #[allow(clippy::significant_drop_tightening)]
  pub fn acquire_blocking(&self) -> ProcCtx {
    let mut g = self.0.inner.lock();
    loop {
      if g.owned {
        // There's a Lock for this ProcSem.
        // Wait until the condvar is triggered so we can check again.
        self.0.signal.wait(&mut g);
      } else {
        // Lock and return a Lock object
        g.owned = true;
        g.endmsg = None;
        if let Some(ref r) = self.0.reporter {
          r.begin();
        }
        break ProcCtx(Arc::clone(&self.0));
      }
    }
  }

  /// Attempt to acquire [`ProcCtx`], returning immediately.
  ///
  /// Returns `Some(ProcCtx)` if the lock was acquired successfully.  Returns
  /// `None` if the lock is already owned by another lock.
  ///
  /// On successful acquire the internal endmsg will be cleared.
  #[must_use]
  #[allow(clippy::significant_drop_tightening)]
  pub fn try_acquire(&self) -> Option<ProcCtx> {
    let mut g = self.0.inner.lock();
    if g.owned {
      // Lock already acquired
      None
    } else {
      g.owned = true;
      g.endmsg = None;
      if let Some(ref r) = self.0.reporter {
        r.begin();
      }
      Some(ProcCtx(Arc::clone(&self.0)))
    }
  }

  /// Get a clone of the internal end message.
  #[must_use]
  pub fn endmsg(&self) -> Option<String> {
    self.0.inner.lock().endmsg.clone()
  }
}


/// Future used to wait to acquire a [`ProcCtx`] in an async environment.
pub struct AcquireFuture {
  shared: Arc<Shared>,

  id: Option<NonZeroUsize>
}

impl Future for AcquireFuture {
  type Output = ProcCtx;

  fn poll(
    mut self: Pin<&mut Self>,
    ctx: &mut Context<'_>
  ) -> Poll<Self::Output> {
    let mut inner = self.shared.inner.lock();
    if inner.owned {
      //
      // Must wait for lock to be released
      //

      // Generate a unique identifier for this waker
      let id = loop {
        let id = self.shared.idgen.fetch_add(1, Ordering::SeqCst);
        // Make sure id is non-zero and unique
        if id == 0 || inner.wakers.contains_key(&id) {
          continue;
        }
        break id;
      };
      inner.wakers.insert(id, ctx.waker().clone());
      drop(inner);

      // SAFETY: The loop above will guarantee that the id is non-zero
      self.id = Some(unsafe { NonZeroUsize::new_unchecked(id) });

      Poll::Pending
    } else {
      //
      // Mark lock as owned and return a Lock object.
      //
      inner.owned = true;
      inner.endmsg = None;
      if let Some(ref r) = self.shared.reporter {
        r.begin();
      }

      drop(inner);
      Poll::Ready(ProcCtx(Arc::clone(&self.shared)))
    }
  }
}

impl Drop for AcquireFuture {
  /// If this future has generated a `Waker`, then remove it.
  fn drop(&mut self) {
    if let Some(id) = self.id {
      let mut inner = self.shared.inner.lock();
      // Remove this future's waker
      let _ = inner.wakers.swap_remove(&id.get());
    }
  }
}


/// Object that denotes ownership of a [`ProcSem`] semaphore.
///
/// When the process is complete this object must be dropped so the semaphore
/// can be acquired again.
#[repr(transparent)]
pub struct ProcCtx(Arc<Shared>);


impl ProcCtx {
  /// Provide a textual describe the current stage of the process.
  pub fn action(&self, text: &str) {
    if let Some(ref r) = self.0.reporter {
      r.action(text);
    }
  }

  /// Set the progress (in percent).
  pub fn progress(&self, percent: Option<u8>) {
    if let Some(ref r) = self.0.reporter {
      r.progress(percent);
    }
  }

  /// Set an end-state message.
  ///
  /// The purpose of the "end message" is to be able to leave a final "result"
  /// message that can be read by the `ProcSem`.
  ///
  /// Calling this function will not cause the [`StateReporter::end()`] to be
  /// called.  However, if an end message has been set, a clone if it will be
  /// passed to the `StateReporter::end()` callback.
  ///
  /// The message will be stored in the internal [`ProcSem`] storage
  /// and can be retrieved using [`ProcSem::endmsg()`].
  pub fn end(&self, text: &str) {
    let mut g = self.0.inner.lock();
    g.endmsg = Some(text.to_string());
  }
}

impl Drop for ProcCtx {
  #[allow(clippy::significant_drop_tightening)]
  fn drop(&mut self) {
    if let Some(ref r) = self.0.reporter {
      let g = self.0.inner.lock();
      let endmsg = g.endmsg.clone();
      drop(g);
      r.end(endmsg);
    }

    let mut g = self.0.inner.lock();
    g.owned = false;

    // Wake up all waiting futures.
    //
    // In an ideal world only one would be woken up, but there are edge cases
    // when using multithreaded executors when an await is performed in a
    // select!{} in which a doomed future might be picked.  In such a case the
    // "wake" would be wasted on a future that will not be able to process the
    // wake request.
    //
    // To avoid this, _all_ wakers are woken up and they'll have to reenter
    // pending state if applicable.
    for (_, w) in g.wakers.drain(..) {
      w.wake();
    }

    // Wake a single blocking thread (blocked threads do not suffer from the
    // same issue).
    self.0.signal.notify_one();
  }
}

// vim: set ft=rust et sw=2 ts=2 sts=2 cinoptions=2 tw=79 :