Merge pull request #2640 from subspace/farmer-support-cpu-core-ranges

Support CPU core ranges in farmer

Author: nazar-pc

crates/subspace-farmer/src/utils.rs

@@ -11,10 +11,9 @@ use futures::channel::oneshot::Canceled;
 use futures::future::Either;
 use rayon::{ThreadBuilder, ThreadPool, ThreadPoolBuildError, ThreadPoolBuilder};
 use std::future::Future;
-use std::num::{NonZeroUsize, ParseIntError};
+use std::num::NonZeroUsize;
 use std::ops::Deref;
 use std::pin::{pin, Pin};
-use std::str::FromStr;
 use std::task::{Context, Poll};
 use std::{io, thread};
 use thread_priority::{set_current_thread_priority, ThreadPriority};
@@ -273,16 +272,32 @@ pub fn all_cpu_cores() -> Vec<CpuCoreSet> {
 
 /// Parse space-separated set of groups of CPU cores (individual cores are coma-separated) into
 /// vector of CPU core sets that can be used for creation of plotting/replotting thread pools.
-pub fn parse_cpu_cores_sets(s: &str) -> Result<Vec<CpuCoreSet>, ParseIntError> {
+pub fn parse_cpu_cores_sets(
+    s: &str,
+) -> Result<Vec<CpuCoreSet>, Box<dyn std::error::Error + Send + Sync>> {
     #[cfg(feature = "numa")]
     let topology = hwlocality::Topology::new().map(std::sync::Arc::new).ok();
 
     s.split(' ')
         .map(|s| {
-            let cores = s
-                .split(',')
-                .map(usize::from_str)
-                .collect::<Result<Vec<usize>, _>>()?;
+            let mut cores = Vec::new();
+            for s in s.split(',') {
+                let mut parts = s.split('-');
+                let range_start = parts
+                    .next()
+                    .ok_or(
+                        "Bad string format, must be comma separated list of CPU cores or ranges",
+                    )?
+                    .parse()?;
+
+                if let Some(range_end) = parts.next() {
+                    let range_end = range_end.parse()?;
+
+                    cores.extend(range_start..=range_end);
+                } else {
+                    cores.push(range_start);
+                }
+            }
 
             Ok(CpuCoreSet {
                 cores,

crates/subspace-farmer/src/utils/tests.rs

@@ -1,4 +1,4 @@
-use crate::utils::run_future_in_dedicated_thread;
+use crate::utils::{parse_cpu_cores_sets, run_future_in_dedicated_thread};
 use std::future;
 use tokio::sync::oneshot;
 
@@ -46,3 +46,35 @@ fn run_future_in_dedicated_thread_tokio_on_drop() {
         ));
     });
 }
+
+#[test]
+fn test_parse_cpu_cores_sets() {
+    {
+        let cores = parse_cpu_cores_sets("0").unwrap();
+        assert_eq!(cores.len(), 1);
+        assert_eq!(cores[0].cores, vec![0]);
+    }
+    {
+        let cores = parse_cpu_cores_sets("0,1,2").unwrap();
+        assert_eq!(cores.len(), 1);
+        assert_eq!(cores[0].cores, vec![0, 1, 2]);
+    }
+    {
+        let cores = parse_cpu_cores_sets("0,1,2 4,5,6").unwrap();
+        assert_eq!(cores.len(), 2);
+        assert_eq!(cores[0].cores, vec![0, 1, 2]);
+        assert_eq!(cores[1].cores, vec![4, 5, 6]);
+    }
+    {
+        let cores = parse_cpu_cores_sets("0-2 4-6,7").unwrap();
+        assert_eq!(cores.len(), 2);
+        assert_eq!(cores[0].cores, vec![0, 1, 2]);
+        assert_eq!(cores[1].cores, vec![4, 5, 6, 7]);
+    }
+
+    assert!(parse_cpu_cores_sets("").is_err());
+    assert!(parse_cpu_cores_sets("a").is_err());
+    assert!(parse_cpu_cores_sets("0,").is_err());
+    assert!(parse_cpu_cores_sets("0,a").is_err());
+    assert!(parse_cpu_cores_sets("0 a").is_err());
+}