operation_queue/

line_token.rs

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Helpers for synchronizing operations (e.g. error handling) across futures.
//!
//! This module revolves around the [`Line`] struct, which is an asynchronous
//! flow control structure that behaves a bit like a mutex, with the exception
//! that consumers waiting for the [`Line`] to be released do not subsequently
//! lock it.
//!
//! The design of a [`Line`] is inspired from the one of a [one-track railway
//! line](https://en.wikipedia.org/wiki/Token_(railway_signalling)). To avoid
//! collisions, conductors must acquire a token at the entrance to the line that
//! ensures they're the only one on it. If the token is being held, traffic
//! around this line stops until it's released again.
//!
//! Similarly, in a context with multiple parallel [`Future`]s, it might be
//! necessary to ensure only one takes care of a given operation. For example,
//! if multiple requests are being performed against the same service, and one
//! of them hits an authentication error, it is likely the others will as well.
//! In this case, it is preferrable to only let one future handle the error than
//! let every request re-authenticate independently (in this example,
//! credentials are the same across requests, and multiple simultaneous
//! authentication attempts might cause issues with complex flows).
//!
//! Each future holds a shared on a [`Line`] (e.g. wrapped in an [`Rc`] or an
//! [`Arc`]). Whenever a future needs to perform an operation that should only
//! be performed once at a time, it attempts to acquire the line's token with
//! [`Line::try_acquire_token`]. This function returns an enum
//! ([`AcquireOutcome`]) describing one of two cases:
//!
//! * The line's token is available and has been acquired, and the future can
//!   start performing the operation immediately. It is granted the line's
//!   [`Token`], which it must hold in scope for the duration of the operation,
//!   as dropping it releases the line.
//! * The line's token has already been acquired by another future, in which
//!   case the future must wait for the line to become available again. When the
//!   line becomes available again, the future does not need to acquire another
//!   token, as another future should have taken care of performing the
//!   operation.
//!
//! [`OperationQueue`]: crate::operation_queue::OperationQueue
//! [`Future`]: std::future::Future
//! [`Rc`]: std::rc::Rc
//! [`Arc`]: std::sync::Arc

use async_lock::Mutex;
use futures::{FutureExt, future::Shared};
use oneshot::{Receiver, Sender};

/// A oneshot channel used internally by a [`Line`] that's been acquired to
/// communicate that the token has been dropped and the line was released.
///
/// The channel's [`Receiver`] is wrapped in a [`Shared`] that can be cloned
/// when a new consumer tries and fails to acquire a token for the line.
struct ReleaseChannel {
    sender: Sender<()>,
    receiver: Shared<Receiver<()>>,
}

/// The current status of a line.
pub enum LineStatus {
    /// The line is free.
    Free,

    /// The line is busy, and the nested future can be used to wait till the
    /// line becomes free again.
    Busy(Shared<Receiver<()>>),
}

/// A [`Line`] from which a [`Token`] can be acquired.
///
/// # Thread safety
///
/// `Line` is thread-safe and can be shared across threads via [`Arc`]:
///
/// ```rust
/// use std::sync::Arc;
/// use operation_queue::line_token::Line;
///
/// fn assert_send_sync<T: Send + Sync>() {}
/// assert_send_sync::<Line>();
/// ```
///
/// [`Arc`]: std::sync::Arc
#[derive(Default)]
pub struct Line {
    channel: Mutex<Option<ReleaseChannel>>,
}

impl Line {
    /// Instantiates a new [`Line`].
    pub fn new() -> Line {
        Line {
            channel: Default::default(),
        }
    }

    /// Retrieve the current status of the [`Line`], i.e. whether it's currently
    /// free or not.
    ///
    /// Unlike [`try_acquire_token`], this method does not attempt to lock the
    /// line if it's free, and does not generate a [`Token`]. If the line is
    /// busy, it returns a future the consumer can use to wait till the line
    /// becomes free again.
    ///
    /// [`try_acquire_token`]: Self::try_acquire_token
    pub async fn status(&self) -> LineStatus {
        match self.channel.lock().await.as_ref() {
            Some(channel) => LineStatus::Busy(channel.receiver.clone()),
            None => LineStatus::Free,
        }
    }

