Skip to content

Commit a5f6569

Browse files
alexcrichtonalexcrichton
committed
Use RAII guards in a few locations instead of closures
This represents no functional difference from what happens today but in the future all of these functions will become `async` functions where RAII guards will then be required, so this goes ahead and switches them to RAII guards to split out the diff.
1 parent aa91737 commit a5f6569

3 files changed

Lines changed: 180 additions & 127 deletions

File tree

‎crates/wasmtime/src/runtime/store/gc.rs‎

Lines changed: 76 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -81,72 +81,93 @@ impl StoreOpaque {
8181
log::trace!("Attempting to grow the GC heap by {bytes_needed} bytes");
8282
assert!(bytes_needed > 0);
8383

84+
let page_size = self.engine().tunables().gc_heap_memory_type().page_size();
85+
8486
// Take the GC heap's underlying memory out of the GC heap, attempt to
8587
// grow it, then replace it.
86-
let mut memory = self.unwrap_gc_store_mut().gc_heap.take_memory();
87-
let mut delta_bytes_grown = 0;
88-
let grow_result: Result<()> = (|| {
89-
let page_size = self.engine().tunables().gc_heap_memory_type().page_size();
88+
let mut heap = TakenGcHeap::new(self);
9089

91-
let current_size_in_bytes = u64::try_from(memory.byte_size()).unwrap();
92-
let current_size_in_pages = current_size_in_bytes / page_size;
90+
let current_size_in_bytes = u64::try_from(heap.memory.byte_size()).unwrap();
91+
let current_size_in_pages = current_size_in_bytes / page_size;
9392

94-
// Aim to double the heap size, amortizing the cost of growth.
95-
let doubled_size_in_pages = current_size_in_pages.saturating_mul(2);
96-
assert!(doubled_size_in_pages >= current_size_in_pages);
97-
let delta_pages_for_doubling = doubled_size_in_pages - current_size_in_pages;
93+
// Aim to double the heap size, amortizing the cost of growth.
94+
let doubled_size_in_pages = current_size_in_pages.saturating_mul(2);
95+
assert!(doubled_size_in_pages >= current_size_in_pages);
96+
let delta_pages_for_doubling = doubled_size_in_pages - current_size_in_pages;
9897

99-
// When doubling our size, saturate at the maximum memory size in pages.
100-
//
101-
// TODO: we should consult the instance allocator for its configured
102-
// maximum memory size, if any, rather than assuming the index
103-
// type's maximum size.
104-
let max_size_in_bytes = 1 << 32;
105-
let max_size_in_pages = max_size_in_bytes / page_size;
106-
let delta_to_max_size_in_pages = max_size_in_pages - current_size_in_pages;
107-
let delta_pages_for_alloc = delta_pages_for_doubling.min(delta_to_max_size_in_pages);
98+
// When doubling our size, saturate at the maximum memory size in pages.
99+
//
100+
// TODO: we should consult the instance allocator for its configured
101+
// maximum memory size, if any, rather than assuming the index
102+
// type's maximum size.
103+
let max_size_in_bytes = 1 << 32;
104+
let max_size_in_pages = max_size_in_bytes / page_size;
105+
let delta_to_max_size_in_pages = max_size_in_pages - current_size_in_pages;
106+
let delta_pages_for_alloc = delta_pages_for_doubling.min(delta_to_max_size_in_pages);
108107

109-
// But always make sure we are attempting to grow at least as many pages
110-
// as needed by the requested allocation. This must happen *after* the
111-
// max-size saturation, so that if we are at the max already, we do not
112-
// succeed in growing by zero delta pages, and then return successfully
113-
// to our caller, who would be assuming that there is now capacity for
114-
// their allocation.
115-
let pages_needed = bytes_needed.div_ceil(page_size);
116-
assert!(pages_needed > 0);
117-
let delta_pages_for_alloc = delta_pages_for_alloc.max(pages_needed);
118-
assert!(delta_pages_for_alloc > 0);
108+
// But always make sure we are attempting to grow at least as many pages
109+
// as needed by the requested allocation. This must happen *after* the
110+
// max-size saturation, so that if we are at the max already, we do not
111+
// succeed in growing by zero delta pages, and then return successfully
112+
// to our caller, who would be assuming that there is now capacity for
113+
// their allocation.
114+
let pages_needed = bytes_needed.div_ceil(page_size);
115+
assert!(pages_needed > 0);
116+
let delta_pages_for_alloc = delta_pages_for_alloc.max(pages_needed);
117+
assert!(delta_pages_for_alloc > 0);
119118

120-
// Safety: we pair growing the GC heap with updating its associated
121-
// `VMMemoryDefinition` in the `VMStoreContext` immediately
122-
// afterwards.
123-
unsafe {
124-
memory
125-
.grow(delta_pages_for_alloc, Some(self.traitobj().as_mut()))?
126-
.ok_or_else(|| anyhow!("failed to grow GC heap"))?;
127-
}
128-
self.vm_store_context.gc_heap = memory.vmmemory();
119+
// Safety: we pair growing the GC heap with updating its associated
120+
// `VMMemoryDefinition` in the `VMStoreContext` immediately
121+
// afterwards.
122+
unsafe {
123+
heap.memory
124+
.grow(delta_pages_for_alloc, Some(heap.store.traitobj().as_mut()))?
125+
.ok_or_else(|| anyhow!("failed to grow GC heap"))?;
126+
}
127+
heap.store.vm_store_context.gc_heap = heap.memory.vmmemory();
129128

130-
let new_size_in_bytes = u64::try_from(memory.byte_size()).unwrap();
131-
assert!(new_size_in_bytes > current_size_in_bytes);
132-
delta_bytes_grown = new_size_in_bytes - current_size_in_bytes;
133-
let delta_bytes_for_alloc = delta_pages_for_alloc.checked_mul(page_size).unwrap();
134-
assert!(
135-
delta_bytes_grown >= delta_bytes_for_alloc,
136-
"{delta_bytes_grown} should be greater than or equal to {delta_bytes_for_alloc}"
137-
);
138-
Ok(())
139-
})();
129+
let new_size_in_bytes = u64::try_from(heap.memory.byte_size()).unwrap();
130+
assert!(new_size_in_bytes > current_size_in_bytes);
131+
heap.delta_bytes_grown = new_size_in_bytes - current_size_in_bytes;
132+
let delta_bytes_for_alloc = delta_pages_for_alloc.checked_mul(page_size).unwrap();
133+
assert!(
134+
heap.delta_bytes_grown >= delta_bytes_for_alloc,
135+
"{} should be greater than or equal to {delta_bytes_for_alloc}",
136+
heap.delta_bytes_grown,
137+
);
138+
return Ok(());
140139

141-
// Regardless whether growing succeeded or failed, place the memory back
142-
// inside the GC heap.
143-
unsafe {
144-
self.unwrap_gc_store_mut()
145-
.gc_heap
146-
.replace_memory(memory, delta_bytes_grown);
140+
struct TakenGcHeap<'a> {
141+
store: &'a mut StoreOpaque,
142+
memory: ManuallyDrop<vm::Memory>,
143+
delta_bytes_grown: u64,
147144
}
148145

149-
grow_result
146+
impl<'a> TakenGcHeap<'a> {
147+
fn new(store: &'a mut StoreOpaque) -> TakenGcHeap<'a> {
148+
TakenGcHeap {
149+
memory: ManuallyDrop::new(store.unwrap_gc_store_mut().gc_heap.take_memory()),
150+
store,
151+
delta_bytes_grown: 0,
152+
}
153+
}
154+
}
155+
156+
impl Drop for TakenGcHeap<'_> {
157+
fn drop(&mut self) {
158+
// SAFETY: this `Drop` guard ensures that this has exclusive
159+
// ownership of fields and is thus safe to take `self.memory`.
160+
// Additionally for `replace_memory` the memory was previously
161+
// taken when this was created so it should be safe to place
162+
// back inside the GC heap.
163+
unsafe {
164+
self.store.unwrap_gc_store_mut().gc_heap.replace_memory(
165+
ManuallyDrop::take(&mut self.memory),
166+
self.delta_bytes_grown,
167+
);
168+
}
169+
}
170+
}
150171
}
151172

