mpk: optimize layout of protected stripes (#7603)

* mpk: optimize layout of protected stripes

While experimenting with the limits of MPK-protected memory pools,
@alexcrichton and I discovered that the current slab layout calculations
were too conservative. This meant that the memory pool could not pack in
as many memories as it should have been able: we were expecting, but not
seeing, ~15x more memory slots over non-MPK memory pools.

The fix ends up being simpler than the original: we must maintain the
codegen constraints that expect a static memory to be inaccessible for
OOB access within a `static_memory_maximum_size +
static_memory_guard_size` region (called `expected_slot_bytes +
guard_bytes` in `memory_pool.rs`). By dividing up that region between
the stripes, we still guarantee that the region is inaccessible by
packing in other MPK-protected stripes. And we still need to make sure
that the `post_slab_guard_bytes` add up to that region. These changes
fix the memory inefficiency issues we were seeing.

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

* mpk: eliminate extra stripe

@alexcrichton pointed out that we know that `slot_bytes /
max_memory_bytes` will at least be 1 due to a `max` comparison above.
Knowing this, we can remove a `+ 1` intended for the case when
`needed_num_stripes == 0`, which should be impossible.

* review: replace `checked_*` with `saturating_*`

This style change is a readability improvement; no calculations should
change.

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

---------

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

crates/runtime/proptest-regressions/instance/allocator/pooling/memory_pool.txt

@@ -6,3 +6,4 @@
 # everyone who runs the test benefits from these saved cases.
 cc 696808084287d5d58b85c60c4720227ab4dd83ada7be6841a67162023aaf4914 # shrinks to c = SlabConstraints { max_memory_bytes: 0, num_memory_slots: 1, num_pkeys_available: 0, guard_bytes: 9223372036854775808 }
 cc cf9f6c36659f7f56ed8ea646e8c699cbf46708cef6911cdd376418ad69ea1388 # shrinks to c = SlabConstraints { max_memory_bytes: 14161452635954640438, num_memory_slots: 0, num_pkeys_available: 0, guard_bytes: 4285291437754911178 }
+cc 58f42405c4fbfb4b464950f372995d4d08a77d6884335e38c2e68d590cacf7d8 # shrinks to c = SlabConstraints { expected_slot_bytes: 0, max_memory_bytes: 8483846582232735745, num_slots: 0, num_pkeys_available: 0, guard_bytes: 0, guard_before_slots: false }

crates/runtime/src/instance/allocator/pooling/memory_pool.rs

@@ -673,8 +673,7 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
     // to define a slot as "all of the memory and guard region."
     let slot_bytes = expected_slot_bytes
         .max(max_memory_bytes)
-        .checked_add(guard_bytes)
-        .unwrap_or(usize::MAX);
+        .saturating_add(guard_bytes);
 
     let (num_stripes, slot_bytes) = if guard_bytes == 0 || max_memory_bytes == 0 || num_slots == 0 {
         // In the uncommon case where the memory/guard regions are empty or we don't need any slots , we
@@ -703,20 +702,22 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
         // 3`), we will run into failures if we attempt to set up more than
         // three stripes.
         let needed_num_stripes =
-            slot_bytes / max_memory_bytes + usize::from(slot_bytes % max_memory_bytes != 0) + 1;
+            slot_bytes / max_memory_bytes + usize::from(slot_bytes % max_memory_bytes != 0);
+        assert!(needed_num_stripes > 0);
         let num_stripes = num_pkeys_available.min(needed_num_stripes).min(num_slots);
 
-        // Next, we try to reduce the slot size by "overlapping" the
-        // stripes: we can make slot `n` smaller since we know that slot
-        // `n+1` and following are in different stripes and will look just
-        // like `PROT_NONE` memory.
-        let next_slots_overlapping_bytes = max_memory_bytes
-            .checked_mul(num_stripes - 1)
-            .unwrap_or(usize::MAX);
+        // Next, we try to reduce the slot size by "overlapping" the stripes: we
+        // must respect the `expected_slot_bytes + guard bytes` expectations
+        // that codegen constrains us to, but, since a stripe of a different
+        // color is just as inaccessible as `PROT_NONE`, we can squeeze them
+        // together, effectively reducing the size of `slot_bytes` but still
+        // guaranteeing the codegen constraints.
         let needed_slot_bytes = slot_bytes
-            .checked_sub(next_slots_overlapping_bytes)
-            .unwrap_or(0)
+            .checked_div(num_stripes)
+            .unwrap_or(slot_bytes)
             .max(max_memory_bytes);
+        assert!(needed_slot_bytes >= max_memory_bytes);
+
         (num_stripes, needed_slot_bytes)
     };
 
@@ -728,13 +729,14 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
         .ok_or_else(|| anyhow!("slot size is too large"))?;
 
     // We may need another guard region (like `pre_slab_guard_bytes`) at the end
-    // of our slab. We could be conservative and just create it as large as
-    // `guard_bytes`, but because we know that the last slot already has
-    // `guard_bytes` factored in to its guard region, we can reduce the final
-    // guard region by that much.
-    let post_slab_guard_bytes = guard_bytes
-        .checked_sub(slot_bytes - max_memory_bytes)
-        .unwrap_or(0);
+    // of our slab. If we have overlapped stripes of different colors to reduce
+    // the slot size, the last stripe will be vulnerable--it is not followed by
+    // other inaccessible stripes. To fix this, we ensure that the
+    // `post_slab_guard_bytes` is large enough to ensure the last slot meets the
+    // `expected_slot_bytes + guard_bytes` guarantees.
+    let post_slab_guard_bytes = expected_slot_bytes
+        .saturating_add(guard_bytes)
+        .saturating_sub(slot_bytes);
 
     // Check that we haven't exceeded the slab we can calculate given the limits
     // of `usize`.