From 19bca4cae4ad8b9ad87aabf380e176b168f6197b Mon Sep 17 00:00:00 2001
From: Joel Dice <joel.dice@fermyon.com>
Date: Wed, 9 Oct 2024 08:22:35 -0600
Subject: [PATCH] add wasm32-wasip2 support

This adds support for the new `wasm32-wasip2` target platform, which includes
more extensive support for sockets than `wasm32-wasip1` (formerly known as
`wasm32-wasi`).

The bulk of the changes are in https://github.com/tokio-rs/mio/pull/1836.  This
patch just tweaks a few `cfg` directives to indicate `wasm32-wasip2`'s
additional capabilities.

In the future, we could consider adding support for `ToSocketAddrs`.  WASIp2
natively supports asynchronous DNS lookups and is single threaded, whereas Tokio
currently assumes DNS lookups are blocking and require multithreading to emulate
async lookups.  A WASIp2-specific implementation could do the lookup directly
without multithreading.

I've tested this end-to-end using https://github.com/dicej/wasi-sockets-tests,
which includes smoke tests for `mio`, `tokio`, `tokio-postgres`, etc.  I'd also
be happy to add tests to this repo if appropriate; it would require adding a
dev-dependency on e.g. `wasmtime` to actually run the test cases.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>
---
 tokio/Cargo.toml            | 3 ++-
 tokio/src/lib.rs            | 1 +
 tokio/src/macros/cfg.rs     | 9 +++++++++
 tokio/src/net/mod.rs        | 2 +-
 tokio/src/net/tcp/stream.rs | 7 +++++--
 5 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml
index b07f50150b7..2ce4ee5e418 100644
--- a/tokio/Cargo.toml
+++ b/tokio/Cargo.toml
@@ -92,7 +92,8 @@ pin-project-lite = "0.2.11"
 
 # Everything else is optional...
 bytes = { version = "1.0.0", optional = true }
-mio = { version = "1.0.1", optional = true, default-features = false }
+# TODO dicej: switch to upstream once WASIp2 support is merged:
+mio = { git = "https://github.com/dicej/mio", branch = "wasip2", optional = true, default-features = false }
 parking_lot = { version = "0.12.0", optional = true }
 
 [target.'cfg(not(target_family = "wasm"))'.dependencies]
diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs
index bd85e28510f..20bf6b74ee5 100644
--- a/tokio/src/lib.rs
+++ b/tokio/src/lib.rs
@@ -19,6 +19,7 @@
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![cfg_attr(docsrs, allow(unused_attributes))]
 #![cfg_attr(loom, allow(dead_code, unreachable_pub))]
+#![cfg_attr(all(target_os = "wasi", target_env = "p2"), feature(wasip2))]
 
 //! A runtime for writing reliable network applications without compromising speed.
 //!
diff --git a/tokio/src/macros/cfg.rs b/tokio/src/macros/cfg.rs
index 3242d3ce2ea..b0b14ce11aa 100644
--- a/tokio/src/macros/cfg.rs
+++ b/tokio/src/macros/cfg.rs
@@ -594,6 +594,15 @@ macro_rules! cfg_not_wasi {
     }
 }
 
+macro_rules! cfg_not_wasip1 {
+    ($($item:item)*) => {
+        $(
+            #[cfg(any(not(target_os = "wasi"), target_env = "p2"))]
+            $item
+        )*
+    }
+}
+
 macro_rules! cfg_is_wasm_not_wasi {
     ($($item:item)*) => {
         $(
diff --git a/tokio/src/net/mod.rs b/tokio/src/net/mod.rs
index d5562ac5d47..c2b0c4e20c0 100644
--- a/tokio/src/net/mod.rs
+++ b/tokio/src/net/mod.rs
@@ -29,7 +29,7 @@
 //! [`AsyncFd`]: crate::io::unix::AsyncFd
 
 mod addr;
-cfg_not_wasi! {
+cfg_not_wasip1! {
     #[cfg(feature = "net")]
     pub(crate) use addr::to_socket_addrs;
 }
diff --git a/tokio/src/net/tcp/stream.rs b/tokio/src/net/tcp/stream.rs
index 2714592d3a3..efd517242ed 100644
--- a/tokio/src/net/tcp/stream.rs
+++ b/tokio/src/net/tcp/stream.rs
@@ -1,7 +1,10 @@
 cfg_not_wasi! {
+    use std::time::Duration;
+}
+
+cfg_not_wasip1! {
     use crate::net::{to_socket_addrs, ToSocketAddrs};
     use std::future::poll_fn;
-    use std::time::Duration;
 }
 
 use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready};
@@ -72,7 +75,7 @@ cfg_net! {
 }
 
 impl TcpStream {
-    cfg_not_wasi! {
+    cfg_not_wasip1! {
         /// Opens a TCP connection to a remote host.
         ///
         /// `addr` is an address of the remote host. Anything which implements the