|
| 1 | +use std::time::Duration; |
| 2 | + |
| 3 | +use tokio::sync::mpsc; |
| 4 | +use tokio::time::timeout; |
| 5 | + |
| 6 | +/// Receives from a Tokio unbounded channel with an adaptive linger policy. |
| 7 | +/// |
| 8 | +/// This helper is intended for single-consumer background workers that want |
| 9 | +/// to avoid parking on `recv()` after every message during bursty traffic. |
| 10 | +/// |
| 11 | +/// The receiver has two modes - hot and cold. In cold mode it blocks on |
| 12 | +/// `recv()` until the next message arrives. In hot mode it prefers to stay |
| 13 | +/// awake, so after receiving a message, it will drain the channel and wait |
| 14 | +/// a short period (linger) before falling back to cold mode. |
| 15 | +/// |
| 16 | +/// The linger policy is as follows: If a message arrives while we are in |
| 17 | +/// the linger window, double the window up to `max_linger`. If the linger |
| 18 | +/// timer expires at any point without receiving a new message, reset the |
| 19 | +/// window to `baseline_linger`. |
| 20 | +/// |
| 21 | +/// Note, messages returned immediately by `try_recv()` do not count as hits, |
| 22 | +/// and do not double the linger window. |
| 23 | +#[derive(Debug)] |
| 24 | +pub struct AdaptiveUnboundedReceiver<T> { |
| 25 | + rx: mpsc::UnboundedReceiver<T>, |
| 26 | + linger: AdaptiveLinger, |
| 27 | + is_hot: bool, |
| 28 | +} |
| 29 | + |
| 30 | +impl<T> AdaptiveUnboundedReceiver<T> { |
| 31 | + /// Create an adaptive receiver around a Tokio unbounded channel. |
| 32 | + /// |
| 33 | + /// `baseline_linger` is the linger window used after a cold wakeup or any |
| 34 | + /// linger miss. `max_linger` caps how far the linger window may grow after |
| 35 | + /// repeated linger hits. |
| 36 | + /// |
| 37 | + /// This constructor does not spawn any tasks and does not alter the |
| 38 | + /// channel's ordering semantics. It only configures how aggressively the |
| 39 | + /// consumer stays awake after work arrives. |
| 40 | + pub fn new(rx: mpsc::UnboundedReceiver<T>, baseline_linger: Duration, max_linger: Duration) -> Self { |
| 41 | + Self { |
| 42 | + rx, |
| 43 | + linger: AdaptiveLinger::new(baseline_linger, max_linger), |
| 44 | + is_hot: false, |
| 45 | + } |
| 46 | + } |
| 47 | + |
| 48 | + /// Receive the next message while adapting how aggressively we linger |
| 49 | + /// before parking again. |
| 50 | + /// |
| 51 | + /// Once a worker has been woken up by one message, subsequent calls try to |
| 52 | + /// stay on the hot path: |
| 53 | + /// |
| 54 | + /// 1. Drain any already-queued work immediately with `try_recv()` |
| 55 | + /// 2. If the queue is empty, wait for the current linger window |
| 56 | + /// 3. On a linger hit, double the window and continue lingering |
| 57 | + /// 4. On a linger miss, reset the window to the baseline and park on `recv()` |
| 58 | + /// |
| 59 | + /// This minimizes latency during periods of low activity but maximizes |
| 60 | + /// throughput during periods of high activity. |
| 61 | + pub async fn recv(&mut self) -> Option<T> { |
| 62 | + loop { |
| 63 | + if !self.is_hot { |
| 64 | + let message = self.rx.recv().await?; |
| 65 | + self.is_hot = true; |
| 66 | + return Some(message); |
| 67 | + } |
| 68 | + |
| 69 | + match self.rx.try_recv() { |
| 70 | + Ok(message) => return Some(message), |
| 71 | + Err(mpsc::error::TryRecvError::Disconnected) => return None, |
| 72 | + Err(mpsc::error::TryRecvError::Empty) => {} |
| 73 | + } |
| 74 | + |
| 75 | + let linger = self.linger.current(); |
| 76 | + if linger.is_zero() { |
| 77 | + self.cool_down(); |
| 78 | + continue; |
| 79 | + } |
| 80 | + |
| 81 | + match timeout(linger, self.rx.recv()).await { |
| 82 | + Ok(Some(message)) => { |
| 83 | + self.linger.on_hit(); |
| 84 | + return Some(message); |
| 85 | + } |
| 86 | + Ok(None) => return None, |
| 87 | + Err(_) => { |
| 88 | + self.cool_down(); |
| 89 | + } |
| 90 | + } |
| 91 | + } |
| 92 | + } |
| 93 | + |
| 94 | + /// Return the receiver to its cold state after a linger miss. |
| 95 | + /// |
| 96 | + /// The next call to [`Self::recv`] will block on the underlying channel |
| 97 | + /// instead of continuing to linger, and the linger policy is reset to its |
| 98 | + /// baseline window. |
| 99 | + fn cool_down(&mut self) { |
| 100 | + self.is_hot = false; |
| 101 | + self.linger.on_miss(); |
| 102 | + } |
| 103 | +} |
| 104 | + |
| 105 | +#[derive(Debug)] |
| 106 | +struct AdaptiveLinger { |
| 107 | + baseline: Duration, |
| 108 | + current: Duration, |
| 109 | + max: Duration, |
| 110 | +} |
| 111 | + |
| 112 | +impl AdaptiveLinger { |
| 113 | + /// Create a linger policy with a baseline window and an upper bound. |
| 114 | + /// |
| 115 | + /// `baseline` is the window restored after any linger miss. `max` caps how |
| 116 | + /// far the window may grow after repeated linger hits. |
| 117 | + fn new(baseline: Duration, max: Duration) -> Self { |
| 118 | + assert!( |
| 119 | + baseline <= max, |
| 120 | + "baseline linger ({baseline:?}) must not exceed max linger ({max:?})" |
| 121 | + ); |
| 122 | + Self { |
| 123 | + baseline, |
| 124 | + current: baseline, |
| 125 | + max, |
| 126 | + } |
| 127 | + } |
| 128 | + |
| 129 | + /// Return the current linger window. |
| 130 | + fn current(&self) -> Duration { |
| 131 | + self.current |
| 132 | + } |
| 133 | + |
| 134 | + /// Record a linger hit by growing the next linger window. |
| 135 | + /// |
| 136 | + /// The window doubles on each hit until it reaches `self.max`. |
| 137 | + fn on_hit(&mut self) { |
| 138 | + self.current = self.current.saturating_mul(2).min(self.max); |
| 139 | + } |
| 140 | + |
| 141 | + /// Record a linger miss by resetting to the baseline window. |
| 142 | + fn on_miss(&mut self) { |
| 143 | + self.current = self.baseline; |
| 144 | + } |
| 145 | +} |
0 commit comments