Beef up documentation about TLS management in Wasmtime (#10547)

This comes out of discussions on the wasip3-prototyping repository where
I was trying to re-learn what this all does. The goal here is to more
exhaustively document how this module works.

Author: alexcrichton

crates/wasmtime/src/runtime/vm/traphandlers.rs

@@ -425,8 +425,27 @@ mod call_thread_state {
     use super::*;
     use crate::runtime::vm::Unwind;
 
-    /// Temporary state stored on the stack which is registered in the `tls` module
-    /// below for calls into wasm.
+    /// Temporary state stored on the stack which is registered in the `tls`
+    /// module below for calls into wasm.
+    ///
+    /// This structure is stored on the stack and allocated during the
+    /// `catch_traps` function above. The purpose of this structure is to track
+    /// the state of an "activation" or a sequence of 0-or-more contiguous
+    /// WebAssembly call frames. A `CallThreadState` always lives on the stack
+    /// and additionally maintains pointers to previous states to form a linked
+    /// list of activations.
+    ///
+    /// One of the primary goals of `CallThreadState` is to store the state of
+    /// various fields in `VMStoreContext` when it was created. This is done
+    /// because calling WebAssembly will clobber these fields otherwise.
+    ///
+    /// Another major purpose of `CallThreadState` is to assist with unwinding
+    /// and track state necessary when an unwind happens for the original
+    /// creator of `CallThreadState` to determine why the unwind happened.
+    ///
+    /// Note that this structure is pointed-to from TLS, hence liberal usage of
+    /// interior mutability here since that only gives access to
+    /// `&CallThreadState`.
     pub struct CallThreadState {
         pub(super) unwind: Cell<Option<(UnwindReason, Option<Backtrace>, Option<CoreDumpStack>)>>,
         pub(super) jmp_buf: Cell<*const u8>,
@@ -525,12 +544,31 @@ mod call_thread_state {
             self.prev.get()
         }
 
+        /// Pushes this `CallThreadState` activation on to the linked list
+        /// stored in TLS.
+        ///
+        /// This method will take the current head of the linked list, stored in
+        /// our TLS pointer, and move it into `prev`. The TLS pointer is then
+        /// updated to `self`.
+        ///
+        /// # Panics
+        ///
+        /// Panics if this activation is already in a linked list (e.g.
+        /// `self.prev` is set).
         #[inline]
         pub(crate) unsafe fn push(&self) {
             assert!(self.prev.get().is_null());
             self.prev.set(tls::raw::replace(self));
         }
 
+        /// Pops this `CallThreadState` from the linked list stored in TLS.
+        ///
+        /// This method will restore `self.prev` into the head of the linked
+        /// list stored in TLS and will additionally null-out `self.prev`.
+        ///
+        /// # Panics
+        ///
+        /// Panics if this activation isn't the head of the list.
         #[inline]
         pub(crate) unsafe fn pop(&self) {
             let prev = self.prev.replace(ptr::null());
@@ -734,11 +772,84 @@ impl CallThreadState {
     }
 }
 
-// A private inner module for managing the TLS state that we require across
-// calls in wasm. The WebAssembly code is called from C++ and then a trap may
-// happen which requires us to read some contextual state to figure out what to
-// do with the trap. This `tls` module is used to persist that information from
-// the caller to the trap site.
+/// A private inner module managing the state of Wasmtime's thread-local storage
+/// (TLS) state.
+///
+/// Wasmtime at this time has a single pointer of TLS. This single pointer of
+/// TLS is the totality of all TLS required by Wasmtime. By keeping this as
+/// small as possible it generally makes it easier to integrate with external
+/// systems and implement features such as fiber context switches. This single
+/// TLS pointer is declared in platform-specific modules to handle platform
+/// differences, so this module here uses getters/setters which delegate to
+/// platform-specific implementations.
+///
+/// The single TLS pointer used by Wasmtime is morally
+/// `Option<&CallThreadState>` meaning that it's a possibly-present pointer to
+/// some state. This pointer is a pointer to the most recent (youngest)
+/// `CallThreadState` activation, or the most recent call into WebAssembly.
+///
+/// This TLS pointer is additionally the head of a linked list of activations
+/// that are all stored on the stack for the current thread. Each time
+/// WebAssembly is recursively invoked by an embedder will push a new entry into
+/// this linked list. This singly-linked list is maintained with its head in TLS
+/// node pointers are stored in `CallThreadState::prev`.
+///
+/// An example stack might look like this:
+///
+/// ```text
+/// ┌─────────────────────┐◄───── highest, or oldest, stack address
+/// │ native stack frames │
+/// │         ...         │
+/// │  ┌───────────────┐◄─┼──┐
+/// │  │CallThreadState│  │  │
+/// │  └───────────────┘  │  p
+/// ├─────────────────────┤  r
+/// │  wasm stack frames  │  e
+/// │         ...         │  v
+/// ├─────────────────────┤  │
+/// │ native stack frames │  │
+/// │         ...         │  │
+/// │  ┌───────────────┐◄─┼──┼── TLS pointer
+/// │  │CallThreadState├──┼──┘
+/// │  └───────────────┘  │
+/// ├─────────────────────┤
+/// │  wasm stack frames  │
+/// │         ...         │
+/// ├─────────────────────┤
+/// │ native stack frames │
+/// │         ...         │
+/// └─────────────────────┘◄───── smallest, or youngest, stack address
+/// ```
+///
+/// # Fibers and async
+///
+/// Wasmtime supports stack-switching with fibers to implement async. This means
+/// that Wasmtime will temporarily execute code on a separate stack and then
+/// suspend from this stack back to the embedder for async operations. Doing
+/// this safely requires manual management of the TLS pointer updated by
+/// Wasmtime.
+///
+/// For example when a fiber is suspended that means that the TLS pointer needs
+/// to be restored to whatever it was when the fiber was resumed. Additionally
+/// this may need to pop multiple `CallThreadState` activations, one for each
+/// one located on the fiber stack itself.
+///
+/// The `AsyncWasmCallState` and `PreviousAsyncWasmCallState` structures in this
+/// module are used to manage this state, namely:
+///
+/// * The `AsyncWasmCallState` structure represents the state of a suspended
+///   fiber. This is a linked list, in reverse order, from oldest activation on
+///   the fiber to youngest activation on the fiber.
+///
+/// * The `PreviousAsyncWasmCallState` structure represents a pointer within our
+///   thread's TLS linked list of activations when a fiber was resumed. This
+///   pointer is used during fiber suspension to know when to stop popping
+///   activations from the thread's linked list.
+///
+/// Note that this means that the directionality of linked list links is
+/// opposite when stored in TLS vs when stored for a suspended fiber. The
+/// thread's current list pointed to by TLS is youngest-to-oldest links, while a
+/// suspended fiber stores oldest-to-youngest links.
 pub(crate) mod tls {
     use super::CallThreadState;
 
@@ -830,6 +941,9 @@ pub(crate) mod tls {
         //
         // When pushed onto a thread this linked list is traversed to get pushed
         // onto the current thread at the time.
+        //
+        // If this pointer is null then that means that the fiber this state is
+        // associated with has no activations.
         state: raw::Ptr,
     }
 
@@ -883,7 +997,7 @@ pub(crate) mod tls {
         ///
         /// This is used when exiting a future in Wasmtime to assert that the
         /// current CallThreadState pointer does not point within the stack
-        /// we're leaving (e.g.  allocated for a fiber).
+        /// we're leaving (e.g. allocated for a fiber).
         pub fn assert_current_state_not_in_range(range: core::ops::Range<usize>) {
             let p = raw::get() as usize;
             assert!(p < range.start || range.end < p);
@@ -892,14 +1006,15 @@ pub(crate) mod tls {
 
     /// Opaque state used to help control TLS state across stack switches for
     /// async support.
+    ///
+    /// This structure is returned from [`AsyncWasmCallState::push`] and
+    /// represents the state of this thread's TLS variable prior to the push
+    /// operation.
     #[cfg(feature = "async")]
     pub struct PreviousAsyncWasmCallState {
-        // The head of a linked list, similar to the TLS state. Note though that
-        // this list is stored in reverse order to assist with `push` and `pop`
-        // below.
-        //
-        // After a `push` call this stores the previous head for the current
-        // thread so we know when to stop popping during a `pop`.
+        // The raw value of this thread's TLS pointer when this structure was
+        // created. This is not dereferenced or inspected but is used to halt
+        // linked list traversal in [`PreviousAsyncWasmCallState::restore`].
         state: raw::Ptr,
     }
 
@@ -909,8 +1024,8 @@ pub(crate) mod tls {
         /// `AsyncWasmCallState`.
         ///
         /// This will pop the top activation of this current thread continuously
-        /// until it reaches whatever the current activation was when `push` was
-        /// originally called.
+        /// until it reaches whatever the current activation was when
+        /// [`AsyncWasmCallState::push`] was originally called.
         ///
         /// # Unsafety
         ///