    /// Attempts to acquire a [`Token`] for this line.
    ///
    /// The [`Token`] automatically releases the line upon leaving the current
    /// scope and getting dropped.
    ///
    /// If a [`Token`] has already been acquired for this line, a future to
    /// `await` is returned instead. It resolves when the current token holder
    /// has finished handling the current error and releases the line.
    pub async fn try_acquire_token<'l>(&'l self) -> AcquireOutcome<'l> {
        let mut channel = self.channel.lock().await;

        if let Some(channel) = channel.as_ref() {
            // Since the oneshot `Receiver` is wrapped in a `Shared`, cloning it
            // will return a new handle on the `Shared` which will resolve at
            // the same time as the others.
            return AcquireOutcome::Failure(channel.receiver.clone());
        }

        // The line is currently available, create a new channel and give the
        // consumer their token.
        let (sender, receiver) = oneshot::channel();
        *channel = Some(ReleaseChannel {
            sender,
            receiver: receiver.shared(),
        });

        AcquireOutcome::Success(Token { line: self })
    }

    /// Releases the line, and resolves the [`Shared`] future other consumers
    /// might be awaiting.
    pub(self) fn release(&self) {
        // "Take" the channel out of the `Mutex`; on top of letting us access
        // its content, we're also making sure that even if something bad
        // happens then the line can be acquired again.
        match self.channel.lock_blocking().take() {
            Some(channel) => match channel.sender.send(()) {
                Ok(_) => (),
                Err(_) => log::error!("trying to release using a closed channel"),
            },
            None => log::error!("trying to release before acquiring"),
        };
    }
}

/// The outcome from trying to acquire a [`Token`] for a [`Line`].
#[must_use = "if the token is unused the line will immediately release again"]
pub enum AcquireOutcome<'ao> {
    /// The line could be acquired and returned a token to hold on to.
    ///
    /// The token must remain in scope, as it will release the line when
    /// dropped.
    Success(Token<'ao>),

    /// The line could not be acquired as another consumer is holding a token
    /// for it.
    ///
    /// This variant includes a [`Shared`] future that resolves when the current
    /// token holder drops it and releases the line.
    Failure(Shared<Receiver<()>>),
}

impl<'ao> AcquireOutcome<'ao> {
    /// Returns the [`AcquireOutcome`] if it's a success, otherwise returns a
    /// success with the provided token if it's not [`None`].
    ///
    /// If the current [`AcquireOutcome`] is a failure, and the provided token
    /// is [`None`], the failure is returned.
    ///
    /// # Design considerations
    ///
    /// One way to make this method more straightforward could have been to make
    /// `token` be a [`Token`], not an [`Option`], but the current signature was
    /// picked to simplify the consumers (which store the token, if any, in an
    /// [`Option`]).
    pub fn or_token(self, token: Option<Token<'ao>>) -> Self {
        match self {
            AcquireOutcome::Success(_) => self,
            AcquireOutcome::Failure(_) => match token {
                Some(token) => AcquireOutcome::Success(token),
                None => self,
            },
        }
    }
}

/// A token that symbolizes the current consumer holds exclusive access to the
/// corresponding [`Line`].
///
/// The [`Line`] is automatically released when this token goes out of scope and
/// is dropped.
#[must_use = "if unused the line will immediately release again"]
pub struct Token<'t> {
    line: &'t Line,
}

impl Drop for Token<'_> {
    fn drop(&mut self) {
        self.line.release();
    }
}

#[cfg(test)]
mod tests {
    use tokio::time::Duration;

    use super::*;

