picodata_plugin/

background.rs

//! Background container for long live jobs.
//! Background container guarantees job liveness across plugin life cycle.

use crate::internal::ffi;
#[allow(unused_imports)]
use crate::plugin::interface::PicoContext;
use crate::plugin::interface::ServiceId;
use crate::util::tarantool_error_to_box_error;
use std::cell::Cell;
use std::rc::Rc;
use std::time::Duration;
use tarantool::error::BoxError;
use tarantool::error::TarantoolErrorCode;
use tarantool::fiber;
use tarantool::fiber::FiberId;
use tarantool::fiber::{Channel, RecvError};
use tarantool::util::IntoClones;

/// Same as [`PicoContext::register_job`].
#[track_caller]
pub fn register_job<F>(service_id: &ServiceId, job: F) -> Result<(), BoxError>
where
    F: FnOnce(CancellationToken) + 'static,
{
    let loc = std::panic::Location::caller();
    let tag = format!("{}:{}", loc.file(), loc.line());

    register_tagged_job(service_id, job, &tag)
}

/// Same as [`PicoContext::register_tagged_job`].
pub fn register_tagged_job<F>(service_id: &ServiceId, job: F, tag: &str) -> Result<(), BoxError>
where
    F: FnOnce(CancellationToken) + 'static,
{
    let (token, handle) = CancellationToken::new();
    let finish_event = Rc::clone(&handle.finish_event);

    let fiber_id = fiber::Builder::new()
        .name(tag)
        .func(move || {
            job(token);
            finish_event.signal();
        })
        .start_non_joinable()
        .map_err(tarantool_error_to_box_error)?;

    let plugin = &service_id.plugin;
    let service = &service_id.service;
    let version = &service_id.version;

    let token = FfiBackgroundJobCancellationToken::new(
        fiber_id,
        handle.cancel_channel,
        handle.finish_event,
    );
    register_background_job_cancellation_token(plugin, service, version, tag, token)?;

    Ok(())
}

fn register_background_job_cancellation_token(
    plugin: &str,
    service: &str,
    version: &str,
    job_tag: &str,
    token: FfiBackgroundJobCancellationToken,
) -> Result<(), BoxError> {
    // SAFETY: safe as long as picodata version is compatible
    let rc = unsafe {
        ffi::pico_ffi_background_register_job_cancellation_token(
            plugin.into(),
            service.into(),
            version.into(),
            job_tag.into(),
            token,
        )
    };

    if rc != 0 {
        return Err(BoxError::last());
    }

    Ok(())
}

////////////////////////////////////////////////////////////////////////////////
// cancel_jobs_by_tag
////////////////////////////////////////////////////////////////////////////////

/// Outcome of the request to cancel a set of background jobs.
#[derive(Clone, Copy, Default, Debug, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct JobCancellationResult {
    /// Attempted to cancel this many background jobs.
    pub n_total: u32,
    /// This many jobs didn't finish in time.
    pub n_timeouts: u32,
}

impl JobCancellationResult {
    #[inline(always)]
    pub fn new(n_total: u32, n_timeouts: u32) -> Self {
        Self {
            n_total,
            n_timeouts,
        }
    }
}

/// Same as [`PicoContext::cancel_tagged_jobs`].
#[inline(always)]
pub fn cancel_jobs_by_tag(
    service_id: &ServiceId,
    job_tag: &str,
    timeout: Duration,
) -> Result<JobCancellationResult, BoxError> {
    cancel_background_jobs_by_tag_inner(
        &service_id.plugin,
        &service_id.service,
        &service_id.version,
        job_tag,
        timeout,
    )
}

/// Same as [`PicoContext::cancel_tagged_jobs`].
pub fn cancel_background_jobs_by_tag_inner(
    plugin: &str,
    service: &str,
    version: &str,
    job_tag: &str,
    timeout: Duration,
) -> Result<JobCancellationResult, BoxError> {
    let mut result = JobCancellationResult::default();
    // SAFETY: safe as long as picodata version is compatible
    let rc = unsafe {
        ffi::pico_ffi_background_cancel_jobs_by_tag(
            plugin.into(),
            service.into(),
            version.into(),
            job_tag.into(),
            timeout.as_secs_f64(),
            &mut result,
        )
    };

    if rc != 0 {
        return Err(BoxError::last());
    }

    Ok(result)
}

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Some of jobs are not fully completed, expected: {0}, completed: {1}")]
    PartialCompleted(usize, usize),
    #[error("timeout")]
    CancellationTimeout,
}

////////////////////////////////////////////////////////////////////////////////
// CancellationToken
////////////////////////////////////////////////////////////////////////////////

/// A token which can be used to signal a cancellation request to the job.
#[non_exhaustive]
#[derive(Debug)]
pub struct CancellationToken {
    cancel_channel: Channel<()>,
}

