MultiProgress: prune 'zombie' progress bars; fixes #426

src/draw_target.rs

@@ -187,6 +187,10 @@ impl ProgressDrawTarget {
             _ => None,
         }
     }
+
+    pub(crate) fn last_line_count(&mut self) -> Option<&mut usize> {
+        self.kind.last_line_count()
+    }
 }
 
 #[derive(Debug)]
@@ -209,6 +213,20 @@ enum TargetKind {
     },
 }
 
+impl TargetKind {
+    fn last_line_count(&mut self) -> Option<&mut usize> {
+        match self {
+            TargetKind::Term {
+                last_line_count, ..
+            } => Some(last_line_count),
+            TargetKind::TermLike {
+                last_line_count, ..
+            } => Some(last_line_count),
+            _ => None,
+        }
+    }
+}
+
 pub(crate) enum Drawable<'a> {
     Term {
         term: &'a Term,

src/multi.rs

@@ -1,9 +1,11 @@
+use std::fmt::{Debug, Formatter};
 use std::io;
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, Mutex, RwLock, Weak};
 use std::time::Instant;
 
 use crate::draw_target::{DrawState, DrawStateWrapper, ProgressDrawTarget};
 use crate::progress_bar::ProgressBar;
+use crate::state::BarState;
 
 /// Manages multiple progress bars from different threads
 #[derive(Debug, Clone)]
@@ -128,8 +130,11 @@ impl MultiProgress {
     }
 
     fn internalize(&self, location: InsertLocation, pb: ProgressBar) -> ProgressBar {
-        let idx = self.state.write().unwrap().insert(location);
+        let mut state = self.state.write().unwrap();
+
+        let idx = state.insert(location);
         pb.set_draw_target(ProgressDrawTarget::new_remote(self.state.clone(), idx));
+        state.elements.get_mut(idx).unwrap().pb = pb.weak_bar_state();
         pb
     }
 
