refactor wait_for_probe_reboot into its own function

Cargo.toml

@@ -14,6 +14,9 @@ anyhow = "1.0"
 thiserror = "1.0"
 indicatif = "0.16"
 
+[profile.release]
+lto = "fat"
+
 [patch.crates-io]
 # Required until dfu-libusb updates its dfu-core dependency.
 dfu-core = { git = "https://github.com/dfu-rs/dfu-core", rev = "refs/pull/9/head" }

src/bmp.rs

@@ -401,50 +401,18 @@ impl BlackmagicProbeDevice
         Ok(())
     }
 
+    /// Requests the Black Magic Probe to detach, and re-initializes this struct with the new
+    /// device.
     pub fn detach_and_enumerate(&mut self) -> Result<(), BmputilError>
     {
-
         // Make sure we have our serial number before we try to detach,
         // so that we can find the DFU-mode device when it re-enumerates.
         let serial = self.serial_number()?.to_string();
         unsafe { self.request_detach()? };
 
-        // Give libusb a moment to at least register the disconnect before we try to do anything
-        // else.
-        thread::sleep(Duration::from_millis(200));
-
         // Now that we've detached, try to find the device again with the same serial number.
 
-        let matcher = BlackmagicProbeMatcher {
-            index: None,
-            serial: Some(serial),
-            port: None,
-        };
-
-        let mut dev = find_matching_probes(&matcher).pop_single("download");
-
-        // TODO: make this configurable?
-        const TIMEOUT: Duration = Duration::from_secs(5);
-        let start = Instant::now();
-
-        while let Err(BmputilError::DeviceNotFoundError) = dev {
-
-            // If it's been more than our timeout length, error out.
-            if start.duration_since(Instant::now()) > TIMEOUT {
-                error!(
-                    "Timed-out waiting for Black Magic Probe device to re-enumerate in DFU mode!"
-                );
-                return Err(BmputilError::DeviceReconfigureError { source: None });
-            }
-
-            // Wait 200 milliseconds between checks. Hardware is a bottleneck and we
-            // don't need to peg the CPU waiting for it to come back up.
-            // TODO: make this configurable and/or optimize?
-            thread::sleep(Duration::from_millis(200));
-            dev = find_matching_probes(&matcher).pop_single("download");
-        }
-
-        let dev = dev?;
+        let dev = wait_for_probe_reboot(&serial, Duration::from_secs(5), "flash")?;
 
         // If we've made it here, then we have successfully re-found the device.
         // Re-initialize this structure from the new data.
@@ -696,6 +664,18 @@ impl BlackmagicProbeMatchResults
 
         Ok(self.found.remove(0))
     }
+
+    /// Like `pop_single()`, but does not print helpful diagnostics for edge cases.
+    pub(crate) fn pop_single_silent(&mut self) -> Result<BlackmagicProbeDevice, BmputilError>
+    {
+        if self.found.len() > 1 {
+            return Err(BmputilError::TooManyDevicesError);
+        } else if self.found.is_empty() {
+            return Err(BmputilError::DeviceNotFoundError);
+        }
+
+        return Ok(self.found.remove(0));
+    }
 }
 
 
@@ -839,6 +819,49 @@ pub fn find_matching_probes(matcher: &BlackmagicProbeMatcher) -> BlackmagicProbe
 }
 
 
