pingora-timeout 0.9.0

// Copyright 2026 Cloudflare, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! The fast and more complicated version of pingora-timeout
//!
//! The following optimizations are applied:
//! - The timeouts lazily initialize their timer when the Future is pending for the first time.
//! - There is no global lock for creating and cancelling timeouts.
//! - Timeout timers are rounded to the next 10ms tick and timers are shared across all timeouts with the same deadline.
//!
//! In order for this to work, a standalone thread is created to arm the timers, which has some
//! overhead. As a general rule, the benefits of this don't outweigh the overhead unless
//! there are more than about 100 timeout() calls/sec in the system. Use regular tokio timeout or
//! [super::tokio_timeout] in the low usage case.

use super::timer::*;
use super::*;
use once_cell::sync::Lazy;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

static TIMER_MANAGER: Lazy<Arc<TimerManager>> = Lazy::new(|| {
    let tm = Arc::new(TimerManager::new());
    check_clock_thread(&tm);
    tm
});

/// Default duration above which [`fast_timeout()`] falls back to Tokio's timeout.
pub const DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD: Duration = Duration::from_secs(15 * 60);
const DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS: u64 =
    DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD.as_secs();
const FAST_TIMEOUT_TO_TOKIO_DISABLED: u64 = u64::MAX;

static FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS: AtomicU64 =
    AtomicU64::new(DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS);

fn check_clock_thread(tm: &Arc<TimerManager>) {
    if tm.should_i_start_clock() {
        std::thread::Builder::new()
            .name("Timer thread".into())
            .spawn(|| TIMER_MANAGER.clock_thread())
            .unwrap();
    }
}

fn should_use_tokio_timeout(duration: Duration) -> bool {
    let threshold_secs = FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS.load(Ordering::Relaxed);
    threshold_secs != FAST_TIMEOUT_TO_TOKIO_DISABLED && duration.as_secs() > threshold_secs
}

/// Set the duration above which [`fast_timeout()`] falls back to Tokio's timeout.
///
/// When set to `None`, [`fast_timeout()`] never falls back to Tokio's timeout.
///
/// Long Pingora fast-timeout timers are kept in [`TimerManager`] until
/// their deadline expires. Tokio's timeout has cancellation cleanup, so it is a
/// better fit for long durations that are often cancelled before expiry.
pub fn set_fast_timeout_to_tokio_threshold(threshold: Option<Duration>) {
    FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS.store(
        threshold.map_or(FAST_TIMEOUT_TO_TOKIO_DISABLED, |threshold| {
            threshold.as_secs()
        }),
        Ordering::Relaxed,
    );
}

/// Return the duration above which [`fast_timeout()`] falls back to Tokio's timeout.
pub fn fast_timeout_to_tokio_threshold() -> Option<Duration> {
    match FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS.load(Ordering::Relaxed) {
        FAST_TIMEOUT_TO_TOKIO_DISABLED => None,
        threshold_secs => Some(Duration::from_secs(threshold_secs)),
    }
}

/// The timeout generated by [fast_timeout()].
///
/// Users don't need to interact with this object.
pub struct FastTimeout(Duration);

impl ToTimeout for FastTimeout {
    fn timeout(&self) -> Pin<Box<dyn Future<Output = ()> + Send + Sync>> {
        Box::pin(TIMER_MANAGER.register_timer(self.0).poll())
    }

    fn create(d: Duration) -> Self {
        FastTimeout(d)
    }
}

enum TimeoutKind {
    Fast,
    Tokio,
}

/// The timeout selector generated by [fast_timeout()].
///
/// Users don't need to interact with this object.
pub struct FastOrTokioTimeout {
    duration: Duration,
    kind: TimeoutKind,
}

impl FastOrTokioTimeout {
    fn use_tokio(&self) -> bool {
        matches!(self.kind, TimeoutKind::Tokio)
    }
}

impl ToTimeout for FastOrTokioTimeout {
    fn timeout(&self) -> Pin<Box<dyn Future<Output = ()> + Send + Sync>> {
        match self.kind {
            TimeoutKind::Fast => FastTimeout::create(self.duration).timeout(),
            TimeoutKind::Tokio => TokioTimeout::create(self.duration).timeout(),
        }
    }

    fn create(duration: Duration) -> Self {
        Self {
            duration,
            kind: if should_use_tokio_timeout(duration) {
                TimeoutKind::Tokio
            } else {
                TimeoutKind::Fast
            },
        }
    }
}

/// Similar to [tokio::time::timeout] but more efficient.
pub fn fast_timeout<T>(duration: Duration, future: T) -> Timeout<T, FastOrTokioTimeout>
where
    T: Future,
{
    let callback = FastOrTokioTimeout::create(duration);
    if !callback.use_tokio() {
        check_clock_thread(&TIMER_MANAGER);
    }

    Timeout::new_with_callback(future, callback)
}

/// Similar to [tokio::time::sleep] but more efficient.
pub async fn fast_sleep(duration: Duration) {
    check_clock_thread(&TIMER_MANAGER);
    TIMER_MANAGER.register_timer(duration).poll().await
}

/// Pause the timer for fork()
///
/// Because RwLock across fork() is undefined behavior, this function makes sure that no one
/// holds any locks.
///
/// This function should be called right before fork().
pub fn pause_for_fork() {
    TIMER_MANAGER.pause_for_fork();
}

/// Unpause the timer after fork()
///
/// This function should be called right after fork().
pub fn unpause() {
    TIMER_MANAGER.unpause();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_timeout() {
        let fut = tokio_sleep(Duration::from_secs(1000));
        let to = fast_timeout(Duration::from_secs(1), fut);
        assert!(to.await.is_err())
    }

    #[tokio::test]
    async fn test_instantly_return() {
        let fut = async { 1 };
        let to = fast_timeout(Duration::from_secs(1), fut);
        assert_eq!(to.await.unwrap(), 1)
    }

    #[tokio::test]
    async fn test_delayed_return() {
        let fut = async {
            tokio_sleep(Duration::from_secs(1)).await;
            1
        };
        let to = fast_timeout(Duration::from_secs(60), fut);
        assert_eq!(to.await.unwrap(), 1)
    }

    #[tokio::test]
    async fn test_sleep() {
        let fut = async {
            fast_sleep(Duration::from_secs(1)).await;
            1
        };
        let to = fast_timeout(Duration::from_secs(60), fut);
        assert_eq!(to.await.unwrap(), 1)
    }

    #[test]
    fn test_default_fast_timeout_to_tokio_threshold() {
        assert_eq!(
            Some(Duration::from_secs(15 * 60)),
            fast_timeout_to_tokio_threshold()
        );
    }

    #[test]
    fn test_fast_timeout_uses_fast_below_threshold() {
        let callback = FastOrTokioTimeout::create(Duration::from_secs(1));
        assert!(!callback.use_tokio());
    }

    #[test]
    fn test_fast_timeout_uses_fast_at_threshold() {
        let callback = FastOrTokioTimeout::create(fast_timeout_to_tokio_threshold().unwrap());
        assert!(!callback.use_tokio());
    }

    #[test]
    fn test_fast_timeout_uses_tokio_above_threshold() {
        let callback = FastOrTokioTimeout::create(
            fast_timeout_to_tokio_threshold().unwrap() + Duration::from_secs(1),
        );
        assert!(callback.use_tokio());
    }
}