1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! Well known uses for provider contexts.
//!
//! This mirrors what you might find with Go's [`context.Context`](https://pkg.go.dev/context).

use core::{any::Provider, future::Future};
use std::time::{Duration, Instant};

use crate::{with_ref, ProviderFut, ProviderFutExt};

pub use shutdown::{ShutdownReceiver, ShutdownSender};

#[cfg(feature = "time")]
pub use shutdown::time::{run_until_signal, SignalOrComplete};

mod shutdown;

/// Extension trait to provide some well known context values
pub trait WellKnownProviderExt: Future + Sized {
    /// Wraps a [`Future`] so that it should expire within the given duration.
    ///
    /// Note, this doesn't guarantee that the future will stop executing, this is up
    /// to the implementation to respect the timeout.
    ///
    /// ```
    /// use std::time::Duration;
    /// use context_rs::well_known::{
    ///     WellKnownProviderExt, ShutdownSender, run_until_signal, SignalOrComplete
    /// };
    ///
    /// async fn do_something() {
    ///     loop {
    ///         // pretend this is more interesting
    ///         let work = tokio::time::sleep(Duration::from_secs(1));
    ///         match run_until_signal(std::pin::pin!(work)).await {
    ///             SignalOrComplete::Completed(_) => continue,
    ///             SignalOrComplete::ShutdownSignal(_) => break,
    ///         }
    ///     }
    /// }
    ///
    /// # #[tokio::main] async fn main() {
    /// do_something().with_timeout(Duration::from_secs(5)).await;
    /// # }
    /// ```
    fn with_timeout(self, duration: Duration) -> ProviderFut<Self, Deadline> {
        self.with_deadline(Instant::now() + duration)
    }

    /// Wraps a [`Future`] o that it should expire at the given deadline.
    ///
    /// See [`with_timeout`](WellKnownProviderExt::with_timeout) for more
    fn with_deadline(self, deadline: Instant) -> ProviderFut<Self, Deadline> {
        self.provide(Deadline(deadline))
    }

    /// Wraps a [`Future`] so it can be shutdown externally.
    ///
    /// ```
    /// use context_rs::well_known::{
    ///     WellKnownProviderExt, ShutdownSender, run_until_signal, SignalOrComplete
    /// };
    ///
    /// async fn generator(tx: tokio::sync::mpsc::Sender<i32>) {
    ///     for i in 0.. {
    ///         match run_until_signal(std::pin::pin!(tx.send(i))).await {
    ///             SignalOrComplete::Completed(_) => continue,
    ///             SignalOrComplete::ShutdownSignal(_) => break,
    ///         }
    ///     }
    /// }
    ///
    /// # #[tokio::main] async fn main() {
    /// let shutdown = ShutdownSender::new();
    /// let (tx, mut rx) = tokio::sync::mpsc::channel(1);
    ///
    /// tokio::spawn(generator(tx).with_shutdown_handler(shutdown.clone().receiver()));
    ///
    /// while let Some(x) = rx.recv().await {
    ///     dbg!(x);
    ///     if x == 5 {
    ///         break
    ///     }
    /// }
    ///
    /// // shutdown now that we're done with the rx
    /// shutdown.shutdown()
    /// # }
    /// ```
    fn with_shutdown_handler(
        self,
        handler: ShutdownReceiver,
    ) -> ProviderFut<Self, ShutdownReceiver> {
        self.provide(handler)
    }
}
impl<F: Future> WellKnownProviderExt for F {}

/// A Deadline type that wraps an `Instant`.
#[derive(Debug, Clone, Copy)]
pub struct Deadline(pub Instant);

impl Provider for Deadline {
    fn provide<'a>(&'a self, demand: &mut core::any::Demand<'a>) {
        demand.provide_ref(self);
    }
}

impl Deadline {
    /// Returns the deadline of the current context, if there is one
    pub async fn get() -> Option<Instant> {
        with_ref(|Deadline(deadline)| *deadline).await
    }
}