Skip to content

Implement async task cancellation #1264

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alexcrichton merged 1 commit into from
Apr 30, 2025
Merged

Implement async task cancellation #1264

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 95 additions & 31 deletions crates/core/src/abi.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -554,6 +554,10 @@ def_instruction! {
blocks: usize,
} : [1] => [0],

/// Deallocates the language-specific handle representation on the top
/// of the stack. Used for async imports.
DropHandle { ty: &'a Type } : [1] => [0],

/// Call `task.return` for an async-lifted export.
///
/// This will call core wasm import `name` which will be mapped to
Expand Down Expand Up @@ -783,33 +787,37 @@ pub fn post_return(resolve: &Resolve, func: &Function, bindgen: &mut impl Bindge
/// a list or a string primarily.
pub fn guest_export_needs_post_return(resolve: &Resolve, func: &Function) -> bool {
func.result
.map(|t| needs_post_return(resolve, &t))
.map(|t| needs_deallocate(resolve, &t, Deallocate::Lists))
.unwrap_or(false)
}

fn needs_post_return(resolve: &Resolve, ty: &Type) -> bool {
fn needs_deallocate(resolve: &Resolve, ty: &Type, what: Deallocate) -> bool {
match ty {
Type::String => true,
Type::ErrorContext => true,
Type::Id(id) => match &resolve.types[*id].kind {
TypeDefKind::List(_) => true,
TypeDefKind::Type(t) => needs_post_return(resolve, t),
TypeDefKind::Handle(_) => false,
TypeDefKind::Type(t) => needs_deallocate(resolve, t, what),
TypeDefKind::Handle(Handle::Own(_)) => what.handles(),
TypeDefKind::Handle(Handle::Borrow(_)) => false,
TypeDefKind::Resource => false,
TypeDefKind::Record(r) => r.fields.iter().any(|f| needs_post_return(resolve, &f.ty)),
TypeDefKind::Tuple(t) => t.types.iter().any(|t| needs_post_return(resolve, t)),
TypeDefKind::Record(r) => r
.fields
.iter()
.any(|f| needs_deallocate(resolve, &f.ty, what)),
TypeDefKind::Tuple(t) => t.types.iter().any(|t| needs_deallocate(resolve, t, what)),
TypeDefKind::Variant(t) => t
.cases
.iter()
.filter_map(|t| t.ty.as_ref())
.any(|t| needs_post_return(resolve, t)),
TypeDefKind::Option(t) => needs_post_return(resolve, t),
.any(|t| needs_deallocate(resolve, t, what)),
TypeDefKind::Option(t) => needs_deallocate(resolve, t, what),
TypeDefKind::Result(t) => [&t.ok, &t.err]
.iter()
.filter_map(|t| t.as_ref())
.any(|t| needs_post_return(resolve, t)),
.any(|t| needs_deallocate(resolve, t, what)),
TypeDefKind::Flags(_) | TypeDefKind::Enum(_) => false,
TypeDefKind::Future(_) | TypeDefKind::Stream(_) => false,
TypeDefKind::Future(_) | TypeDefKind::Stream(_) => what.handles(),
TypeDefKind::Unknown => unreachable!(),
},

Expand All @@ -836,7 +844,18 @@ pub fn deallocate_lists_in_types<B: Bindgen>(
ptr: B::Operand,
bindgen: &mut B,
) {
Generator::new(resolve, bindgen).deallocate_lists_in_types(types, ptr);
Generator::new(resolve, bindgen).deallocate_in_types(types, ptr, Deallocate::Lists);
}

/// Generate instructions in `bindgen` to deallocate all lists in `ptr` where
/// that's a pointer to a sequence of `types` stored in linear memory.
pub fn deallocate_lists_and_own_in_types<B: Bindgen>(
resolve: &Resolve,
types: &[Type],
ptr: B::Operand,
bindgen: &mut B,
) {
Generator::new(resolve, bindgen).deallocate_in_types(types, ptr, Deallocate::ListsAndOwn);
}

#[derive(Copy, Clone)]
Expand All @@ -845,6 +864,25 @@ pub enum Realloc {
Export(&'static str),
}

/// What to deallocate in various `deallocate_*` methods.
#[derive(Copy, Clone)]
enum Deallocate {
/// Only deallocate lists.
Lists,
/// Deallocate lists and owned resources such as `own<T>` and
/// futures/streams.
ListsAndOwn,
}

impl Deallocate {
fn handles(&self) -> bool {
match self {
Deallocate::Lists => false,
Deallocate::ListsAndOwn => true,
}
}
}

struct Generator<'a, B: Bindgen> {
bindgen: &'a mut B,
resolve: &'a Resolve,
Expand Down Expand Up @@ -1166,14 +1204,14 @@ impl<'a, B: Bindgen> Generator<'a, B> {

let mut types = Vec::new();
types.extend(func.result);
self.deallocate_lists_in_types(&types, addr);
self.deallocate_in_types(&types, addr, Deallocate::Lists);

self.emit(&Instruction::Return { func, amt: 0 });
}

fn deallocate_lists_in_types(&mut self, types: &[Type], addr: B::Operand) {
fn deallocate_in_types(&mut self, types: &[Type], addr: B::Operand, what: Deallocate) {
for (offset, ty) in self.bindgen.sizes().field_offsets(types) {
self.deallocate(ty, addr.clone(), offset);
self.deallocate(ty, addr.clone(), offset, what);
}

assert!(
Expand Down Expand Up @@ -1973,12 +2011,18 @@ impl<'a, B: Bindgen> Generator<'a, B> {
});
}

fn deallocate(&mut self, ty: &Type, addr: B::Operand, offset: ArchitectureSize) {
fn deallocate(
&mut self,
ty: &Type,
addr: B::Operand,
offset: ArchitectureSize,
what: Deallocate,
) {
use Instruction::*;

// No need to execute any instructions if this type itself doesn't
// require any form of post-return.
if !needs_post_return(self.resolve, ty) {
if !needs_deallocate(self.resolve, ty, what) {
return;
}

Expand Down Expand Up @@ -2008,7 +2052,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
| Type::ErrorContext => {}

Type::Id(id) => match &self.resolve.types[id].kind {
TypeDefKind::Type(t) => self.deallocate(t, addr, offset),
TypeDefKind::Type(t) => self.deallocate(t, addr, offset, what),

TypeDefKind::List(element) => {
self.stack.push(addr.clone());
Expand All @@ -2021,30 +2065,36 @@ impl<'a, B: Bindgen> Generator<'a, B> {
self.push_block();
self.emit(&IterBasePointer);
let elemaddr = self.stack.pop().unwrap();
self.deallocate(element, elemaddr, Default::default());
self.deallocate(element, elemaddr, Default::default(), what);
self.finish_block(0);

self.emit(&Instruction::GuestDeallocateList { element });
}

TypeDefKind::Handle(_) => {
todo!()
TypeDefKind::Handle(Handle::Own(_))
| TypeDefKind::Future(_)
| TypeDefKind::Stream(_)
if what.handles() =>
{
self.read_from_memory(ty, addr, offset);
self.emit(&DropHandle { ty });
}

TypeDefKind::Resource => {
todo!()
}
TypeDefKind::Handle(Handle::Own(_)) => unreachable!(),
TypeDefKind::Handle(Handle::Borrow(_)) => unreachable!(),
TypeDefKind::Resource => unreachable!(),

TypeDefKind::Record(record) => {
self.deallocate_fields(
&record.fields.iter().map(|f| f.ty).collect::<Vec<_>>(),
addr,
offset,
what,
);
}

TypeDefKind::Tuple(tuple) => {
self.deallocate_fields(&tuple.types, addr, offset);
self.deallocate_fields(&tuple.types, addr, offset, what);
}

TypeDefKind::Flags(_) => {}
Expand All @@ -2055,26 +2105,33 @@ impl<'a, B: Bindgen> Generator<'a, B> {
addr,
variant.tag(),
variant.cases.iter().map(|c| c.ty.as_ref()),
what,
);
self.emit(&GuestDeallocateVariant {
blocks: variant.cases.len(),
});
}

TypeDefKind::Option(t) => {
self.deallocate_variant(offset, addr, Int::U8, [None, Some(t)]);
self.deallocate_variant(offset, addr, Int::U8, [None, Some(t)], what);
self.emit(&GuestDeallocateVariant { blocks: 2 });
}

TypeDefKind::Result(e) => {
self.deallocate_variant(offset, addr, Int::U8, [e.ok.as_ref(), e.err.as_ref()]);
self.deallocate_variant(
offset,
addr,
Int::U8,
[e.ok.as_ref(), e.err.as_ref()],
what,
);
self.emit(&GuestDeallocateVariant { blocks: 2 });
}

TypeDefKind::Enum(_) => {}

TypeDefKind::Future(_) => todo!("read future from memory"),
TypeDefKind::Stream(_) => todo!("read stream from memory"),
TypeDefKind::Future(_) => unreachable!(),
TypeDefKind::Stream(_) => unreachable!(),
TypeDefKind::Unknown => unreachable!(),
},
}
Expand All @@ -2086,22 +2143,29 @@ impl<'a, B: Bindgen> Generator<'a, B> {
addr: B::Operand,
tag: Int,
cases: impl IntoIterator<Item = Option<&'b Type>> + Clone,
what: Deallocate,
) {
self.stack.push(addr.clone());
self.load_intrepr(offset, tag);
let payload_offset = offset + (self.bindgen.sizes().payload_offset(tag, cases.clone()));
for ty in cases {
self.push_block();
if let Some(ty) = ty {
self.deallocate(ty, addr.clone(), payload_offset);
self.deallocate(ty, addr.clone(), payload_offset, what);
}
self.finish_block(0);
}
}

fn deallocate_fields(&mut self, tys: &[Type], addr: B::Operand, offset: ArchitectureSize) {
fn deallocate_fields(
&mut self,
tys: &[Type],
addr: B::Operand,
offset: ArchitectureSize,
what: Deallocate,
) {
for (field_offset, ty) in self.bindgen.sizes().field_offsets(tys) {
self.deallocate(ty, addr.clone(), offset + (field_offset));
self.deallocate(ty, addr.clone(), offset + (field_offset), what);
}
}
}
Expand Down
4 changes: 3 additions & 1 deletion crates/csharp/src/function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1262,7 +1262,9 @@ impl Bindgen for FunctionBindgen<'_, '_> {
| Instruction::StreamLower { .. }
| Instruction::StreamLift { .. }
| Instruction::ErrorContextLower { .. }
| Instruction::ErrorContextLift { .. } => todo!(),
| Instruction::ErrorContextLift { .. }
| Instruction::DropHandle { .. }
=> todo!(),
}
}

Expand Down
Loading
Loading