    async fn get_token(line: &Line) -> Token<'_> {
        match line.try_acquire_token().await {
            AcquireOutcome::Success(token) => token,
            AcquireOutcome::Failure(_) => panic!("expected a token from try_acquire_token()"),
        }
    }

    #[tokio::test(flavor = "current_thread")]
    async fn acquire_token() {
        let line = Line::new();

        let _token = get_token(&line).await;

        match line.try_acquire_token().await {
            AcquireOutcome::Success(_) => {
                panic!("should not be able to acquire the line while the token is in scope")
            }
            AcquireOutcome::Failure(_) => (),
        }
    }

    #[tokio::test(flavor = "current_thread")]
    async fn token_out_of_scope() {
        let line = Line::new();

        {
            let _token = get_token(&line).await;

            match line.try_acquire_token().await {
                AcquireOutcome::Success(_) => {
                    panic!("should not be able to acquire the line while the token is in scope")
                }
                AcquireOutcome::Failure(_) => (),
            }
        }

        match line.try_acquire_token().await {
            AcquireOutcome::Success(_) => (),
            AcquireOutcome::Failure(_) => {
                panic!("expected a token now that the previous token has been dropped")
            }
        }
    }

    #[tokio::test(flavor = "current_thread")]
    async fn or_token() {
        let line = Line::new();

        let token = get_token(&line).await;

        match line.try_acquire_token().await.or_token(Some(token)) {
            AcquireOutcome::Success(_) => (),
            AcquireOutcome::Failure(_) => panic!("we should have kept our token"),
        }
    }

    #[tokio::test(flavor = "current_thread")]
    async fn line_release_on_drop() {
        let line = Line::new();

        // A mutable variable that will act as the test's success flag and will
        // only be true if it succeeds.
        let mut success = false;

        // Acquire the line's token, sleep for a bit (10ms) and then drop it.
        // The reason we sleep here is to give some time to `wait_for_line` to
        // try (and fail) to acquire the line's token before we drop it.
        async fn acquire_sleep_and_drop(line: &Line) {
            let _token = get_token(&line).await;
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        // Try (and fail) to acquire the token, then wait for the line to become
        // available again. This function sets the success flag.
        async fn wait_for_line(line: &Line, success: &mut bool) {
            let shared = match line.try_acquire_token().await {
                AcquireOutcome::Success(_) => {
                    panic!("should not be able to acquire the line while the token is in scope")
                }
                AcquireOutcome::Failure(shared) => shared,
            };

            shared.await.unwrap();
            *success = true;
        }

        // Run both futures in parallel. `biased;` ensures the futures are
        // polled in order (meaning `acquire_sleep_and_drop` is run first).
        tokio::join! {
            biased;
            acquire_sleep_and_drop(&line),
            wait_for_line(&line, &mut success),
        };

        assert!(success)
    }

    /// Tests the [`Line::status`] method.
    ///
    /// This test behaves a lot like [`line_release_on_drop`], except it waits
    /// for the line to be freed via [`Line::status`] rather than
    /// [`Line::try_acquire_token`].
    #[tokio::test(flavor = "current_thread")]
    async fn line_status() {
        let line = Line::new();
        assert!(matches!(line.status().await, LineStatus::Free));

        // A mutable variable that will act as the test's success flag and will
        // only be true if it succeeds.
        let mut success = false;

        // Acquire the line's token, sleep for a bit (10ms) and then drop it.
        // The reason we sleep here is to give some time to `wait_for_line` to
        // try (and fail) to acquire the line's token before we drop it.
        async fn acquire_sleep_and_drop(line: &Line) {
            let _token = get_token(&line).await;
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        // Check that the line is busy, then wait for the line to become
        // available again. This function sets the success flag.
        async fn wait_for_line(line: &Line, success: &mut bool) {
            let shared = match line.status().await {
                LineStatus::Free => {
                    panic!("line should be busy")
                }
                LineStatus::Busy(shared) => shared,
            };

            shared.await.unwrap();
            *success = true;
        }

        // Run both futures in parallel. `biased;` ensures the futures are
        // polled in order (meaning `acquire_sleep_and_drop` is run first).
        tokio::join! {
            biased;
            acquire_sleep_and_drop(&line),
            wait_for_line(&line, &mut success),
        };

        assert!(matches!(line.status().await, LineStatus::Free));
        assert!(success)
    }
}