impl CancellationToken {
    /// Create a cancellation token and cancellation token handle pair.
    /// User should use cancellation token for graceful shutdown their job.
    /// Cancellation token handle used by `picodata` for sending cancel signal to a user job.
    pub fn new() -> (CancellationToken, CancellationTokenHandle) {
        let (cancel_tx, cancel_rx) = Channel::new(1).into_clones();
        (
            CancellationToken {
                cancel_channel: cancel_rx,
            },
            CancellationTokenHandle {
                cancel_channel: cancel_tx,
                finish_event: Rc::new(OnceEvent::new()),
            },
        )
    }

    /// Wait until token are canceled.
    ///
    /// # Arguments
    ///
    /// * `timeout`: if cancellation does not occur within the specified timeout - return error
    ///
    /// # Errors
    ///
    /// Return `Ok` if token is canceled or [`Error::CancellationTimeout`] if timeout is reached.
    pub fn wait_timeout(&self, timeout: Duration) -> Result<(), Error> {
        match self.cancel_channel.recv_timeout(timeout) {
            Err(RecvError::Timeout) => Err(Error::CancellationTimeout),
            _ => Ok(()),
        }
    }
}

#[derive(Debug)]
pub struct CancellationTokenHandle {
    cancel_channel: Channel<()>,
    finish_event: Rc<OnceEvent>,
}

impl CancellationTokenHandle {
    /// Cancel related job and return a finish event.
    /// Caller should wait on the finish event to make sure the job
    /// is completed successfully (graceful shutdown occurred).
    pub fn cancel(self) -> Rc<OnceEvent> {
        let Self {
            cancel_channel,
            finish_event,
        } = self;
        _ = cancel_channel.send(());
        finish_event
    }
}

////////////////////////////////////////////////////////////////////////////////
// set_background_jobs_shutdown_timeout
////////////////////////////////////////////////////////////////////////////////

/// In case when jobs were canceled by `picodata` use this function to determine
/// a shutdown timeout - time duration that `picodata` uses to ensure that all
/// jobs gracefully end.
///
/// By default, 5-second timeout are used.
///
/// Consider using [`PicoContext::set_jobs_shutdown_timeout`] instead
pub fn set_jobs_shutdown_timeout(plugin: &str, service: &str, version: &str, timeout: Duration) {
    // SAFETY: safe as long as picodata version is compatible
    let rc = unsafe {
        ffi::pico_ffi_background_set_jobs_shutdown_timeout(
            plugin.into(),
            service.into(),
            version.into(),
            timeout.as_secs_f64(),
        )
    };
    debug_assert!(
        rc == 0,
        "return code is only for future compatibility at the moment"
    );
}

////////////////////////////////////////////////////////////////////////////////
// OnceEvent
////////////////////////////////////////////////////////////////////////////////

/// A one-time event that can be signaled once and waited on by multiple fibers.
/// Late waiters (after signal) return immediately.
#[derive(Debug)]
pub struct OnceEvent {
    finished: Cell<bool>,
    cond: fiber::Cond,
}

impl OnceEvent {
    fn new() -> Self {
        Self {
            finished: Cell::new(false),
            cond: fiber::Cond::new(),
        }
    }

    /// Signal that the job has finished. Wakes all waiting fibers.
    pub fn signal(&self) {
        self.finished.set(true);
        self.cond.broadcast();
    }

    #[inline]
    pub fn is_finished(&self) -> bool {
        self.finished.get()
    }

    pub fn wait_timeout(&self, timeout: Duration) {
        self.cond.wait_timeout(timeout);
    }
}

////////////////////////////////////////////////////////////////////////////////
// CancellationCallbackState
////////////////////////////////////////////////////////////////////////////////

#[derive(Debug)]
pub struct CancellationCallbackState {
    cancel_channel: Channel<()>,
    finish_event: Rc<OnceEvent>,
    status: Cell<CancellationCallbackStatus>,
}

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
pub enum CancellationCallbackStatus {
    #[default]
    Initial = 0,
    JobCancelled = 1,
    JobFinished = 2,
}

impl CancellationCallbackState {
    fn new(cancel_channel: Channel<()>, finish_event: Rc<OnceEvent>) -> Self {
        Self {
            cancel_channel,
            finish_event,
            status: Cell::new(CancellationCallbackStatus::Initial),
        }
    }

