From b9fe575878a9748a4bfac63e8d21d68f2472e34f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mi=C5=82osz=20Rachwa=C5=82?= <me@milek7.pl>
Date: Wed, 7 Feb 2024 14:32:22 +0100
Subject: [PATCH] Allocate thread IDs for wasi_thread_spawn sequentially

---
 crates/wasi-threads/src/lib.rs | 39 +++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/crates/wasi-threads/src/lib.rs b/crates/wasi-threads/src/lib.rs
index 76f12edae86a..26fb7676faff 100644
--- a/crates/wasi-threads/src/lib.rs
+++ b/crates/wasi-threads/src/lib.rs
@@ -3,8 +3,8 @@
 //! [`wasi-threads`]: https://github.com/WebAssembly/wasi-threads
 
 use anyhow::{anyhow, Result};
-use rand::Rng;
 use std::panic::{catch_unwind, AssertUnwindSafe};
+use std::sync::atomic::{AtomicI32, Ordering};
 use std::sync::Arc;
 use std::thread;
 use wasmtime::{Caller, ExternType, InstancePre, Linker, Module, SharedMemory, Store, ValType};
@@ -16,12 +16,14 @@ const WASI_ENTRY_POINT: &str = "wasi_thread_start";
 
 pub struct WasiThreadsCtx<T> {
     instance_pre: Arc<InstancePre<T>>,
+    tid: AtomicI32,
 }
 
 impl<T: Clone + Send + 'static> WasiThreadsCtx<T> {
     pub fn new(module: Module, linker: Arc<Linker<T>>) -> Result<Self> {
         let instance_pre = Arc::new(linker.instantiate_pre(&module)?);
-        Ok(Self { instance_pre })
+        let tid = AtomicI32::new(0);
+        Ok(Self { instance_pre, tid })
     }
 
     pub fn spawn(&self, host: T, thread_start_arg: i32) -> Result<i32> {
@@ -45,8 +47,14 @@ impl<T: Clone + Send + 'static> WasiThreadsCtx<T> {
             return Ok(-1);
         }
 
+        let wasi_thread_id = self.next_thread_id();
+        if wasi_thread_id.is_none() {
+            log::error!("ran out of valid thread IDs");
+            return Ok(-1);
+        }
+        let wasi_thread_id = wasi_thread_id.unwrap();
+
         // Start a Rust thread running a new instance of the current module.
-        let wasi_thread_id = random_thread_id();
         let builder = thread::Builder::new().name(format!("wasi-thread-{}", wasi_thread_id));
         builder.spawn(move || {
             // Catch any panic failures in host code; e.g., if a WASI module
@@ -89,16 +97,23 @@ impl<T: Clone + Send + 'static> WasiThreadsCtx<T> {
 
         Ok(wasi_thread_id)
     }
-}
 
-/// Helper for generating valid WASI thread IDs (TID).
-///
-/// Callers of `wasi_thread_spawn` expect a TID >=0 to indicate a successful
-/// spawning of the thread whereas a negative return value indicates an
-/// failure to spawn.
-fn random_thread_id() -> i32 {
-    let tid: u32 = rand::thread_rng().gen();
-    (tid >> 1) as i32
+    /// Helper for generating valid WASI thread IDs (TID).
+    ///
+    /// Callers of `wasi_thread_spawn` expect a TID in range of 0 < TID <= 0x1FFFFFFF
+    /// to indicate a successful spawning of the thread whereas a negative
+    /// return value indicates an failure to spawn.
+    fn next_thread_id(&self) -> Option<i32> {
+        match self
+            .tid
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| match v {
+                ..=0x1ffffffe => Some(v + 1),
+                _ => None,
+            }) {
+            Ok(v) => Some(v + 1),
+            Err(_) => None,
+        }
+    }
 }
 
 /// Manually add the WASI `thread_spawn` function to the linker.