Skip to main content

regulatory_region_lib/
pub_sub_hub.rs

1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! A simple publish-and-subscribe facility.
6
7use serde::{Deserialize, Serialize};
8use std::cell::RefCell;
9use std::collections::BTreeMap;
10use std::fs::{self, File};
11use std::future::Future;
12use std::io;
13use std::path::{Path, PathBuf};
14use std::pin::Pin;
15use std::task::{Context, Poll, Waker};
16
17/// A rendezvous point for publishers and subscribers.
18pub struct PubSubHub {
19    // To minimize the risk of run-time errors, we never store the borrowed `inner` in a named
20    // variable. By only borrowing `inner` via temporaries, we make any simultaneous borrows easier
21    // to spot in review. And spotting simultaneous borrows enables us to spot conflicting borrows
22    // (simultaneous `borrow()` and `borrow_mut()`.)
23    inner: RefCell<PubSubHubInner>,
24    // The path of the file to load from and write the current value to.
25    storage_path: PathBuf,
26}
27
28/// The `Future` used by a subscriber to `await` updates.
29pub struct PubSubFuture<'a> {
30    // See comment for `PubSubHub::inner`, about how to borrow from `hub`.
31    hub: &'a RefCell<PubSubHubInner>,
32    id: usize,
33    last_value: Option<String>,
34}
35
36struct PubSubHubInner {
37    item: Option<String>,
38    next_future_id: usize,
39    wakers: BTreeMap<usize, Waker>,
40}
41
42impl PubSubHub {
43    pub fn new(storage_path: PathBuf) -> Self {
44        let initial_value = load_region_code(&storage_path);
45        Self {
46            inner: RefCell::new(PubSubHubInner {
47                item: initial_value,
48                next_future_id: 0,
49                wakers: BTreeMap::new(),
50            }),
51            storage_path,
52        }
53    }
54
55    /// Publishes `new_value`.
56    /// * All pending futures are woken.
57    /// * Later calls to `watch_for_change()` will be evaluated against `new_value`.
58    pub fn publish<S>(&self, new_value: S)
59    where
60        S: Into<String>,
61    {
62        let hub = &self.inner;
63        let new_value = new_value.into();
64        hub.borrow_mut().item = Some(new_value.clone());
65        hub.borrow_mut().wakers.values().for_each(|w| w.wake_by_ref());
66        hub.borrow_mut().wakers.clear();
67        // Store the value that should be loaded at startup.
68        write_region_code(new_value, &self.storage_path);
69    }
70
71    /// Watches the value stored in this hub, resolving when the
72    /// stored value differs from `last_value`.
73    pub fn watch_for_change<S>(&self, last_value: Option<S>) -> PubSubFuture<'_>
74    where
75        S: Into<String>,
76    {
77        let hub = &self.inner;
78        let id = hub.borrow().next_future_id;
79        hub.borrow_mut().next_future_id = id.checked_add(1).expect("`id` is impossibly large");
80        PubSubFuture { hub, id, last_value: last_value.map(|s| s.into()) }
81    }
82
83    pub fn get_value(&self) -> Option<String> {
84        let hub = &self.inner;
85        hub.borrow().get_value()
86    }
87}
88
89/// The regulatory region code as a struct, to be used for reading and writing the value as JSON.
90#[derive(Debug, Deserialize, Serialize)]
91struct RegulatoryRegion {
92    region_code: String,
93}
94
95// Try to load the stored region code from a file at the specified path. If an error occurs, it
96// will not cause a failure because the cache is not necessary.
97// TODO(67860) Add metric for failures reading cache.
98fn load_region_code(path: impl AsRef<Path>) -> Option<String> {
99    let file = match File::open(path.as_ref()) {
100        Ok(file) => file,
101        Err(e) => match e.kind() {
102            io::ErrorKind::NotFound => return None,
103            _ => {
104                log::info!(
105                    "Failed to read cached regulatory region, will initialize with none: {}",
106                    e
107                );
108                try_delete_file(path);
109                return None;
110            }
111        },
112    };
113    match serde_json::from_reader::<_, RegulatoryRegion>(io::BufReader::new(file)) {
114        Ok(region) => Some(region.region_code),
115        Err(e) => {
116            log::info!("Error parsing stored regulatory region code: {}", e);
117            try_delete_file(path);
118            None
119        }
120    }
121}
122
123/// Try to write the region code as a JSON at the specified file location. For example, the file
124/// contents may look like "{"region_code": "US"}". Errors saving the region code will not cause a
125/// failure because the cache is not necessary.
126// TODO(67860) Add metric for failures writing cache.
127fn write_region_code(region_code: String, storage_path: impl AsRef<Path>) {
128    let write_val = RegulatoryRegion { region_code };
129    let file = match File::create(storage_path.as_ref()) {
130        Ok(file) => file,
131        Err(e) => {
132            log::info!("Failed to open file to write regulatory region: {}", e);
133            try_delete_file(storage_path);
134            return;
135        }
136    };
137    if let Err(e) = serde_json::to_writer(io::BufWriter::new(file), &write_val) {
138        log::info!("Failed to write regulatory region: {}", e);
139        try_delete_file(storage_path);
140    }
141}
142
143fn try_delete_file(storage_path: impl AsRef<Path>) {
144    if let Err(e) = fs::remove_file(&storage_path) {
145        log::info!("Failed to delete previously cached regulatory region: {}", e);
146    }
147}
148
149impl Future for PubSubFuture<'_> {
150    type Output = Option<String>;
151
152    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
153        let hub = &self.hub;
154        if hub.borrow().has_value(&self.last_value) {
155            hub.borrow_mut().set_waker_for_future(self.id, context.waker().clone());
156            Poll::Pending
157        } else {
158            Poll::Ready(hub.borrow().get_value())
159        }
160    }
161}
162
163impl PubSubHubInner {
164    fn set_waker_for_future(&mut self, future_id: usize, waker: Waker) {
165        self.wakers.insert(future_id, waker);
166    }
167
168    fn has_value(&self, expected: &Option<String>) -> bool {
169        self.item == *expected
170    }
171
172    fn get_value(&self) -> Option<String> {
173        self.item.clone()
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use assert_matches::assert_matches;
181    use futures_test::task::new_count_waker;
182    use std::io::Write;
183    use tempfile::TempDir;
184
185    #[fuchsia::test(allow_stalls = false)]
186    async fn watch_for_change_future_is_pending_when_both_values_are_none() {
187        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
188        let path = temp_dir.path().join("regulatory_region.json");
189        let hub = PubSubHub::new(path);
190        let (waker, count) = new_count_waker();
191        let mut context = Context::from_waker(&waker);
192        let mut future = hub.watch_for_change(Option::<String>::None);
193        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
194        assert_eq!(0, count.get());
195    }
196
197    #[fuchsia::test(allow_stalls = false)]
198    async fn watch_for_change_future_is_pending_when_values_are_same_and_not_none() {
199        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
200        let path = temp_dir.path().join("regulatory_region.json");
201        let hub = PubSubHub::new(path);
202        let (waker, count) = new_count_waker();
203        let mut context = Context::from_waker(&waker);
204        hub.publish("US");
205
206        let mut future = hub.watch_for_change(Some("US"));
207        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
208        assert_eq!(0, count.get());
209    }
210
211    #[fuchsia::test(allow_stalls = false)]
212    async fn watch_for_change_future_is_immediately_ready_when_argument_differs_from_published_value()
213     {
214        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
215        let path = temp_dir.path().join("regulatory_region.json");
216        let hub = PubSubHub::new(path);
217        let (waker, count) = new_count_waker();
218        let mut context = Context::from_waker(&waker);
219        hub.publish("US");
220
221        let mut future = hub.watch_for_change(Option::<String>::None);
222        assert_eq!(Poll::Ready(Some("US".to_string())), Pin::new(&mut future).poll(&mut context));
223        assert_eq!(0, count.get());
224    }
225
226    #[fuchsia::test(allow_stalls = false)]
227    async fn single_watcher_is_woken_correctly_on_change_from_none_to_some() {
228        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
229        let path = temp_dir.path().join("regulatory_region.json");
230        let hub = PubSubHub::new(path);
231        let (waker, count) = new_count_waker();
232        let mut context = Context::from_waker(&waker);
233        let mut future = hub.watch_for_change(Option::<String>::None);
234        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
235
236        // Change value, and expect wake, and new value.
237        hub.publish("US");
238        assert_eq!(1, count.get());
239        assert_eq!(Poll::Ready(Some("US".to_string())), Pin::new(&mut future).poll(&mut context));
240    }
241
242    #[fuchsia::test(allow_stalls = false)]
243    async fn single_watcher_is_woken_correctly_on_change_from_some_to_new_some() {
244        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
245        let path = temp_dir.path().join("regulatory_region.json");
246        let hub = PubSubHub::new(path);
247        let (waker, count) = new_count_waker();
248        let mut context = Context::from_waker(&waker);
249        hub.publish("US");
250
251        let mut future = hub.watch_for_change(Some("US"));
252        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
253
254        // Change value, and expect wake, and new value.
255        hub.publish("SU");
256        assert_eq!(1, count.get());
257        assert_eq!(Poll::Ready(Some("SU".to_string())), Pin::new(&mut future).poll(&mut context));
258    }
259
260    #[fuchsia::test(allow_stalls = false)]
261    async fn multiple_watchers_are_woken_correctly_on_change_from_some_to_new_some() {
262        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
263        let path = temp_dir.path().join("regulatory_region.json");
264        let hub = PubSubHub::new(path);
265        let (waker_a, wake_count_a) = new_count_waker();
266        let (waker_b, wake_count_b) = new_count_waker();
267        let mut context_a = Context::from_waker(&waker_a);
268        let mut context_b = Context::from_waker(&waker_b);
269        hub.publish("US");
270
271        let mut future_a = hub.watch_for_change(Some("US"));
272        let mut future_b = hub.watch_for_change(Some("US"));
273        assert_eq!(Poll::Pending, Pin::new(&mut future_a).poll(&mut context_a), "for future a");
274        assert_eq!(Poll::Pending, Pin::new(&mut future_b).poll(&mut context_b), "for future b");
275
276        // Change value, and expect wakes, and new value for both futures.
277        hub.publish("SU");
278        assert_eq!(1, wake_count_a.get(), "for waker a");
279        assert_eq!(1, wake_count_b.get(), "for waker b");
280        assert_eq!(
281            Poll::Ready(Some("SU".to_string())),
282            Pin::new(&mut future_a).poll(&mut context_a),
283            "for future a"
284        );
285        assert_eq!(
286            Poll::Ready(Some("SU".to_string())),
287            Pin::new(&mut future_b).poll(&mut context_b),
288            "for future b"
289        );
290    }
291
292    #[fuchsia::test(allow_stalls = false)]
293    async fn multiple_watchers_are_woken_correctly_after_spurious_update() {
294        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
295        let path = temp_dir.path().join("regulatory_region.json");
296        let hub = PubSubHub::new(path);
297        let (waker_a, wake_count_a) = new_count_waker();
298        let (waker_b, wake_count_b) = new_count_waker();
299        let mut context_a = Context::from_waker(&waker_a);
300        let mut context_b = Context::from_waker(&waker_b);
301        hub.publish("US");
302
303        let mut future_a = hub.watch_for_change(Some("US"));
304        let mut future_b = hub.watch_for_change(Some("US"));
305        assert_eq!(Poll::Pending, Pin::new(&mut future_a).poll(&mut context_a), "for future a");
306        assert_eq!(Poll::Pending, Pin::new(&mut future_b).poll(&mut context_b), "for future b");
307
308        // Generate spurious update.
309        hub.publish("US");
310        assert_eq!(Poll::Pending, Pin::new(&mut future_a).poll(&mut context_a), "for future a");
311        assert_eq!(Poll::Pending, Pin::new(&mut future_b).poll(&mut context_b), "for future b");
312
313        // Generate a real update. Expect wakes, and new value for both futures.
314        let old_wake_count_a = wake_count_a.get();
315        let old_wake_count_b = wake_count_b.get();
316        hub.publish("SU");
317        assert_eq!(1, wake_count_a.get() - old_wake_count_a);
318        assert_eq!(1, wake_count_b.get() - old_wake_count_b);
319        assert_eq!(
320            Poll::Ready(Some("SU".to_string())),
321            Pin::new(&mut future_a).poll(&mut context_a),
322            "for future a"
323        );
324        assert_eq!(
325            Poll::Ready(Some("SU".to_string())),
326            Pin::new(&mut future_b).poll(&mut context_b),
327            "for future b"
328        );
329    }
330
331    #[fuchsia::test(allow_stalls = false)]
332    async fn multiple_watchers_can_share_a_waker() {
333        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
334        let path = temp_dir.path().join("regulatory_region.json");
335        let hub = PubSubHub::new(path);
336        let (waker, count) = new_count_waker();
337        let mut context = Context::from_waker(&waker);
338        let mut future_a = hub.watch_for_change(Option::<String>::None);
339        let mut future_b = hub.watch_for_change(Option::<String>::None);
340        assert_eq!(Poll::Pending, Pin::new(&mut future_a).poll(&mut context), "for future a");
341        assert_eq!(Poll::Pending, Pin::new(&mut future_b).poll(&mut context), "for future b");
342
343        // Change value, and expect wakes, and new value for both futures.
344        hub.publish("US");
345        assert_eq!(2, count.get());
346        assert_eq!(
347            Poll::Ready(Some("US".to_string())),
348            Pin::new(&mut future_a).poll(&mut context),
349            "for future a"
350        );
351        assert_eq!(
352            Poll::Ready(Some("US".to_string())),
353            Pin::new(&mut future_b).poll(&mut context),
354            "for future b"
355        );
356    }
357
358    #[fuchsia::test(allow_stalls = false)]
359    async fn single_watcher_is_not_woken_again_after_future_is_ready() {
360        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
361        let path = temp_dir.path().join("regulatory_region.json");
362        let hub = PubSubHub::new(path);
363        let (waker, count) = new_count_waker();
364        let mut context = Context::from_waker(&waker);
365        let mut future = hub.watch_for_change(Option::<String>::None);
366        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
367
368        // Publish an update, which resolves `future`.
369        hub.publish("US");
370        assert_eq!(1, count.get());
371        assert_eq!(Poll::Ready(Some("US".to_string())), Pin::new(&mut future).poll(&mut context));
372
373        // Further updates should leave `count` unchanged, since they should not wake `waker`.
374        hub.publish("SU");
375        assert_eq!(1, count.get());
376    }
377
378    #[fuchsia::test(allow_stalls = false)]
379    async fn second_watcher_is_woken_for_second_update() {
380        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
381        let path = temp_dir.path().join("regulatory_region.json");
382        let hub = PubSubHub::new(path);
383        let (waker, count) = new_count_waker();
384        let mut context = Context::from_waker(&waker);
385        let mut future = hub.watch_for_change(Option::<String>::None);
386        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
387
388        // Publish first update, which resolves `future`.
389        hub.publish("US");
390        assert_eq!(1, count.get());
391        assert_eq!(Poll::Ready(Some("US".to_string())), Pin::new(&mut future).poll(&mut context));
392
393        // Create a new `future`, and verify that a second update resolves the new `future`.
394        let mut future = hub.watch_for_change(Some("US"));
395        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
396        hub.publish("SU");
397        assert!(count.get() > 1, "Count should be >1, but is {}", count.get());
398        assert_eq!(Poll::Ready(Some("SU".to_string())), Pin::new(&mut future).poll(&mut context));
399    }
400
401    #[fuchsia::test(allow_stalls = false)]
402    async fn multiple_polls_of_single_watcher_do_not_cause_multiple_wakes_when_waker_is_reused() {
403        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
404        let path = temp_dir.path().join("regulatory_region.json");
405        let hub = PubSubHub::new(path);
406        let (waker, count) = new_count_waker();
407        let mut context = Context::from_waker(&waker);
408        let mut future = hub.watch_for_change(Option::<String>::None);
409        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
410        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context));
411
412        // Publish an update, which resolves `future`.
413        hub.publish("US");
414        assert_eq!(1, count.get());
415    }
416
417    #[fuchsia::test(allow_stalls = false)]
418    async fn multiple_polls_of_single_watcher_do_not_cause_multiple_wakes_when_waker_is_replaced() {
419        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
420        let path = temp_dir.path().join("regulatory_region.json");
421        let hub = PubSubHub::new(path);
422        let (waker_a, wake_count_a) = new_count_waker();
423        let (waker_b, wake_count_b) = new_count_waker();
424        let mut context_a = Context::from_waker(&waker_a);
425        let mut context_b = Context::from_waker(&waker_b);
426        let mut future = hub.watch_for_change(Option::<String>::None);
427        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context_a));
428        assert_eq!(Poll::Pending, Pin::new(&mut future).poll(&mut context_b));
429
430        // Publish an update, which resolves `future`.
431        hub.publish("US");
432        assert_eq!(0, wake_count_a.get());
433        assert_eq!(1, wake_count_b.get());
434    }
435
436    #[test]
437    fn get_value_is_none() {
438        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
439        let path = temp_dir.path().join("regulatory_region.json");
440        let hub = PubSubHub::new(path);
441        assert_eq!(None, hub.get_value());
442    }
443
444    #[test]
445    fn get_value_is_some() {
446        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
447        let path = temp_dir.path().join("regulatory_region.json");
448        let hub = PubSubHub::new(path);
449        hub.publish("US");
450        assert_eq!(Some("US".to_string()), hub.get_value());
451    }
452
453    #[test]
454    fn published_value_is_saved_and_loaded_on_creation() {
455        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
456        let path = temp_dir.path().join("regulatory_region.json");
457        let hub = PubSubHub::new(path.to_path_buf());
458        assert_eq!(hub.get_value(), None);
459        hub.publish("WW");
460        assert_eq!(hub.get_value(), Some("WW".to_string()));
461
462        // Create a new PubSubHub with the same storage path and verify that the initial value is
463        // the last thing published to previous PubSubHub.
464        let hub = PubSubHub::new(path.to_path_buf());
465        assert_eq!(hub.get_value(), Some("WW".to_string()));
466
467        // Verify that the files is unaffected.
468        let file = File::open(&path).expect("Failed to open file");
469        assert_matches!(
470            serde_json::from_reader(io::BufReader::new(file)),
471            Ok(RegulatoryRegion{ region_code }) if region_code.as_str() == "WW"
472        );
473    }
474
475    #[test]
476    fn publishing_over_previously_saved_value_overwrites_cache() {
477        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
478        let path = temp_dir.path().join("regulatory_region.json");
479
480        // Write some value to the cache.
481        let cache_val = RegulatoryRegion { region_code: "WW".to_string() };
482        let file = File::create(&path).expect("failed to create file");
483        serde_json::to_writer(io::BufWriter::new(file), &cache_val)
484            .expect("Failed to write JSON to file");
485
486        // Check that PubSubHub loads the correct value.
487        let hub = PubSubHub::new(path.to_path_buf());
488        assert_eq!(hub.get_value(), Some("WW".to_string()));
489
490        // Publish a new value and check that the file has the new value.
491        hub.publish("US");
492        let file = File::open(&path).expect("Failed to open file");
493        assert_matches!(
494            serde_json::from_reader(io::BufReader::new(file)),
495            Ok(RegulatoryRegion{ region_code }) if region_code.as_str() == "US"
496        );
497        let hub = PubSubHub::new(path.to_path_buf());
498        assert_eq!(hub.get_value(), Some("US".to_string()));
499    }
500
501    #[test]
502    fn load_as_none_if_cache_file_is_bad() {
503        let temp_dir = TempDir::new_in("/cache/").expect("failed to create temporary directory");
504        let path = temp_dir.path().join("regulatory_region.json");
505        assert!(!path.exists());
506        let mut file = File::create(&path).expect("failed to create file");
507        let bad_contents = b"{\"region_code\": ";
508        file.write_all(bad_contents).expect("failed to write to file");
509        file.flush().expect("failed to flush file");
510
511        // Check that PubSubHub is initialized with an unset value.
512        let hub = PubSubHub::new(path.to_path_buf());
513        assert_eq!(hub.get_value(), None);
514
515        // Check that the bad file was deleted.
516        assert_matches!(File::open(&path), Err(io_err) if io_err.kind() == io::ErrorKind::NotFound);
517    }
518}