@@ -169,8 +174,8 @@ impl MultiProgress {
 pub(crate) struct MultiState {
     /// The collection of states corresponding to progress bars
     /// the state is None for bars that have not yet been drawn or have been removed
-    draw_states: Vec<Option<DrawState>>,
-    /// Set of removed bars, should have corresponding `None` elements in the `draw_states` vector
+    elements: Vec<MultiStateElement>,
+    /// Set of removed bars, should have corresponding `None` elements in the `elements` vector
     free_set: Vec<usize>,
     /// Indices to the `draw_states` to maintain correct visual order
     ordering: Vec<usize>,
@@ -187,7 +192,7 @@ pub(crate) struct MultiState {
 impl MultiState {
     fn new(draw_target: ProgressDrawTarget) -> Self {
         Self {
-            draw_states: vec![],
+            elements: vec![],
             free_set: vec![],
             ordering: vec![],
             draw_target,
@@ -203,6 +208,32 @@ impl MultiState {
         extra_lines: Option<Vec<String>>,
         now: Instant,
     ) -> io::Result<()> {
+        // Reap all consecutive 'zombie' progress bars from head of elements
+        let mut zombies = vec![];
+        let mut adjust = 0;
+        for index in self.ordering.iter() {
+            let element = &self.elements[*index];
+            if element.is_zombie {
+                zombies.push(*index);
+                adjust += element
+                    .draw_state
+                    .as_ref()
+                    .map(|d| d.lines.len())
+                    .unwrap_or_default();
+            } else {
+                break;
+            }
+        }
+
+        for index in zombies {
+            self.remove_idx(index);
+        }
+
+        // Adjust last line count so we don't clear too many lines
+        if let Some(last_line_count) = self.draw_target.last_line_count() {
+            *last_line_count -= adjust;
+        }
+
         let orphan_lines_count = self.orphan_lines.len();
         force_draw |= orphan_lines_count > 0;
         let mut drawable = match self.draw_target.drawable(force_draw, now) {
@@ -222,9 +253,15 @@ impl MultiState {
         draw_state.lines.append(&mut self.orphan_lines);
 
         for index in self.ordering.iter() {
-            if let Some(state) = &self.draw_states[*index] {
+            let element = &mut self.elements[*index];
+            if let Some(state) = &element.draw_state {
                 draw_state.lines.extend_from_slice(&state.lines[..]);
             }
+
+            // Mark the dead progress bar as a zombie - will be reaped on next draw
+            if element.pb.upgrade().is_none() {
+                element.is_zombie = true;
+            }
         }
 
         drop(draw_state);
@@ -244,20 +281,14 @@ impl MultiState {
     }
 
     pub(crate) fn draw_state(&mut self, idx: usize) -> DrawStateWrapper<'_> {
-        let (states, orphans) = (&mut self.draw_states, &mut self.orphan_lines);
-        let state = match states.get_mut(idx) {
-            Some(Some(draw_state)) => draw_state,
-            Some(inner) => {
-                *inner = Some(DrawState::default());
-                let state = inner.as_mut().unwrap();
-                state.move_cursor = self.move_cursor;
-                state.alignment = self.alignment;
-                state
-            }
-            _ => unreachable!(),
-        };
+        let element = self.elements.get_mut(idx).unwrap();
+        let state = element.draw_state.get_or_insert(DrawState {
+            move_cursor: self.move_cursor,
+            alignment: self.alignment,
+            ..Default::default()
+        });
 
-        DrawStateWrapper::for_multi(state, orphans)
+        DrawStateWrapper::for_multi(state, &mut self.orphan_lines)
     }
 
     pub(crate) fn is_hidden(&self) -> bool {
@@ -278,12 +309,12 @@ impl MultiState {
     fn insert(&mut self, location: InsertLocation) -> usize {
         let idx = match self.free_set.pop() {
             Some(idx) => {
-                self.draw_states[idx] = None;
+                self.elements[idx] = MultiStateElement::default();
                 idx
             }
             None => {
-                self.draw_states.push(None);
-                self.draw_states.len() - 1
+                self.elements.push(MultiStateElement::default());
+                self.elements.len() - 1
             }
         };
 
@@ -328,7 +359,7 @@ impl MultiState {
             return;
         }
 
-        self.draw_states[idx].take();
+        self.elements[idx] = MultiStateElement::default();
         self.free_set.push(idx);
         self.ordering.retain(|&x| x != idx);
 
@@ -340,7 +371,23 @@ impl MultiState {
     }
 
     fn len(&self) -> usize {
-        self.draw_states.len() - self.free_set.len()
+        self.elements.len() - self.free_set.len()
+    }
+}
+
+#[derive(Default)]
+struct MultiStateElement {
+    draw_state: Option<DrawState>,
+    pb: Weak<Mutex<BarState>>,
+    is_zombie: bool,
+}
+
+impl Debug for MultiStateElement {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MultiStateElement")
+            .field("draw_state", &self.draw_state)
+            .field("is_zombie", &self.is_zombie)
+            .finish_non_exhaustive()
     }
 }
 
@@ -422,19 +469,19 @@ mod tests {
 
         let state = mp.state.read().unwrap();
         // the removed place for p1 is reused
-        assert_eq!(state.draw_states.len(), 4);
+        assert_eq!(state.elements.len(), 4);
         assert_eq!(state.len(), 3);
 
         // free_set may contain 1 or 2
         match state.free_set.last() {
             Some(1) => {
                 assert_eq!(state.ordering, vec![0, 2, 3]);
-                assert!(state.draw_states[1].is_none());
+                assert!(state.elements[1].draw_state.is_none());
                 assert_eq!(p4.index().unwrap(), 2);
             }
             Some(2) => {
                 assert_eq!(state.ordering, vec![0, 1, 3]);
-                assert!(state.draw_states[2].is_none());
+                assert!(state.elements[2].draw_state.is_none());
                 assert_eq!(p4.index().unwrap(), 1);
             }
             _ => unreachable!(),
@@ -534,10 +581,10 @@ mod tests {
 
         let state = mp.state.read().unwrap();
         // the removed place for p1 is reused
-        assert_eq!(state.draw_states.len(), 2);
+        assert_eq!(state.elements.len(), 2);
         assert_eq!(state.free_set.len(), 1);
         assert_eq!(state.len(), 1);
-        assert!(state.draw_states[0].is_none());
+        assert!(state.elements[0].draw_state.is_none());
         assert_eq!(state.free_set.last(), Some(&0));
 
         assert_eq!(state.ordering, vec![1]);

src/progress_bar.rs

@@ -250,6 +250,10 @@ impl ProgressBar {
         }
     }
 
+    pub(crate) fn weak_bar_state(&self) -> Weak<Mutex<BarState>> {
+        Arc::downgrade(&self.state)
+    }
+
     /// Resets the ETA calculation
     ///
     /// This can be useful if the progress bars made a large jump or was paused for a prolonged