use enum types to represent status and event codes

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -61,23 +61,27 @@ mod table;
 // currently haphazard.  We need a refactor to manage yielding, backpressure, and event polling and delivery in a
 // more unified and structured way.
 
-// TODO: move these into an enum:
-const STATUS_STARTING: u32 = 0;
-const STATUS_STARTED: u32 = 1;
-const STATUS_RETURNED: u32 = 2;
-const STATUS_DONE: u32 = 3;
-
-mod events {
-    // TODO: move these into an enum:
-    pub const _EVENT_CALL_STARTING: u32 = 0;
-    pub const EVENT_CALL_STARTED: u32 = 1;
-    pub const EVENT_CALL_RETURNED: u32 = 2;
-    pub const EVENT_CALL_DONE: u32 = 3;
-    pub const _EVENT_YIELDED: u32 = 4;
-    pub const EVENT_STREAM_READ: u32 = 5;
-    pub const EVENT_STREAM_WRITE: u32 = 6;
-    pub const EVENT_FUTURE_READ: u32 = 7;
-    pub const EVENT_FUTURE_WRITE: u32 = 8;
+#[derive(Clone, Copy, Eq, PartialEq)]
+#[repr(u32)]
+enum Status {
+    Starting,
+    Started,
+    Returned,
+    Done,
+}
+
+#[derive(Clone, Copy, Eq, PartialEq, Debug)]
+#[repr(u32)]
+enum Event {
+    _Starting,
+    Started,
+    Returned,
+    Done,
+    _Yielded,
+    StreamRead,
+    StreamWrite,
+    FutureRead,
+    FutureWrite,
 }
 
 const EXIT_FLAG_ASYNC_CALLER: u32 = 1 << 0;
@@ -149,7 +153,7 @@ impl<T: 'static> PromisesUnordered<T> {
 }
 
 struct HostTaskResult {
-    event: u32,
+    event: Event,
     param: u32,
     caller: TableId<GuestTask>,
 }
