1 files changed, 49 insertions, 53 deletions

diff --git a/crates/tor-dirmgr/src/event.rs b/crates/tor-dirmgr/src/event.rs
index cb8e2184d..c6dc27693 100644
--- a/crates/tor-dirmgr/src/event.rs
+++ b/crates/tor-dirmgr/src/event.rs
@@ -18,6 +18,7 @@ use std::{
 
 use educe::Educe;
 use futures::{stream::Stream, Future, StreamExt};
+use itertools::chain;
 use time::OffsetDateTime;
 use tor_basic_utils::skip_fmt;
 use tor_netdir::DirEvent;
@@ -457,9 +458,6 @@ impl DirBootstrapStatus {
     }
 
     /// Return the next DirStatus, if there is one.
-    ///
-    /// Testing-only.
-    #[cfg(test)]
     fn next(&self) -> Option<&DirStatus> {
         match &self.0 {
             StatusEnum::Replacing { next, .. } => Some(&next.status),
@@ -467,6 +465,21 @@ impl DirBootstrapStatus {
         }
     }
 
+    /// Return the contained `DirStatus`es, in order: `current`, then `next`
+    fn statuses(&self) -> impl Iterator<Item = &DirStatus> + DoubleEndedIterator {
+        chain!(self.current(), self.next(),)
+    }
+
+    /// Return the contained `StatusEntry`s mutably, in order: `current`, then `next`
+    fn entries_mut(&mut self) -> impl Iterator<Item = &mut StatusEntry> + DoubleEndedIterator {
+        let (current, next) = match &mut self.0 {
+            StatusEnum::NoActivity => (None, None),
+            StatusEnum::Single { current } => (Some(current), None),
+            StatusEnum::Replacing { current, next } => (Some(current), Some(next)),
+        };
+        chain!(current, next,)
+    }
+
     /// Return the fraction of completion for directory download, in a form
     /// suitable for a progress bar at some particular time.
     ///
@@ -476,15 +489,10 @@ impl DirBootstrapStatus {
     /// Callers _should not_ depend on the specific meaning of any particular
     /// fraction; we may change these fractions in the future.
     pub fn frac_at(&self, when: SystemTime) -> f32 {
-        match &self.0 {
-            StatusEnum::NoActivity => 0.0,
-            StatusEnum::Single { current } => current.status.frac_at(when).unwrap_or(0.0),
-            StatusEnum::Replacing { current, next } => current
-                .status
-                .frac_at(when)
-                .or_else(|| next.status.frac_at(when))
-                .unwrap_or(0.0),
-        }
+        self.statuses()
+            .filter_map(|st| st.frac_at(when))
+            .next()
+            .unwrap_or(0.0)
     }
 
     /// Return true if this status indicates that we have a current usable
@@ -508,55 +516,43 @@ impl DirBootstrapStatus {
             }
         }
 
-        match &self.0 {
-            // We're not trying to fetch anything, so it can't be blocked.
-            StatusEnum::NoActivity => None,
-            // We have only one attempt: its blockage is the only relevant one.
-            StatusEnum::Single { current } => current.status.blockage(),
-            // We know about two attempts: any blockage in "current" is more
-            // serious.
-            StatusEnum::Replacing { current, next } => {
-                current.status.blockage().or_else(|| next.status.blockage())
-            }
-        }
+        // Any blockage in "current" is more serious, so return that if there is one
+        self.statuses().filter_map(|st| st.blockage()).next()
     }
 
     /// Return the appropriate DirStatus for `AttemptId`, constructing it if
     /// necessary.
     ///
     /// Return None if all relevant attempts are more recent than this Id.
+    #[allow(clippy::search_is_some)] // tpo/core/arti/-/merge_requests/599#note_2816368
     fn mut_status_for(&mut self, attempt_id: AttemptId) -> Option<&mut DirStatus> {
-        // First, we add a status for this attempt_id if appropriate.
-        //
-        // TODO: should make sure that the compiler is smart enough to optimize
-        // this mem::take() and replacement away, and turn it into a conditional
-        // replacement?
-        self.0 = match std::mem::take(&mut self.0) {
-            StatusEnum::NoActivity => StatusEnum::Single {
-                current: StatusEntry::new(attempt_id),
-            },
-            StatusEnum::Single { current } if current.id < attempt_id => StatusEnum::Replacing {
-                current,
-                next: StatusEntry::new(attempt_id),
-            },
-            StatusEnum::Replacing { current, next } if next.id < attempt_id => {
-                StatusEnum::Replacing {
-                    current,
-                    next: StatusEntry::new(attempt_id),
-                }
-            }
-            other => other,
-        };
-
-        // Now return the correct status.
-        match &mut self.0 {
-            StatusEnum::Single { current } if current.id == attempt_id => Some(&mut current.status),
-            StatusEnum::Replacing { current, .. } if current.id == attempt_id => {
-                Some(&mut current.status)
-            }
-            StatusEnum::Replacing { next, .. } if next.id == attempt_id => Some(&mut next.status),
-            _ => None,
+        // First, ensure that we have a *recent enough* attempt
+        // Look for the latest attempt, and see if it's new enough; if not, start a new one.
+        if self
+            .entries_mut()
+            .rev()
+            .take(1)
+            .find(|entry| entry.id >= attempt_id)
+            .is_none()
+        {
+            let current = match std::mem::take(&mut self.0) {
+                StatusEnum::NoActivity => None,
+                StatusEnum::Single { current } => Some(current),
+                StatusEnum::Replacing { current, .. } => Some(current),
+            };
+            // If we have a `current` already, we keep it, and restart `next`.
+            let next = StatusEntry::new(attempt_id);
+            self.0 = match current {
+                None => StatusEnum::Single { current: next },
+                Some(current) => StatusEnum::Replacing { current, next },
+            };
         }
+
+        // Find the entry with `attempt_id` and return it.
+        // (Despite the above, there might not be one: maybe `attempt_id` is old.)
+        self.entries_mut()
+            .find(|entry| entry.id == attempt_id)
+            .map(|entry| &mut entry.status)
     }
 
     /// If the "next" status is usable, replace the current status with it.