+pub fn wait_for_probe_reboot(serial: &str, timeout: Duration, operation: &str) -> Result<BlackmagicProbeDevice, BmputilError>
+{
+    let silence_timeout = timeout / 2;
+
+    let matcher = BlackmagicProbeMatcher {
+        index: None,
+        serial: Some(serial.to_string()),
+        port: None,
+    };
+
+    let start = Instant::now();
+
+    let mut dev = find_matching_probes(&matcher).pop_single_silent();
+
+    while let Err(BmputilError::DeviceNotFoundError) = dev {
+
+        // If it's been more than the timeout length, error out.
+        if start.duration_since(Instant::now()) > timeout {
+            error!(
+                "Timed-out waiting for Black Magic Probe to re-enumerate!"
+            );
+            return Err(BmputilError::DeviceRebootError { source: None });
+        }
+
+        // Wait 200 milliseconds between checks. Hardware is a bottleneck and we
+        // don't need to peg the CPU waiting for it to come back up.
+        // TODO: make this configurable and/or optimize?
+        thread::sleep(Duration::from_millis(200));
+
+        // If we've been trying for over half the full timeout, start logging warnings.
+        if start.duration_since(Instant::now()) > silence_timeout {
+            dev = find_matching_probes(&matcher).pop_single(operation);
+        } else {
+            dev = find_matching_probes(&matcher).pop_single_silent();
+        }
+    }
+
+    let dev = dev?;
+
+    Ok(dev)
+}
+
+
 pub struct BmpVidPid;
 impl BmpVidPid
 {

src/error.rs

@@ -6,7 +6,7 @@ use thiserror::Error;
 #[allow(dead_code)] // XXX FIXME
 pub enum BmputilError
 {
-    #[error("Failed to read firmware file")]
+    #[error("Failed to read firmware file {filename}")]
     FirmwareFileIOError
     {
         #[source]
@@ -46,7 +46,6 @@ pub enum BmputilError
         context: String,
     },
 
-    #[allow(dead_code)] // FIXME: this will presumably be used once we have more dynamic detection of re-enumeration.
     #[error("Blackmagic Probe device did not re-enumerate after requesting to switch to DFU mode")]
     DeviceReconfigureError
     {
@@ -78,6 +77,9 @@ pub enum BmputilError
 
     #[error("Other/unhandled libusb error (please report this so we can add better handling!)")]
     LibusbError(#[from] rusb::Error),
+
+    #[error("Other/unhandled dfu_libusb error (please report this so we can add better error handling!")]
+    DfuLibusbError(#[from] dfu_libusb::Error),
 }
 
 

src/main.rs

@@ -1,5 +1,6 @@
-use std::str::FromStr;
 use std::rc::Rc;
+use std::str::FromStr;
+use std::time::Duration;
 
 use clap::{Command, Arg, ArgMatches};
 
@@ -58,6 +59,7 @@ fn flash(matches: &ArgMatches) -> Result<(), BmputilError>
     let matcher = BlackmagicProbeMatcher::from_clap_matches(matches);
     let mut results = find_matching_probes(&matcher);
     let mut dev: BlackmagicProbeDevice = results.pop_single("flash")?;
+    let serial = dev.serial_number()?.to_string();
 
     println!("Found: {}", dev);
 
@@ -78,12 +80,55 @@ fn flash(matches: &ArgMatches) -> Result<(), BmputilError>
     let progress_bar = Rc::new(progress_bar);
     let enclosed = Rc::clone(&progress_bar);
 
-    dev.download(firmware_file, file_size, move |delta| {
+    match dev.download(firmware_file, file_size, move |delta| {
         enclosed.inc(delta as u64);
-    })?;
+    }) {
+        Ok(v) => Ok(v),
+        Err(BmputilError::DfuLibusbError(dfu_libusb::Error::Io(source))) => Err(BmputilError::FirmwareFileIOError {
+            source,
+            filename: filename.to_string(),
+        }),
+        Err(e) => Err(e),
+    }?;
 
     progress_bar.finish();
 
+    // Now that we've flashed, try and re-enumerate the device one more time.
+    let mut dev = bmp::wait_for_probe_reboot(&serial, Duration::from_secs(5), "flash")
+        .map_err(|e| {
+            error!("Black Magic Probe did not re-enumerate after flashing! Invalid firmware?");
+            e
+        })?;
+
+    let languages = dev
+        .handle()
+        .read_languages(Duration::from_secs(2))
+        .map_err(|e| {
+            error!("Error reading firmware version after flash! Invalid firmware?");
+            e
+        })?;
+
+    let desc = dev.device().device_descriptor().unwrap();
+
+    let product_string = dev
+        .handle()
+        .read_product_string(
+            *languages.first().unwrap(),
+            &desc,
+            Duration::from_secs(2),
+        )
+        .map_err(|e| {
+            error!("Error reading firmware version after flash! Invalid firmware?");
+            e
+        })?;
+
+    let version_string = product_string
+        .chars()
+        .skip("Black Magic Probe ".len())
+        .collect::<String>();
+
+    println!("Black Magic Probe successfully rebooted into firmware version {}", version_string);
+
     Ok(())
 }