common_batcher/
notifier.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::future::Future;
16use std::num::NonZeroUsize;
17
18use futures::{StreamExt, stream};
19use tokio::sync::Semaphore;
20use tokio::sync::mpsc::error::TrySendError;
21use tokio::sync::mpsc::{self, Receiver, Sender};
22
23/// Best-effort notification admission, independent of the write result.
24///
25/// The caller owns the notification payload, delivery handler and failure metrics.
26pub struct Notifier<T> {
27    sender: Sender<T>,
28}
29
30impl<T> Clone for Notifier<T> {
31    fn clone(&self) -> Self {
32        Self {
33            sender: self.sender.clone(),
34        }
35    }
36}
37
38impl<T> Notifier<T> {
39    /// Returns `None` if the queue capacity is zero or exceeds Tokio's supported capacity.
40    pub fn try_new(capacity: usize) -> Option<(Self, Receiver<T>)> {
41        if capacity == 0 || capacity > Semaphore::MAX_PERMITS {
42            return None;
43        }
44        let (sender, receiver) = mpsc::channel(capacity);
45        Some((Self { sender }, receiver))
46    }
47
48    /// Configured queue capacity, independent of its current occupancy.
49    pub fn max_capacity(&self) -> usize {
50        self.sender.max_capacity()
51    }
52
53    /// Never waits for delivery. Returns the payload on a full or closed queue
54    /// so the caller can record domain-specific diagnostics.
55    pub fn try_notify(&self, notification: T) -> Result<(), TrySendError<T>> {
56        self.sender.try_send(notification)
57    }
58}
59
60/// Delivers queued notifications with bounded concurrency.
61///
62/// The caller schedules this future on its runtime. Closing all senders drains
63/// queued and in-flight notifications. Handler failures are handled by the caller;
64/// this worker does not retry or propagate them into write completion.
65pub async fn run_notifier<T, F, Fut>(
66    receiver: Receiver<T>,
67    max_concurrent_notifications: NonZeroUsize,
68    handler: F,
69) where
70    F: FnMut(T) -> Fut,
71    Fut: Future<Output = ()>,
72{
73    stream::unfold(receiver, |mut receiver| async move {
74        receiver.recv().await.map(|item| (item, receiver))
75    })
76    .for_each_concurrent(max_concurrent_notifications.get(), handler)
77    .await;
78}
79
80#[cfg(test)]
81mod tests {
82    use std::future::{Future, poll_fn};
83    use std::num::NonZeroUsize;
84    use std::sync::Arc;
85    use std::sync::atomic::{AtomicUsize, Ordering};
86    use std::task::Poll;
87
88    use tokio::sync::Semaphore;
89    use tokio::sync::mpsc::error::TrySendError;
90
91    use crate::notifier::{Notifier, run_notifier};
92
93    #[test]
94    fn test_admission() {
95        assert!(Notifier::<usize>::try_new(0).is_none());
96        assert!(Notifier::<usize>::try_new(usize::MAX).is_none());
97        let (notifier, receiver) = Notifier::try_new(1).unwrap();
98        let other = notifier.clone();
99        assert_eq!(other.max_capacity(), 1);
100        assert_eq!(notifier.try_notify(1), Ok(()));
101        assert_eq!(other.try_notify(2), Err(TrySendError::Full(2)));
102        drop(receiver);
103        assert_eq!(other.try_notify(3), Err(TrySendError::Closed(3)));
104    }
105
106    #[tokio::test]
107    async fn test_concurrency_and_drain() {
108        // One-factor comparison: only delivery concurrency changes.
109        for _ in 0..2 {
110            for concurrency in [1, 2] {
111                let (notifier, receiver) = Notifier::try_new(3).unwrap();
112                for item in 0..3 {
113                    notifier.try_notify(item).unwrap();
114                }
115                drop(notifier);
116                let started = Arc::new(AtomicUsize::new(0));
117                let completed = Arc::new(AtomicUsize::new(0));
118                let gate = Arc::new(Semaphore::new(0));
119                let worker =
120                    run_notifier(receiver, NonZeroUsize::new(concurrency).unwrap(), |_| {
121                        let started = started.clone();
122                        let completed = completed.clone();
123                        let gate = gate.clone();
124                        async move {
125                            started.fetch_add(1, Ordering::SeqCst);
126                            gate.acquire().await.unwrap().forget();
127                            completed.fetch_add(1, Ordering::SeqCst);
128                        }
129                    });
130                tokio::pin!(worker);
131                poll_fn(|cx| {
132                    assert!(worker.as_mut().poll(cx).is_pending());
133                    Poll::Ready(())
134                })
135                .await;
136                assert_eq!(started.load(Ordering::SeqCst), concurrency);
137                assert_eq!(completed.load(Ordering::SeqCst), 0);
138                gate.add_permits(3);
139                worker.await;
140                assert_eq!(started.load(Ordering::SeqCst), 3);
141                assert_eq!(completed.load(Ordering::SeqCst), 3);
142            }
143        }
144    }
145}
common_batcher/notifier.rs

common_batcher/
notifier.rs
