From cbe6600149b53dcb97272a1940850bfa2e45401e Mon Sep 17 00:00:00 2001 From: Arash Sahebolamri Date: Mon, 15 Jul 2024 13:04:01 -0700 Subject: [PATCH] Allow specifying number of worker threads used to run an enclave --- intel-sgx/enclave-runner/src/command.rs | 5 +++- intel-sgx/enclave-runner/src/loader.rs | 27 ++++++++++++++----- intel-sgx/enclave-runner/src/usercalls/mod.rs | 4 +-- 3 files changed, 27 insertions(+), 9 deletions(-) diff --git a/intel-sgx/enclave-runner/src/command.rs b/intel-sgx/enclave-runner/src/command.rs index 666ccb05..df09fb3b 100644 --- a/intel-sgx/enclave-runner/src/command.rs +++ b/intel-sgx/enclave-runner/src/command.rs @@ -23,6 +23,7 @@ pub struct Command { usercall_ext: Option>, forward_panics: bool, cmd_args: Vec>, + num_worker_threads: usize, } impl MappingInfo for Command { @@ -45,6 +46,7 @@ impl Command { usercall_ext: Option>, forward_panics: bool, cmd_args: Vec>, + 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) } } diff --git a/intel-sgx/enclave-runner/src/loader.rs b/intel-sgx/enclave-runner/src/loader.rs index daca85e3..6e523e5e 100644 --- a/intel-sgx/enclave-runner/src/loader.rs +++ b/intel-sgx/enclave-runner/src/loader.rs @@ -69,6 +69,7 @@ pub struct EnclaveBuilder<'a> { hash_enclave: Option) -> Result>>, forward_panics: bool, cmd_args: Option>>, + num_worker_threads: Option, } #[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,7 +300,6 @@ 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>(&mut self, arg: S) -> &mut Self { let arg = arg.as_ref().to_owned(); @@ -306,6 +307,15 @@ impl<'a> EnclaveBuilder<'a> { 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( mut self, loader: &mut T, @@ -332,21 +342,26 @@ impl<'a> EnclaveBuilder<'a> { } pub fn build(mut self, loader: &mut T) -> Result { + 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(mut self, loader: &mut T) -> Result { - 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)) diff --git a/intel-sgx/enclave-runner/src/usercalls/mod.rs b/intel-sgx/enclave-runner/src/usercalls/mod.rs index a1340ddc..5cfc6708 100644 --- a/intel-sgx/enclave-runner/src/usercalls/mod.rs +++ b/intel-sgx/enclave-runner/src/usercalls/mod.rs @@ -1086,7 +1086,9 @@ impl EnclaveState { usercall_ext: Option>, forward_panics: bool, cmd_args: Vec>, + 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,