diff --git a/src/enclave_proc/resource_manager.rs b/src/enclave_proc/resource_manager.rs index d67d6eea7..4153d0a42 100644 --- a/src/enclave_proc/resource_manager.rs +++ b/src/enclave_proc/resource_manager.rs @@ -589,7 +589,7 @@ impl EnclaveHandle { .len(); // Update the poll timeout based on the eif size or allocated memory - let poll_timeout = calculate_necessary_timeout(eif_size); + let poll_timeout = calculate_necessary_timeout(eif_size, self.allocated_memory_mib * MiB); enclave_ready(listener, poll_timeout).map_err(|err| { let err_msg = format!("Waiting on enclave to boot failed with error {:?}", err); @@ -1188,21 +1188,26 @@ pub fn between_packets_delay() -> Option { /// # Arguments /// /// * `eif_size` - The EIF size in bytes +/// * `allocated_memory` - The memory size in bytes /// /// # Examples /// /// ``` /// use nitro_cli::enclave_proc::resource_manager::calculate_necessary_timeout; /// use nitro_cli::enclave_proc::utils::GiB; -/// // Returns the timeout based on the 8GiB EIF size -/// let timeout = calculate_necessary_timeout(8 * GiB); +/// // Returns the timeout based on the 8GiB EIF size and 32GiB of allocated memory +/// let timeout = calculate_necessary_timeout(8 * GiB, 32 * GiB); /// ``` -pub fn calculate_necessary_timeout(eif_size: u64) -> c_int { +pub fn calculate_necessary_timeout(eif_size: u64, allocated_memory: u64) -> c_int { // in case we have a valid eif_size give TIMEOUT_MINUTE_MS ms for each 6GiB - let poll_timeout: c_int = + let efi_size_timeout: c_int = ((1 + (eif_size - 1) / (6 * GiB)) as i32).saturating_mul(TIMEOUT_MINUTE_MS); - poll_timeout + // Update the poll timeout to be TIMEOUT_MINUTE_MS per 100 GiB of enclave memory. + let allocated_memory_timeout: c_int = + ((1 + (allocated_memory - 1) / (100 * GiB)) as i32).saturating_mul(TIMEOUT_MINUTE_MS); + + std::cmp::max(efi_size_timeout, allocated_memory_timeout) } #[cfg(test)] @@ -1213,8 +1218,30 @@ mod tests { #[test] fn test_timeout_calculation() { - assert_eq!(calculate_necessary_timeout(2 * GiB), TIMEOUT_MINUTE_MS); - assert_eq!(calculate_necessary_timeout(6 * GiB), TIMEOUT_MINUTE_MS); - assert_eq!(calculate_necessary_timeout(10 * GiB), 2 * TIMEOUT_MINUTE_MS); + assert_eq!( + calculate_necessary_timeout(2 * GiB, 32 * GiB), + TIMEOUT_MINUTE_MS + ); + assert_eq!( + calculate_necessary_timeout(6 * GiB, 32 * GiB), + TIMEOUT_MINUTE_MS + ); + assert_eq!( + calculate_necessary_timeout(10 * GiB, 32 * GiB), + 2 * TIMEOUT_MINUTE_MS + ); + + assert_eq!( + calculate_necessary_timeout(2 * GiB, 128 * GiB), + 2 * TIMEOUT_MINUTE_MS + ); + assert_eq!( + calculate_necessary_timeout(6 * GiB, 128 * GiB), + 2 * TIMEOUT_MINUTE_MS + ); + assert_eq!( + calculate_necessary_timeout(10 * GiB, 128 * GiB), + 2 * TIMEOUT_MINUTE_MS + ); } }