main: add --cpu-affinity option to pin VCPUs

This allows setting the affinity of the VCPU threads to specific host
CPUs.  Note that each individual CPU has its affinity set to the full
set of CPUs specified, so the host kernel may still reschedule VCPU
threads on whichever host CPUs it sees fit (within the specified set).

BUG=chromium:909793
TEST=build_test

Change-Id: I09b893901caf91368b64f5329a6e9f39027fef23
Signed-off-by: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1554865
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Chirantan Ekbote <chirantan@chromium.org>
Reviewed-by: Dylan Reid <dgreid@chromium.org>

1 files changed, 64 insertions, 0 deletions

diff --git a/sys_util/src/affinity.rs b/sys_util/src/affinity.rs
new file mode 100644
index 0000000..d166562
--- /dev/null
+++ b/sys_util/src/affinity.rs
@@ -0,0 +1,64 @@
+// Copyright 2019 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+//! Wrappers for CPU affinity functions.
+
+use std::iter::FromIterator;
+use std::mem;
+
+use libc::{cpu_set_t, sched_setaffinity, CPU_SET, CPU_SETSIZE, CPU_ZERO, EINVAL};
+
+use crate::{errno_result, Error, Result};
+
+// This is needed because otherwise the compiler will complain that the
+// impl doesn't reference any types from inside this crate.
+struct CpuSet(cpu_set_t);
+
+impl FromIterator<usize> for CpuSet {
+    fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
+        // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
+        let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
+        // Safe because we pass a valid cpuset pointer.
+        unsafe { CPU_ZERO(&mut cpuset) };
+        for cpu in cpus.into_iter() {
+            // Safe because we pass a valid cpuset pointer and cpu index.
+            unsafe { CPU_SET(cpu, &mut cpuset) };
+        }
+        CpuSet(cpuset)
+    }
+}
+
+/// Set the CPU affinity of the current thread to a given set of CPUs.
+///
+/// # Examples
+///
+/// Set the calling thread's CPU affinity so it will run on only CPUs
+/// 0, 1, 5, and 6.
+///
+/// ```
+/// # use sys_util::set_cpu_affinity;
+///   set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
+/// ```
+pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
+    let CpuSet(cpuset) = cpus
+        .into_iter()
+        .map(|cpu| {
+            if cpu < CPU_SETSIZE as usize {
+                Ok(cpu)
+            } else {
+                Err(Error::new(EINVAL))
+            }
+        })
+        .collect::<Result<CpuSet>>()?;
+
+    // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
+    // used for the duration of this call.
+    let res = unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) };
+
+    if res != 0 {
+        errno_result()
+    } else {
+        Ok(())
+    }
+}