152173
/// Attempt an allocation, if it fails due to GC OOM, then do a GC and

‎crates/wasmtime/src/runtime/vm/instance/allocator/pooling/memory_pool.rs‎

Lines changed: 68 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -344,62 +344,79 @@ impl MemoryPool {
344344
format!("memory stripe {stripe_index}"),
345345
)
346346
})?;
347+
let mut guard = DeallocateIndexGuard {
348+
pool: self,
349+
stripe_index,
350+
striped_allocation_index,
351+
active: true,
352+
};
353+
347354
let allocation_index =
348355
striped_allocation_index.as_unstriped_slot_index(stripe_index, self.stripes.len());
349356

350-
match (|| {
351-
// Double-check that the runtime requirements of the memory are
352-
// satisfied by the configuration of this pooling allocator. This
353-
// should be returned as an error through `validate_memory_plans`
354-
// but double-check here to be sure.
355-
assert!(
356-
tunables.memory_reservation + tunables.memory_guard_size
357-
<= u64::try_from(self.layout.bytes_to_next_stripe_slot().byte_count()).unwrap()
358-
);
357+
// Double-check that the runtime requirements of the memory are
358+
// satisfied by the configuration of this pooling allocator. This
359+
// should be returned as an error through `validate_memory_plans`
360+
// but double-check here to be sure.
361+
assert!(
362+
tunables.memory_reservation + tunables.memory_guard_size
363+
<= u64::try_from(self.layout.bytes_to_next_stripe_slot().byte_count()).unwrap()
364+
);
359365

360-
let base = self.get_base(allocation_index);
361-
let base_capacity = self.layout.max_memory_bytes;
362-
363-
let mut slot = self.take_memory_image_slot(allocation_index);
364-
let image = match memory_index {
365-
Some(memory_index) => request.runtime_info.memory_image(memory_index)?,
366-
None => None,
367-
};
368-
let initial_size = ty
369-
.minimum_byte_size()
370-
.expect("min size checked in validation");
371-
372-
// If instantiation fails, we can propagate the error
373-
// upward and drop the slot. This will cause the Drop
374-
// handler to attempt to map the range with PROT_NONE
375-
// memory, to reserve the space while releasing any
376-
// stale mappings. The next use of this slot will then
377-
// create a new slot that will try to map over
378-
// this, returning errors as well if the mapping
379-
// errors persist. The unmap-on-drop is best effort;
380-
// if it fails, then we can still soundly continue
381-
// using the rest of the pool and allowing the rest of
382-
// the process to continue, because we never perform a
383-
// mmap that would leave an open space for someone
384-
// else to come in and map something.
385-
let initial_size = usize::try_from(initial_size).unwrap();
386-
slot.instantiate(initial_size, image, ty, tunables)?;
387-
388-
Memory::new_static(
389-
ty,
390-
tunables,
391-
MemoryBase::Mmap(base),
392-
base_capacity.byte_count(),
393-
slot,
394-
unsafe { &mut *request.store.get().unwrap() },
395-
)
396-
})() {
397-
Ok(memory) => Ok((allocation_index, memory)),
398-
Err(e) => {
399-
self.stripes[stripe_index]
366+
let base = self.get_base(allocation_index);
367+
let base_capacity = self.layout.max_memory_bytes;
368+
369+
let mut slot = self.take_memory_image_slot(allocation_index);
370+
let image = match memory_index {
371+
Some(memory_index) => request.runtime_info.memory_image(memory_index)?,
372+
None => None,
373+
};
374+
let initial_size = ty
375+
.minimum_byte_size()
376+
.expect("min size checked in validation");
377+
378+
// If instantiation fails, we can propagate the error
379+
// upward and drop the slot. This will cause the Drop
380+
// handler to attempt to map the range with PROT_NONE
381+
// memory, to reserve the space while releasing any
382+
// stale mappings. The next use of this slot will then
383+
// create a new slot that will try to map over
384+
// this, returning errors as well if the mapping
385+
// errors persist. The unmap-on-drop is best effort;
386+
// if it fails, then we can still soundly continue
387+
// using the rest of the pool and allowing the rest of
388+
// the process to continue, because we never perform a
389+
// mmap that would leave an open space for someone
390+
// else to come in and map something.
391+
let initial_size = usize::try_from(initial_size).unwrap();
392+
slot.instantiate(initial_size, image, ty, tunables)?;
393+
394+
let memory = Memory::new_static(
395+
ty,
396+
tunables,
397+
MemoryBase::Mmap(base),
398+
base_capacity.byte_count(),
399+
slot,
400+
unsafe { &mut *request.store.get().unwrap() },
401+
)?;
402+
guard.active = false;
403+
return Ok((allocation_index, memory));
404+
405+
struct DeallocateIndexGuard<'a> {
406+
pool: &'a MemoryPool,
407+
stripe_index: usize,
408+
striped_allocation_index: StripedAllocationIndex,
409+
active: bool,
410+
}
411+
412+
impl Drop for DeallocateIndexGuard<'_> {
413+
fn drop(&mut self) {
414+
if !self.active {
415+
return;
416+
}
417+
self.pool.stripes[self.stripe_index]
400418
.allocator
401-
.free(SlotId(striped_allocation_index.0));
402-
Err(e)
419+
.free(SlotId(self.striped_allocation_index.0));
403420
}
404421
}
405422
}