@@ -214,7 +218,7 @@ struct GuestTask {
     lift_result: Option<(RawLift, TypeTaskReturnIndex)>,
     result: Option<LiftedResult>,
     callback: Option<Callback>,
-    events: VecDeque<(u32, AnyTask, u32)>,
+    events: VecDeque<(Event, AnyTask, u32)>,
     caller: Caller,
     deferred: Deferred,
     should_yield: bool,
@@ -264,9 +268,7 @@ impl AnyTask {
                     .get(*task)?
                     .events
                     .iter()
-                    .filter_map(|(event, call, _)| {
-                        (*event == events::EVENT_CALL_DONE).then_some(*call)
-                    })
+                    .filter_map(|(event, call, _)| (*event == Event::Done).then_some(*call))
                     .collect::<Vec<_>>();
 
                 for call in finished {
@@ -535,7 +537,7 @@ pub(crate) fn first_poll<T, R: Send + 'static>(
                 let result = fun(store.as_context_mut())?;
                 lower(store, result)?;
                 Ok(HostTaskResult {
-                    event: events::EVENT_CALL_DONE,
+                    event: Event::Done,
                     param: 0u32,
                     caller,
                 })
@@ -606,7 +608,7 @@ pub(crate) fn poll_and_block<'a, T, R: Send + Sync + 'static>(
                 store.concurrent_state().table.get_mut(caller)?.result =
                     Some(Box::new(result) as _);
                 Ok(HostTaskResult {
-                    event: events::EVENT_CALL_DONE,
+                    event: Event::Done,
                     param: 0u32,
                     caller,
                 })
@@ -698,7 +700,7 @@ pub(crate) async fn on_fiber<'a, R: Send + Sync + 'static, T: Send>(
 fn maybe_send_event<'a, T>(
     mut store: StoreContextMut<'a, T>,
     guest_task: TableId<GuestTask>,
-    event: u32,
+    event: Event,
     call: AnyTask,
     result: u32,
 ) -> Result<StoreContextMut<'a, T>> {
@@ -718,15 +720,15 @@ fn maybe_send_event<'a, T>(
             bail!("handle not found for waitable rep {}", call.rep());
         };
         log::trace!(
-            "use callback to deliver event {event} to {} for {} (handle {handle}): {:?} {}",
+            "use callback to deliver event {event:?} to {} for {} (handle {handle}): {:?} {}",
             guest_task.rep(),
             call.rep(),
             callback.function,
             callback.context
         );
         let params = &mut [
             ValRaw::u32(callback.context),
-            ValRaw::u32(event),
+            ValRaw::u32(event as u32),
             ValRaw::u32(handle),
             ValRaw::u32(result),
         ];
@@ -741,13 +743,8 @@ fn maybe_send_event<'a, T>(
             match &store.concurrent_state().table.get(guest_task)?.caller {
                 Caller::Guest { task, .. } => {
                     let task = *task;
-                    store = maybe_send_event(
-                        store,
-                        task,
-                        events::EVENT_CALL_DONE,
-                        AnyTask::Guest(guest_task),
-                        0,
-                    )?;
+                    store =
+                        maybe_send_event(store, task, Event::Done, AnyTask::Guest(guest_task), 0)?;
                 }
                 Caller::Host(_) => {
                     log::trace!("maybe_send_event will delete {}", call.rep());
@@ -765,7 +762,7 @@ fn maybe_send_event<'a, T>(
             .events
             .push_back((event, call, result));
 
-        let resumed = if event == events::EVENT_CALL_DONE {
+        let resumed = if event == Event::Done {
             if let Some(fiber) = store
                 .concurrent_state()
                 .table
@@ -774,7 +771,7 @@ fn maybe_send_event<'a, T>(
                 .take_fiber()
             {
                 log::trace!(
-                    "use fiber to deliver event {event} to {} for {}",
+                    "use fiber to deliver event {event:?} to {} for {}",
                     guest_task.rep(),
                     call.rep()
                 );
@@ -791,7 +788,7 @@ fn maybe_send_event<'a, T>(
 
         if !resumed {
             log::trace!(
-                "queue event {event} to {} for {}",
+                "queue event {event:?} to {} for {}",
                 guest_task.rep(),
                 call.rep()
             );
@@ -819,7 +816,7 @@ fn resume_stackful<'a, T>(
                         .with_context(|| format!("bad handle: {}", guest_task.rep()))?
                         .events,
                 ) {
-                    if event == events::EVENT_CALL_DONE {
+                    if event == Event::Done {
                         log::trace!("resume_stackful will delete call {}", call.rep());
                         call.delete_all_from(store.as_context_mut())?;
                     }
@@ -832,13 +829,7 @@ fn resume_stackful<'a, T>(
                     }
                     Caller::Guest { task, .. } => {
                         let task = *task;
-                        maybe_send_event(
-                            store,
-                            task,
-                            events::EVENT_CALL_DONE,
-                            AnyTask::Guest(guest_task),
-                            0,
-                        )
+                        maybe_send_event(store, task, Event::Done, AnyTask::Guest(guest_task), 0)
                     }
                 }
             } else {
@@ -867,11 +858,11 @@ fn resume_stackless<'a, T>(
 
     let task = store.concurrent_state().table.get_mut(guest_task)?;
     let event = if task.lift_result.is_some() {
-        events::EVENT_CALL_STARTED
+        Event::Started
     } else if guest_context != 0 {
-        events::EVENT_CALL_RETURNED
+        Event::Returned
     } else {
-        events::EVENT_CALL_DONE
+        Event::Done
     };
     if guest_context != 0 {
         log::trace!("set callback for {}", guest_task.rep());
@@ -915,11 +906,10 @@ fn handle_ready<'a, T>(
         let result = fun(vm_store)?;
         store = unsafe { StoreContextMut::<T>(&mut *vm_store.cast()) };
         let task = match result.event {
-            events::EVENT_CALL_DONE => AnyTask::Host(TableId::<HostTask>::new(task)),
-            events::EVENT_STREAM_READ
-            | events::EVENT_FUTURE_READ
-            | events::EVENT_STREAM_WRITE
-            | events::EVENT_FUTURE_WRITE => AnyTask::Transmit(TableId::<TransmitState>::new(task)),
+            Event::Done => AnyTask::Host(TableId::<HostTask>::new(task)),
+            Event::StreamRead | Event::FutureRead | Event::StreamWrite | Event::FutureWrite => {
+                AnyTask::Transmit(TableId::<TransmitState>::new(task))
+            }
             _ => unreachable!(),
         };
         store = maybe_send_event(store, result.caller, result.event, task, result.param)?;
@@ -1306,7 +1296,7 @@ unsafe fn task_check<T>(cx: *mut VMOpaqueContext, async_: bool, check: TaskCheck
                 .ok_or_else(|| anyhow!("no tasks to wait for"))?;
 
             log::trace!(
-                "deliver event {event} via task.wait to {} for {}",
+                "deliver event {event:?} via task.wait to {} for {}",
                 guest_task.rep(),
                 call.rep()
             );
@@ -1332,7 +1322,7 @@ unsafe fn task_check<T>(cx: *mut VMOpaqueContext, async_: bool, check: TaskCheck
             handle.store(&mut lower, InterfaceType::U32, ptr)?;
             result.store(&mut lower, InterfaceType::U32, ptr + 4)?;
 
-            Ok(event)
+            Ok(event as u32)
         }
         TaskCheck::Poll(memory, payload, caller_instance) => {
             if let Some((event, call, result)) = cx
@@ -1343,7 +1333,7 @@ unsafe fn task_check<T>(cx: *mut VMOpaqueContext, async_: bool, check: TaskCheck
                 .pop_front()
             {
                 log::trace!(
-                    "deliver event {event} via task.poll to {} for {}",
+                    "deliver event {event:?} via task.poll to {} for {}",
                     guest_task.rep(),
                     call.rep()
                 );
@@ -1368,7 +1358,7 @@ unsafe fn task_check<T>(cx: *mut VMOpaqueContext, async_: bool, check: TaskCheck
                     &ValRaw::u32(payload),
                 )?;
                 let mut lower = LowerContext::new(cx, &options, types, instance);
-                event.store(&mut lower, InterfaceType::U32, ptr)?;
+                (event as u32).store(&mut lower, InterfaceType::U32, ptr)?;
                 handle.store(&mut lower, InterfaceType::U32, ptr + 4)?;
                 result.store(&mut lower, InterfaceType::U32, ptr + 8)?;
 
@@ -1601,7 +1591,7 @@ pub(crate) extern "C" fn async_enter<T>(
                         maybe_send_event(
                             cx,
                             TableId::new(rep),
-                            events::EVENT_CALL_STARTED,
+                            Event::Started,
                             AnyTask::Guest(task),
                             0,
                         )?;
@@ -1623,7 +1613,7 @@ pub(crate) extern "C" fn async_enter<T>(
                             maybe_send_event(
                                 cx,
                                 TableId::new(rep),
-                                events::EVENT_CALL_RETURNED,
+                                Event::Returned,
                                 AnyTask::Guest(task),
                                 0,
                             )?;
@@ -1902,29 +1892,29 @@ pub(crate) extern "C" fn async_exit<T>(
             let task = cx.concurrent_state().table.get(guest_task)?;
 
             let mut status = if task.lower_params.is_some() {
-                STATUS_STARTING
+                Status::Starting
             } else if task.lift_result.is_some() {
-                STATUS_STARTED
+                Status::Started
             } else if guest_context != 0 || callback.is_null() {
-                STATUS_RETURNED
+                Status::Returned
             } else {
-                STATUS_DONE
+                Status::Done
             };
 
-            let call = if status != STATUS_DONE {
+            let call = if status != Status::Done {
                 if (flags & EXIT_FLAG_ASYNC_CALLER) != 0 {
                     (*instance).component_waitable_tables()[caller_instance]
                         .insert(guest_task.rep(), WaitableState::Task)?
                 } else {
                     poll_for_result(cx)?;
-                    status = STATUS_DONE;
+                    status = Status::Done;
                     0
                 }
             } else {
                 0
             };
 
-            Ok((status << 30) | call)
+            Ok(((status as u32) << 30) | call)
         })
     }
 }

crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs

@@ -1,6 +1,6 @@
 use {
     super::{
-        call_host_and_handle_result, events, table::TableId, GuestTask, HostTaskFuture,
+        call_host_and_handle_result, table::TableId, Event, GuestTask, HostTaskFuture,
         HostTaskResult, Promise,
     },
     crate::{
@@ -66,7 +66,7 @@ fn state_table(instance: &mut ComponentInstance, ty: TableIndex) -> &mut StateTa
 fn push_event<T>(
     mut store: StoreContextMut<T>,
     rep: u32,
-    event: u32,
+    event: Event,
     param: usize,
     caller: TableId<GuestTask>,
 ) {
@@ -281,8 +281,8 @@ fn host_write<T: func::Lower + Send + Sync + 'static, U, S: AsContextMut<Data =
                     store.as_context_mut(),
                     transmit_id.rep(),
                     match ty {
-                        TableIndex::Future(_) => events::EVENT_FUTURE_READ,
-                        TableIndex::Stream(_) => events::EVENT_STREAM_READ,
+                        TableIndex::Future(_) => Event::FutureRead,
+                        TableIndex::Stream(_) => Event::StreamRead,
                     },
                     count,
                     caller,
@@ -417,8 +417,8 @@ pub fn host_read<T: func::Lift + Sync + Send + 'static, U, S: AsContextMut<Data
                 store,
                 transmit_id.rep(),
                 match ty {
-                    TableIndex::Future(_) => events::EVENT_FUTURE_WRITE,
-                    TableIndex::Stream(_) => events::EVENT_STREAM_WRITE,
+                    TableIndex::Future(_) => Event::FutureWrite,
+                    TableIndex::Stream(_) => Event::StreamWrite,
                 },
                 count,
                 caller,
@@ -546,8 +546,8 @@ fn host_close_writer<U, S: AsContextMut<Data = U>>(mut store: S, rep: u32) -> Re
                 store,
                 transmit_id.rep(),
                 match ty {
-                    TableIndex::Future(_) => events::EVENT_FUTURE_READ,
-                    TableIndex::Stream(_) => events::EVENT_STREAM_READ,
+                    TableIndex::Future(_) => Event::FutureRead,
+                    TableIndex::Stream(_) => Event::StreamRead,
                 },
                 CLOSED,
                 caller,
@@ -598,8 +598,8 @@ fn host_close_reader<U, S: AsContextMut<Data = U>>(mut store: S, rep: u32) -> Re
                 store.as_context_mut(),
                 transmit_id.rep(),
                 match ty {
-                    TableIndex::Future(_) => events::EVENT_FUTURE_WRITE,
-                    TableIndex::Stream(_) => events::EVENT_STREAM_WRITE,
+                    TableIndex::Future(_) => Event::FutureRead,
+                    TableIndex::Stream(_) => Event::StreamRead,
                 },
                 CLOSED,
                 caller,
@@ -1421,8 +1421,8 @@ fn guest_write<T>(
                         cx,
                         transmit_id.rep(),
                         match read_ty {
-                            TableIndex::Future(_) => events::EVENT_FUTURE_READ,
-                            TableIndex::Stream(_) => events::EVENT_STREAM_READ,
+                            TableIndex::Future(_) => Event::FutureRead,
+                            TableIndex::Stream(_) => Event::StreamRead,
                         },
                         count,
                         read_caller,
@@ -1576,8 +1576,8 @@ fn guest_read<T>(
                         cx,
                         transmit_id.rep(),
                         match write_ty {
-                            TableIndex::Future(_) => events::EVENT_FUTURE_WRITE,
-                            TableIndex::Stream(_) => events::EVENT_STREAM_WRITE,
+                            TableIndex::Future(_) => Event::FutureWrite,
+                            TableIndex::Stream(_) => Event::StreamWrite,
                         },
                         count,
                         write_caller,