    fn cancellation_callback(&self, action: u64, timeout: Duration) -> Result<(), BoxError> {
        use CancellationCallbackStatus::*;
        let next_status = action;

        if next_status == JobCancelled as u64 {
            // Cancellation is idempotent.
            if self.status.get() == JobCancelled || self.status.get() == JobFinished {
                return Ok(());
            }
            debug_assert_eq!(self.status.get(), Initial);
            _ = self.cancel_channel.send(());

            self.status.set(JobCancelled);
        } else if next_status == JobFinished as u64 {
            if self.status.get() != JobCancelled && self.status.get() != JobFinished {
                return Err(BoxError::new(
                    TarantoolErrorCode::IllegalParams,
                    "cannot wait for job that wasn't cancelled",
                ));
            }

            let deadline = fiber::clock().saturating_add(timeout);

            loop {
                if self.finish_event.is_finished() {
                    self.status.set(JobFinished);
                    return Ok(());
                }

                let remaining = deadline.duration_since(fiber::clock());
                if remaining.is_zero() {
                    return Err(BoxError::new(TarantoolErrorCode::Timeout, "timeout"));
                }

                self.finish_event.wait_timeout(remaining);
            }
        } else {
            return Err(BoxError::new(
                TarantoolErrorCode::IllegalParams,
                format!("unexpected action: {action}"),
            ));
        }

        Ok(())
    }
}

////////////////////////////////////////////////////////////////////////////////
// ffi wrappers
////////////////////////////////////////////////////////////////////////////////

/// **For internal use**.
#[repr(C)]
pub struct FfiBackgroundJobCancellationToken {
    /// This is just a way to do arbitrarily complex things across ABI boundary
    /// without introducing a large amount of FFI-safe wrappers. This will
    /// always be [`CancellationCallbackState::cancellation_callback`].
    callback: extern "C-unwind" fn(data: *const Self, action: u64, timeout: f64) -> i32,
    drop: extern "C-unwind" fn(*mut Self),

    /// The pointer to the closure object.
    closure_pointer: *mut (),

    /// This is the background job fiber which will be cancelled by this token.
    pub fiber_id: FiberId,
}

impl Drop for FfiBackgroundJobCancellationToken {
    #[inline(always)]
    fn drop(&mut self) {
        (self.drop)(self)
    }
}

impl FfiBackgroundJobCancellationToken {
    fn new(fiber_id: FiberId, cancel_channel: Channel<()>, finish_event: Rc<OnceEvent>) -> Self {
        let callback_state = CancellationCallbackState::new(cancel_channel, finish_event);
        let callback = move |action, timeout| {
            let res = callback_state.cancellation_callback(action, timeout);
            if let Err(e) = res {
                e.set_last();
                return -1;
            }

            0
        };

        Self::new_inner(fiber_id, callback)
    }

    /// This function is needed, because we need this `F` type parameter so that
    /// we can specialize the `callback` and `drop` with it inside this function.
    /// If rust supported something like `type F = type_of(callback);` we
    /// wouldn't need this additional function and would just write this code in
    /// the [`Self::new`] above.
    // FIXME: just define an explicit extern "C" fn for cancellation_callback?
    fn new_inner<F>(fiber_id: FiberId, f: F) -> Self
    where
        F: FnMut(u64, Duration) -> i32,
    {
        let closure = Box::new(f);
        let closure_pointer: *mut F = Box::into_raw(closure);

        Self {
            callback: Self::trampoline::<F>,
            drop: Self::drop_handler::<F>,
            closure_pointer: closure_pointer.cast(),

            fiber_id,
        }
    }

    /// An ABI-safe wrapper which calls the rust closure stored in `handler`.
    ///
    /// The result of the closure is copied onto the fiber's region allocation
    /// and the pointer to that allocation is written into `output`.
    extern "C-unwind" fn trampoline<F>(data: *const Self, action: u64, timeout: f64) -> i32
    where
        F: FnMut(u64, Duration) -> i32,
    {
        // This is safe. To verify see `register_rpc_handler` above.
        let closure_pointer: *mut F = unsafe { (*data).closure_pointer.cast::<F>() };
        let closure = unsafe { &mut *closure_pointer };

        closure(action, Duration::from_secs_f64(timeout))
    }

    extern "C-unwind" fn drop_handler<F>(handler: *mut Self) {
        unsafe {
            let closure_pointer: *mut F = (*handler).closure_pointer.cast::<F>();
            let closure = Box::from_raw(closure_pointer);
            drop(closure);

            if cfg!(debug_assertions) {
                // Overwrite the pointer with garbage so that we fail loudly is case of a bug
                (*handler).closure_pointer = 0xcccccccccccccccc_u64 as _;
            }
        }
    }

    /// The error is returned via [`BoxError::set_last`].
    #[inline(always)]
    pub fn cancel_job(&self) {
        let rc = (self.callback)(self, CancellationCallbackStatus::JobCancelled as _, 0.0);
        debug_assert!(rc == 0);
    }

    /// The error is returned via [`BoxError::set_last`].
    #[inline(always)]
    #[allow(clippy::result_unit_err)]
    pub fn wait_job_finished(&self, timeout: Duration) -> Result<(), ()> {
        let rc = (self.callback)(
            self,
            CancellationCallbackStatus::JobFinished as _,
            timeout.as_secs_f64(),
        );
        if rc == -1 {
            // Actual error is passed through tarantool. Can't return BoxError
            // here, because tarantool-module version may be different in picodata.
            return Err(());
        }

        Ok(())
    }
}