update_metadata: Introduce the Wave enum

Add the Wave enum which defines the wave for a given Updog seed as
either the Initial wave, a General wave, or the Last wave. The
associated functions has_started() and has_passed() use these types to
simplify the logic in jitter() and update_ready().

Signed-off-by: Samuel Mendoza-Jonas <samjonas@amazon.com>

workspaces/updater/update_metadata/src/lib.rs

@@ -14,6 +14,37 @@ use std::ops::Bound::{Excluded, Included};
 
 pub const MAX_SEED: u32 = 2048;
 
+#[derive(Debug, PartialEq, Eq)]
+pub enum Wave {
+    Initial {
+        end: DateTime<Utc>,
+    },
+    General {
+        start: DateTime<Utc>,
+        end: DateTime<Utc>,
+    },
+    Last {
+        start: DateTime<Utc>,
+    },
+}
+
+impl Wave {
+    pub fn has_started(&self) -> bool {
+        match self {
+            Self::Initial { .. } => true,
+            Self::General { start, .. } | Self::Last { start } => *start <= Utc::now(),
+        }
+    }
+
+    pub fn has_passed(&self) -> bool {
+        match self {
+            Self::Initial { end } => *end <= Utc::now(),
+            Self::General { end, .. } => *end <= Utc::now(),
+            Self::Last { start } => *start <= Utc::now(),
+        }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Images {
     pub boot: String,
@@ -48,49 +79,51 @@ impl Update {
     /// - The "0th" wave, which has an "end" time but no specified start time.
     /// - The last wave, which has a start time but no specified end time.
     /// - Nothing, if no waves are configured.
-    pub fn update_wave(&self, seed: u32) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
-        (
-            self.waves
-                .range((Included(0), Excluded(seed)))
-                .last()
-                .map(|(_, wave)| *wave),
-            self.waves
-                .range((Included(seed), Included(MAX_SEED)))
-                .next()
-                .map(|(_, wave)| *wave),
-        )
+    pub fn update_wave(&self, seed: u32) -> Option<Wave> {
+        let start = self
+            .waves
+            .range((Included(0), Excluded(seed)))
+            .last()
+            .map(|(_, wave)| *wave);
+        let end = self
+            .waves
+            .range((Included(seed), Included(MAX_SEED)))
+            .next()
+            .map(|(_, wave)| *wave);
+
+        match (start, end) {
+            (None, Some(end)) => Some(Wave::Initial { end }),
+            (Some(start), Some(end)) => Some(Wave::General { start, end }),
+            (Some(start), None) => Some(Wave::Last { start }),
+            _ => None,
+        }
     }
 
     pub fn update_ready(&self, seed: u32) -> bool {
         // Has this client's wave started
-        match self.update_wave(seed) {
-            // some wave with time bounds
-            (Some(start), Some(_)) => return start < Utc::now(),
-            // 0th wave with no minimum time
-            (None, Some(_)) => return true,
-            _ => (),
-        }
-
-        // Alternately have all waves passed
-        if let Some((_, wave)) = self.waves.iter().last() {
-            return *wave <= Utc::now();
+        if let Some(wave) = self.update_wave(seed) {
+            return wave.has_started();
         }
 
         // Or there are no waves
         true
     }
 
     pub fn jitter(&self, seed: u32) -> Option<DateTime<Utc>> {
-        let (start, end) = self.update_wave(seed);
-        if let Some(end) = end {
-            if end < Utc::now() {
-                // this wave has already passed
+        if let Some(wave) = self.update_wave(seed) {
+            if wave.has_passed() {
                 return None;
             }
-            let start = start.unwrap_or_else(Utc::now);
-            let mut rng = thread_rng();
-            if let Some(range) = end.timestamp().checked_sub(start.timestamp()) {
-                return Some(start + Duration::seconds(rng.gen_range(1, range)));
+            let bounds = match self.update_wave(seed) {
+                Some(Wave::Initial { end }) => Some((Utc::now(), end)),
+                Some(Wave::General { start, end }) => Some((start, end)),
+                Some(Wave::Last { start: _ }) | None => None,
+            };
+            if let Some((start, end)) = bounds {
+                let mut rng = thread_rng();
+                if let Some(range) = end.timestamp().checked_sub(start.timestamp()) {
+                    return Some(start + Duration::seconds(rng.gen_range(1, range)));
+                }
             }
         }
         None

workspaces/updater/updog/src/main.rs

@@ -570,7 +570,7 @@ mod tests {
     use chrono::Duration as TestDuration;
     use std::collections::BTreeMap;
     use std::str::FromStr;
-    use update_metadata::Images;
+    use update_metadata::{Images, Wave};
 
     #[test]
     fn test_manifest_json() {
@@ -845,43 +845,52 @@ mod tests {
         };
 
         // | ---- (100, "now") ---
-        u.waves.insert(100, Utc::now());
-        let (start, end) = u.update_wave(1);
-        assert!(start.is_none() && end.is_some(), "Expected to be 0th wave");
+        let first_bound = Utc::now();
+        u.waves.insert(100, first_bound);
+        assert!(
+            u.update_wave(1).unwrap() == Wave::Initial { end: first_bound },
+            "Expected to be 0th wave"
+        );
         assert!(u.jitter(1).is_none(), "Expected immediate update");
-        let (start, end) = u.update_wave(101);
         assert!(
-            start.is_some() && end.is_none(),
+            u.update_wave(101).unwrap() == Wave::Last { start: first_bound },
             "Expected to be final wave"
         );
         assert!(u.jitter(101).is_none(), "Expected immediate update");
 
         // | ---- (100, "now") ---- (200, "+1hr") ---
-        u.waves.insert(200, Utc::now() + TestDuration::hours(1));
-        let (start, end) = u.update_wave(1);
-        assert!(start.is_none() && end.is_some(), "Expected to be 0th wave");
+        let second_bound = Utc::now() + TestDuration::hours(1);
+        u.waves.insert(200, second_bound);
+        assert!(
+            u.update_wave(1).unwrap() == Wave::Initial { end: first_bound },
+            "Expected to be 0th wave"
+        );
         assert!(u.jitter(1).is_none(), "Expected immediate update");
 
-        let (start, end) = u.update_wave(100);
         assert!(
-            start.is_none() && end.is_some(),
+            u.update_wave(100).unwrap() == Wave::Initial { end: first_bound },
             "Expected to be 0th wave (just!)"
         );
         assert!(u.jitter(100).is_none(), "Expected immediate update");
 
-        let (start, end) = u.update_wave(150);
         assert!(
-            start.is_some() && end.is_some(),
+            u.update_wave(150).unwrap()
+                == Wave::General {
+                    start: first_bound,
+                    end: second_bound
+                },
             "Expected to be some bounded wave"
         );
         assert!(
             u.jitter(150).is_some(),
             "Expected to have to wait for update"
         );
 
-        let (start, end) = u.update_wave(201);
         assert!(
-            start.is_some() && end.is_none(),
+            u.update_wave(201).unwrap()
+                == Wave::Last {
+                    start: second_bound
+                },
             "Expected to be final wave"
         );
         assert!(u.jitter(201).is_none(), "Expected immediate update");