lambda_runtime_context 0.1.0

//! Task-local runtime deadline helpers for async Rust workers.
//!
//! `lambda_runtime_context` stores an invocation deadline in a Tokio task-local and
//! exposes small helpers for checking whether enough execution time remains to
//! start more work. It is useful for AWS Lambda handlers, queue consumers, and
//! other bounded async runtimes where a worker should stop early and re-enqueue
//! remaining work before the platform terminates the process.
//!
//! # Example
//!
//! ```rust
//! use std::time::Duration;
//!
//! # async fn enqueue_more_work() {}
//! # async fn process_next_item() {}
//! # async fn run(deadline_ms: u64) {
//! lambda_runtime_context::scope_deadline(deadline_ms, async {
//!     while lambda_runtime_context::is_duration_available(Duration::from_secs(30)) {
//!         process_next_item().await;
//!     }
//!
//!     enqueue_more_work().await;
//! })
//! .await;
//! # }
//! ```
//!
//! # Tokio task-local propagation
//!
//! Tokio task-locals are scoped to the current task. If you spawn a new task,
//! capture [`current_deadline_ms`] and call [`scope_deadline`] inside the
//! spawned future.
//!
//! ```rust
//! # async fn do_spawned_work() {}
//! # async fn run() {
//! if let Some(deadline_ms) = lambda_runtime_context::current_deadline_ms() {
//!     tokio::spawn(async move {
//!         lambda_runtime_context::scope_deadline(deadline_ms, async {
//!             do_spawned_work().await;
//!         })
//!         .await;
//!     });
//! }
//! # }
//! ```

use std::time::Duration;

tokio::task_local! {
    static INVOCATION_DEADLINE_MS: u64;
}

fn now_millis() -> u64 {
    let now = time::OffsetDateTime::now_utc();
    now.unix_timestamp() as u64 * 1_000 + now.millisecond() as u64
}

/// Execute the provided future with the Lambda deadline stored in a task-local slot.
pub async fn scope_deadline<F>(deadline_ms: u64, fut: F) -> F::Output
where
    F: Future,
{
    INVOCATION_DEADLINE_MS.scope(deadline_ms, fut).await
}

/// Returns the raw deadline timestamp (milliseconds since Unix epoch) for the current invocation.
///
/// Intended for propagating the deadline into `tokio::task::spawn` closures, which do not
/// inherit task-locals from the spawning task. Pass the returned value to `scope_deadline`
/// inside the spawned future.
pub fn current_deadline_ms() -> Option<u64> {
    INVOCATION_DEADLINE_MS.try_with(|&ms| ms).ok()
}

/// Returns how much time is left before the Lambda runtime terminates the invocation.
///
/// When running outside Lambda, or before `scope_deadline` is invoked, this returns `None`.
pub fn time_remaining() -> Option<Duration> {
    INVOCATION_DEADLINE_MS
        .try_with(|deadline_ms| {
            let now = now_millis();
            deadline_ms.saturating_sub(now)
        })
        .ok()
        .map(Duration::from_millis)
}

/// Returns the total time limit for the current invocation.
pub fn total_time_limit() -> Option<Duration> {
    INVOCATION_DEADLINE_MS
        .try_with(|deadline_ms| {
            Duration::from_millis(*deadline_ms).saturating_sub(Duration::from_millis(now_millis()))
        })
        .ok()
}

/// Returns true if the given time is still available before the deadline.
pub fn is_duration_available(duration: Duration) -> bool {
    time_remaining()
        .map(|remaining| remaining > Duration::ZERO && remaining >= duration)
        .unwrap_or(false)
}

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

    use super::*;

    #[tokio::test]
    async fn scope_provides_remaining_time() {
        let now = now_millis();
        let fut = async move {
            let remaining = time_remaining().expect("deadline present");
            println!("Remaining time: {:?}", remaining.as_nanos());

            for i in 0..5 {
                sleep(Duration::from_millis(100)).await;
                println!("Slept {}00 ms", i + 1);

                let remaining = time_remaining().expect("deadline present");
                println!("Remaining time: {:?}", remaining.as_nanos());
            }

            assert!(remaining >= Duration::from_millis(950));
        };

        scope_deadline(now + 1_000, fut).await;
    }

    #[tokio::test]
    async fn test_total_time_limit() {
        let fut = async move {
            let total_time = total_time_limit().expect("deadline present");
            assert_eq!(total_time, Duration::from_millis(2_000));
        };

        scope_deadline(now_millis() + 2_000, fut).await;
    }

    #[test]
    fn missing_deadline_returns_none() {
        assert!(time_remaining().is_none());
    }

    #[tokio::test]
    async fn test_is_runtime_time_available() {
        let fut = async move {
            assert!(is_duration_available(Duration::from_millis(500)));
            assert!(!is_duration_available(Duration::from_millis(1_500)));
        };
        scope_deadline(now_millis() + 1_000, fut).await;
    }
}