Allow specifying number of worker threads used to run an enclave

intel-sgx/enclave-runner/src/command.rs

@@ -23,6 +23,7 @@ pub struct Command {
     usercall_ext: Option<Box<dyn UsercallExtension>>,
     forward_panics: bool,
     cmd_args: Vec<Vec<u8>>,
+    num_worker_threads: usize,
 }
 
 impl MappingInfo for Command {
@@ -45,6 +46,7 @@ impl Command {
         usercall_ext: Option<Box<dyn UsercallExtension>>,
         forward_panics: bool,
         cmd_args: Vec<Vec<u8>>,
+        num_worker_threads: usize,
     ) -> Command {
         let main = tcss.remove(0);
         Command {
@@ -55,6 +57,7 @@ impl Command {
             usercall_ext,
             forward_panics,
             cmd_args,
+            num_worker_threads,
         }
     }
 
@@ -63,6 +66,6 @@ impl Command {
     }
 
     pub fn run(self) -> Result<(), Error> {
-        EnclaveState::main_entry(self.main, self.threads, self.usercall_ext, self.forward_panics, self.cmd_args)
+        EnclaveState::main_entry(self.main, self.threads, self.usercall_ext, self.forward_panics, self.cmd_args, self.num_worker_threads)
     }
 }

intel-sgx/enclave-runner/src/loader.rs

@@ -69,6 +69,7 @@ pub struct EnclaveBuilder<'a> {
     hash_enclave: Option<Box<dyn FnOnce(&mut EnclaveSource<'_>) -> Result<EnclaveHash, anyhow::Error>>>,
     forward_panics: bool,
     cmd_args: Option<Vec<Vec<u8>>>,
+    num_worker_threads: Option<usize>,
 }
 
 #[derive(Debug, ThisError)]
@@ -146,6 +147,7 @@ impl<'a> EnclaveBuilder<'a> {
             hash_enclave: None,
             forward_panics: false,
             cmd_args: None,
+            num_worker_threads: None,
         };
 
         let _ = ret.coresident_signature();
@@ -298,14 +300,22 @@ impl<'a> EnclaveBuilder<'a> {
     /// Adding command arguments and then calling [`build_library`] will cause
     /// a panic.
     ///
-    /// [`Command`]: struct.Command.html
     /// [`build_library`]: struct.EnclaveBuilder.html#method.build_library
     pub fn arg<S: AsRef<[u8]>>(&mut self, arg: S) -> &mut Self {
         let arg = arg.as_ref().to_owned();
         self.initialized_args_mut().push(arg);
         self
     }
 
+    /// Sets the number of wroker threads used to run the enclave.
+    ///
+    /// **NOTE:** This is only applicable to [`Command`] enclaves.
+    /// Setting this and then calling [`build_library`](Self::build_library) will cause a panic.
+    pub fn num_worker_threads(&mut self, num_worker_threads: usize) -> &mut Self {
+        self.num_worker_threads = Some(num_worker_threads);
+        self
+    }
+
     fn load<T: Load>(
         mut self,
         loader: &mut T,
@@ -332,21 +342,26 @@ impl<'a> EnclaveBuilder<'a> {
     }
 
     pub fn build<T: Load>(mut self, loader: &mut T) -> Result<Command, anyhow::Error> {
+        let num_cpus = num_cpus::get();
+        let num_worker_threads = self.num_worker_threads.unwrap_or(num_cpus);
+        anyhow::ensure!(num_worker_threads > 0, "`num_worker_threads` cannot be zero");
+        anyhow::ensure!(num_worker_threads <= num_cpus, "`num_worker_threads` cannot exceed number of logical cores (`num_cpus::get()`)");
+
         self.initialized_args_mut();
         let args = self.cmd_args.take().unwrap_or_default();
         let c = self.usercall_ext.take();
         self.load(loader)
-            .map(|(t, a, s, fp)| Command::internal_new(t, a, s, c, fp, args))
+            .map(|(t, a, s, fp)| Command::internal_new(t, a, s, c, fp, args, num_worker_threads))
     }
 
-    /// Panics if you have previously called [`arg`] or [`args`].
+    /// Panics if you have previously called [`arg`], [`args`], or [`num_worker_threads`].
     ///
     /// [`arg`]: struct.EnclaveBuilder.html#method.arg
     /// [`args`]: struct.EnclaveBuilder.html#method.args
+    /// [`num_worker_threads`]: Self::num_worker_threads()
     pub fn build_library<T: Load>(mut self, loader: &mut T) -> Result<Library, anyhow::Error> {
-        if self.cmd_args.is_some() {
-            panic!("Command arguments do not apply to Library enclaves.");
-        }
+        assert!(self.cmd_args.is_none(), "Command arguments do not apply to Library enclaves.");
+        assert!(self.num_worker_threads.is_none(), "`num_worker_threads` cannot be specified for Library enclaves.");
         let c = self.usercall_ext.take();
         self.load(loader)
             .map(|(t, a, s, fp)| Library::internal_new(t, a, s, c, fp))

intel-sgx/enclave-runner/src/usercalls/mod.rs

@@ -1086,7 +1086,9 @@ impl EnclaveState {
         usercall_ext: Option<Box<dyn UsercallExtension>>,
         forward_panics: bool,
         cmd_args: Vec<Vec<u8>>,
+        num_of_worker_threads: usize,
     ) -> StdResult<(), anyhow::Error> {
+        assert!(num_of_worker_threads > 0, "worker_threads cannot be zero");
         let mut event_queues =
             FnvHashMap::with_capacity_and_hasher(threads.len() + 1, Default::default());
         let main = Self::event_queue_add_tcs(&mut event_queues, main);
@@ -1109,8 +1111,6 @@ impl EnclaveState {
             entry: CoEntry::Initial(main.tcs, argv as _, argc as _, 0, 0, 0),
         };
 
-        let num_of_worker_threads = num_cpus::get();
-
         let kind = EnclaveKind::Command(Command {
             panic_reason: Mutex::new(PanicReason {
                 primary_panic_reason: None,