‎crates/wasmtime/src/runtime/vm/instance/allocator/pooling/table_pool.rs‎

Lines changed: 36 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -143,29 +143,44 @@ impl TablePool {
143143
.ok_or_else(|| {
144144
super::PoolConcurrencyLimitError::new(self.max_total_tables, "tables")
145145
})?;
146+
let mut guard = DeallocateIndexGuard {
147+
pool: self,
148+
allocation_index,
149+
active: true,
150+
};
146151

147-
match (|| {
148-
let base = self.get(allocation_index);
149-
let data_size = self.data_size(crate::vm::table::wasm_to_table_type(ty.ref_type));
150-
unsafe {
151-
commit_pages(base, data_size)?;
152-
}
152+
let base = self.get(allocation_index);
153+
let data_size = self.data_size(crate::vm::table::wasm_to_table_type(ty.ref_type));
154+
unsafe {
155+
commit_pages(base, data_size)?;
156+
}
153157

154-
let ptr =
155-
NonNull::new(std::ptr::slice_from_raw_parts_mut(base.cast(), data_size)).unwrap();
156-
unsafe {
157-
Table::new_static(
158-
ty,
159-
tunables,
160-
SendSyncPtr::new(ptr),
161-
&mut *request.store.get().unwrap(),
162-
)
163-
}
164-
})() {
165-
Ok(table) => Ok((allocation_index, table)),
166-
Err(e) => {
167-
self.index_allocator.free(SlotId(allocation_index.0));
168-
Err(e)
158+
let ptr = NonNull::new(std::ptr::slice_from_raw_parts_mut(base.cast(), data_size)).unwrap();
159+
let table = unsafe {
160+
Table::new_static(
161+
ty,
162+
tunables,
163+
SendSyncPtr::new(ptr),
164+
&mut *request.store.get().unwrap(),
165+
)?
166+
};
167+
guard.active = false;
168+
return Ok((allocation_index, table));
169+
170+
struct DeallocateIndexGuard<'a> {
171+
pool: &'a TablePool,
172+
allocation_index: TableAllocationIndex,
173+
active: bool,
174+
}
175+
176+
impl Drop for DeallocateIndexGuard<'_> {
177+
fn drop(&mut self) {
178+
if !self.active {
179+
return;
180+
}
181+
self.pool
182+
.index_allocator
183+
.free(SlotId(self.allocation_index.0));
169184
}
170185
}
171186
}

0 commit comments

Comments
 (0)