1use 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
17pub struct PubSubHub {
19 inner: RefCell<PubSubHubInner>,
24 storage_path: PathBuf,
26}
27
28pub struct PubSubFuture<'a> {
30 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 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 write_region_code(new_value, &self.storage_path);
69 }
70
71 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#[derive(Debug, Deserialize, Serialize)]
91struct RegulatoryRegion {
92 region_code: String,
93}
94
95fn 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
123fn 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 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 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 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 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 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 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 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 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 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 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 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 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 let hub = PubSubHub::new(path.to_path_buf());
465 assert_eq!(hub.get_value(), Some("WW".to_string()));
466
467 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 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 let hub = PubSubHub::new(path.to_path_buf());
488 assert_eq!(hub.get_value(), Some("WW".to_string()));
489
490 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 let hub = PubSubHub::new(path.to_path_buf());
513 assert_eq!(hub.get_value(), None);
514
515 assert_matches!(File::open(&path), Err(io_err) if io_err.kind() == io::ErrorKind::NotFound);
517 }
518}