risingwave_meta/barrier/
notifier.rs

1// Copyright 2022 RisingWave Labs
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 anyhow::anyhow;
16use futures::future::join_all;
17use tokio::sync::oneshot;
18
19use crate::{MetaError, MetaResult};
20
21pub(crate) type CollectionReceiver = oneshot::Receiver<MetaResult<()>>;
22pub(crate) type StartReceiver = oneshot::Receiver<MetaResult<Vec<CollectionReceiver>>>;
23
24pub(crate) async fn wait_collection(receivers: Vec<CollectionReceiver>) -> MetaResult<()> {
25    let mut first_error = None;
26    for result in join_all(receivers).await {
27        let result = match result {
28            Ok(result) => result,
29            Err(_) => Err(anyhow!("failed to collect barrier: notifier dropped").into()),
30        };
31        if let Err(err) = result
32            && first_error.is_none()
33        {
34            first_error = Some(err);
35        }
36    }
37    match first_error {
38        Some(err) => Err(err),
39        None => Ok(()),
40    }
41}
42
43/// Used for notifying the status of a scheduled command/barrier.
44#[derive(Debug)]
45pub(crate) struct Notifier {
46    started: oneshot::Sender<MetaResult<Vec<CollectionReceiver>>>,
47}
48
49impl Notifier {
50    pub fn new() -> (Self, StartReceiver) {
51        let (started, receiver) = oneshot::channel();
52        (Self { started }, receiver)
53    }
54
55    pub fn start(self) -> NotifierStarter {
56        NotifierStarter {
57            started: self.started,
58            pending_collection: vec![],
59        }
60    }
61
62    pub fn notify_start_failed(self, err: MetaError) {
63        self.started.send(Err(err)).ok();
64    }
65}
66
67/// Builds the set of collection notifications before publishing that the command has started.
68#[derive(Debug)]
69pub(crate) struct NotifierStarter {
70    started: oneshot::Sender<MetaResult<Vec<CollectionReceiver>>>,
71    pending_collection: Vec<CollectionReceiver>,
72}
73
74impl NotifierStarter {
75    pub fn add_notify(&mut self) -> CollectionNotifier {
76        let (collected, receiver) = oneshot::channel();
77        self.pending_collection.push(receiver);
78        CollectionNotifier { collected }
79    }
80
81    pub fn started(self) {
82        self.started.send(Ok(self.pending_collection)).ok();
83    }
84
85    pub fn notify_start_failed(self, err: MetaError) {
86        self.started.send(Err(err)).ok();
87    }
88}
89
90/// Notifies the completion of one part of a started command.
91#[derive(Debug)]
92pub(crate) struct CollectionNotifier {
93    collected: oneshot::Sender<MetaResult<()>>,
94}
95
96impl CollectionNotifier {
97    pub fn notify_collected(self) {
98        self.collected.send(Ok(())).ok();
99    }
100
101    /// Notify when we failed to collect a barrier. This function consumes `self`.
102    pub fn notify_collection_failed(self, err: MetaError) {
103        self.collected.send(Err(err)).ok();
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use tokio::time::{Duration, timeout};
110
111    use super::*;
112
113    #[tokio::test]
114    async fn test_zero_collection_notifier() {
115        let (notifier, started_rx) = Notifier::new();
116        notifier.start().started();
117
118        let receivers = started_rx.await.unwrap().unwrap();
119        assert!(receivers.is_empty());
120        wait_collection(receivers).await.unwrap();
121    }
122
123    #[tokio::test]
124    async fn test_one_collection_notifier() {
125        let (notifier, started_rx) = Notifier::new();
126        let mut start = notifier.start();
127        let collection = start.add_notify();
128        start.started();
129
130        collection.notify_collected();
131        let receivers = started_rx.await.unwrap().unwrap();
132        wait_collection(receivers).await.unwrap();
133    }
134
135    #[tokio::test]
136    async fn test_wait_all_collection_notifiers() {
137        let (notifier, started_rx) = Notifier::new();
138        let mut start = notifier.start();
139        let first = start.add_notify();
140        let second = start.add_notify();
141        start.started();
142
143        let receivers = started_rx.await.unwrap().unwrap();
144        let mut wait = Box::pin(wait_collection(receivers));
145        first.notify_collected();
146        assert!(timeout(Duration::from_millis(10), &mut wait).await.is_err());
147        second.notify_collected();
148        wait.await.unwrap();
149    }
150
151    #[tokio::test]
152    async fn test_waits_for_all_collection_notifiers_after_failure() {
153        let (notifier, started_rx) = Notifier::new();
154        let mut start = notifier.start();
155        let first = start.add_notify();
156        let second = start.add_notify();
157        start.started();
158
159        let receivers = started_rx.await.unwrap().unwrap();
160        let mut wait = Box::pin(wait_collection(receivers));
161        first.notify_collection_failed(anyhow!("first part failed").into());
162        assert!(timeout(Duration::from_millis(10), &mut wait).await.is_err());
163        second.notify_collection_failed(anyhow!("second part failed").into());
164        let err = wait.await.unwrap_err();
165        assert!(err.to_string().contains("first part failed"));
166    }
167
168    #[tokio::test]
169    async fn test_dropped_collection_notifier_fails() {
170        let (notifier, started_rx) = Notifier::new();
171        let mut start = notifier.start();
172        let dropped = start.add_notify();
173        start.started();
174        drop(dropped);
175
176        let receivers = started_rx.await.unwrap().unwrap();
177        assert!(wait_collection(receivers).await.is_err());
178    }
179
180    #[tokio::test]
181    async fn test_start_failure() {
182        let (notifier, started_rx) = Notifier::new();
183        notifier.notify_start_failed(anyhow!("start failed").into());
184        assert!(started_rx.await.unwrap().is_err());
185    }
186
187    #[tokio::test]
188    async fn test_start_failure_after_entering_start_phase() {
189        let (notifier, started_rx) = Notifier::new();
190        let mut start = notifier.start();
191        let collection = start.add_notify();
192        start.notify_start_failed(anyhow!("start failed").into());
193
194        assert!(started_rx.await.unwrap().is_err());
195        collection.notify_collected();
196    }
197
198    #[tokio::test]
199    async fn test_dropped_notifier_cancels_start() {
200        let (notifier, started_rx) = Notifier::new();
201        drop(notifier);
202        assert!(started_rx.await.is_err());
203    }
204}
risingwave_meta/barrier/notifier.rs

risingwave_meta/barrier/
notifier.rs
