nitro-cli: Update enclave boot timeout based on allocated memory

Before pull request #310 the enlave boot timeout was calculated based on
the allocated memory size. We reintroduce this heuristic in addition to
the heuristic based on the eif size, since we observed timeout
violations for large allocated memory sizes.

Signed-off-by: Marius Knaust <[email protected]>

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<Duration> {
 /// # 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 eif_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(eif_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
+        );
     }
 }