From d8f63c22375fe23d7ef7595646a12a236acc8223 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9my=20Greinhofer?= <remy.greinhofer@gmail.com>
Date: Sat, 19 Mar 2022 20:03:04 -0500
Subject: [PATCH] Fix add with overflow panic

While working on adding a new example to showcase the use of indicatif
with async multi-progress bars, I faced a random panic bug:

`thread 'main' panicked at 'attempt to add with overflow',
       src/state.rs:424:40`

As a result this patch provides a new example and a fix for a bug it
exposed.
---
 Cargo.toml                   |  3 +-
 examples/async-multi-main.rs | 60 ++++++++++++++++++++++++++++++++++++
 src/state.rs                 |  9 +++++-
 3 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 examples/async-multi-main.rs

diff --git a/Cargo.toml b/Cargo.toml
index ce8c8901..d0cc8355 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,7 @@ exclude = ["screenshots/*"]
 
 [dependencies]
 console = { version = "0.15", default-features = false, features = ["ansi-parsing"] }
+futures = "0.3.21"
 number_prefix = "0.4"
 rayon = { version = "1.1", optional = true }
 tokio = { version = "1", optional = true, features = ["fs", "io-util"] }
@@ -24,7 +25,7 @@ vt100 = { version = "0.15.1", optional = true }
 once_cell = "1"
 rand = "0.8"
 structopt = "0.3"
-tokio = { version = "1", features = ["time", "rt"] }
+tokio = { version = "1", features = ["macros", "time", "rt", "rt-multi-thread"] }
 
 [features]
 default = ["unicode-width", "console/unicode-width"]
diff --git a/examples/async-multi-main.rs b/examples/async-multi-main.rs
new file mode 100644
index 00000000..6532139c
--- /dev/null
+++ b/examples/async-multi-main.rs
@@ -0,0 +1,60 @@
+//! Example of asynchronous multiple progress bars.
+//!
+//! The child bars are added to the main one. Once a child bar is complete it
+//! gets removed from the rendering and the main bar advances by 1 unit.
+//!
+//! Run with
+//!
+//! ```not_rust
+//! cargo run --example async-multi-main.rs
+//! ```
+//!
+use futures::stream::{self, StreamExt};
+use rand::{thread_rng, Rng};
+use std::sync::Arc;
+use tokio::time::{sleep, Duration};
+
+use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
+
+const MAX_CONCURRENT_ITEMS: usize = 5;
+
+#[tokio::main]
+async fn main() {
+    // Creates a new multi-progress object.
+    let multi = Arc::new(MultiProgress::new());
+    let style = ProgressStyle::with_template("{bar:40.green/yellow} {pos:>7}/{len:7}").unwrap();
+
+    // Create the main progress bar.
+    let mut rng = thread_rng();
+    let items = rng.gen_range(MAX_CONCURRENT_ITEMS..MAX_CONCURRENT_ITEMS * 3);
+    let main = Arc::new(multi.add(ProgressBar::new(items as u64).with_style(style.clone())));
+    main.tick();
+
+    // Add the child progress bars.
+    let _pbs = stream::iter(0..items)
+        .map(|_i| add_bar(main.clone(), multi.clone()))
+        .buffer_unordered(MAX_CONCURRENT_ITEMS)
+        .collect::<Vec<_>>()
+        .await;
+    main.finish_with_message("done");
+}
+
+async fn add_bar(main: Arc<ProgressBar>, multi: Arc<MultiProgress>) {
+    // Create a child bar and add it to the main one.
+    let mut rng = thread_rng();
+    let length: u64 = rng.gen_range(128..1024);
+    let sleep_ms: u64 = rng.gen_range(5..10);
+    let style = ProgressStyle::with_template("{bar:40.cyan/blue} {pos:>7}/{len:7}").unwrap();
+    let pb = multi.add(ProgressBar::new(length).with_style(style.clone()));
+
+    // Simulate some work.
+    for _ in 0..length {
+        pb.inc(1);
+        sleep(Duration::from_millis(sleep_ms)).await;
+    }
+    // Remove the bar once complete.
+    pb.finish_and_clear();
+
+    // Advance the main progress bar.
+    main.inc(1);
+}
diff --git a/src/state.rs b/src/state.rs
index 10396eb3..3da8c595 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -421,7 +421,14 @@ impl AtomicPosition {
         let (new, remainder) = ((diff / INTERVAL), (diff % INTERVAL));
         // We add `new` to `capacity`, subtract one for returning `true` from here,
         // then make sure it does not exceed a maximum of `MAX_BURST`.
-        capacity = Ord::min(MAX_BURST, capacity + new as u8 - 1);
+        capacity = Ord::min(
+            MAX_BURST,
+            capacity
+                .checked_add(new as u8)
+                .unwrap_or_default()
+                .checked_sub(1)
+                .unwrap_or_default(),
+        );
 
         // Then, we just store `capacity` and `prev` atomically for the next iteration
         self.capacity.store(capacity, Ordering::Release);