[stdlib] [proposal] When Reference and Coroutine are ready, implement safe & fully async socket package

[martinvuyk]

The basic Idea would be to keep as close an API to Python as possible. But we can do better and actually support async IO from the start and think of WASM and other niche use cases, without raising so much on every function call and making use of all of the power of the Mojo type system.

Every Platform would have it's own private socket implementation with the following interface (will evolve over time):

trait _SocketInterface:
    """Interface for Sockets."""

    # var fd: Arc[FileDescriptor]
    # """The Socket's `Arc[FileDescriptor]`."""

    fn __init__(inout self) raises:
        """Create a new socket object."""
        ...

    fn close(owned self) raises:
        """Closes the Socket."""
        ...

    fn __del__(owned self):
        """Closes the Socket if it's the last reference to its
        `Arc[FileDescriptor]`.
        """
        ...

    @staticmethod
    async fn socketpair() raises -> (Self, Self):
        """Create a pair of socket objects from the sockets returned by the
        platform `socketpair()` function."""
        ...

    async fn send_fds(self, fds: List[FileDescriptor]) -> Bool:
        """Send file descriptor to the socket."""
        ...

    async fn recv_fds(self, maxfds: Int) -> Optional[List[Arc[FileDescriptor]]]:
        """Receive file descriptors from the socket."""
        ...

    @staticmethod
    fn gethostname() -> Optional[String]:
        """Return the current hostname."""
        ...

    @staticmethod
    fn gethostbyname(name: String) -> Optional[Address]:
        """Map a hostname to its Address."""
        ...

    @staticmethod
    fn gethostbyaddr(address: Address) -> Optional[String]:
        """Map an Address to DNS info."""
        ...

    @staticmethod
    fn getservbyname(
        name: String, proto: SockProtocol = SockProtocol.TCP
    ) -> Optional[Address]:
        """Map a service name and a protocol name to a port number."""
        ...

    fn getdefaulttimeout(self) -> Optional[SockTimeout]:
        """Get the default timeout value."""
        ...

    fn setdefaulttimeout(self, value: SockTimeout) -> Bool:
        """Set the default timeout value."""
        ...

    async fn accept(self) -> (Self, Address):
        """Return a new socket representing the connection, and the address of
        the client.
        """
        ...

    @staticmethod
    fn create_connection(
        address: Address,
        timeout: SockTimeout = _DEFAULT_SOCKET_TIMEOUT,
        source_address: Optional[Address] = None,
        *,
        all_errors: Bool = False,
    ) raises -> Self:
        """Connects to an address, with an optional timeout and
        optional source address."""
        ...

    @staticmethod
    fn create_server(
        address: Address,
        *,
        backlog: Optional[Int] = None,
        reuse_port: Bool = False,
        dualstack_ipv6: Bool = False,
    ) raises -> Self:
        """Create a TCP socket and bind it to a specified address."""
        ...

An overarching type must provide as similar an interface to Python's as possible:

struct Socket[
    sock_family: SockFamily = SockFamily.AF_INET,
    sock_type: SockType = SockType.SOCK_STREAM,
    sock_protocol: SockProtocol = SockProtocol.TCP,
    sock_platform: SockPlatform = _get_current_platform(),
](CollectionElement):
    """Struct for using Sockets.

    Parameters:
        sock_family: The socket family e.g. `SockFamily.AF_INET`.
        sock_type: The socket type e.g. `SockType.SOCK_STREAM`.
        sock_protocol: The socket protocol e.g. `SockProtocol.TCP`.
        sock_platform: The socket platform e.g. `SockPlatform.LINUX`.
    """
    ...

Examples:

from socket import Socket


async def main():
    with Socket.create_server(("0.0.0.0", 8000)) as server:
        while True:
            conn, addr = await server.accept()
            ...  # handle new connection

        # TODO: once we have async generators:
        # async for conn, addr in server:
        #     ...  # handle new connection

In the future something like this should be possible:

from multiprocessing import Pool
from socket import Socket, SockAddr


async fn handler(conn: Socket, addr: SockAddr):
    ...

async def main():
    with Socket.create_server(("0.0.0.0", 8000)) as server:
        with Pool() as pool:
            _ = await pool.starmap(handler, server)

Thoughts and Questions I have for discussion

Compatibility with Python

• How much compatibility and at what API layer do we enforce it?
  • Is the endgoal ASGI compatibility or that the socket interface remains as similar to Python's as possible ?
• How do we prepare the infrastructure for those higher level APIs while still
  allowing lower level control if desired for other protocols ?
• Do we develop an async only interface ?

Decisions on the choice of kernel IO protocols

• Should we even develop a syncronous poll model like Python's?
• Is it worth it using io_uring (Linux), kqueue (Unix), IOCP (Windows) ?
  • How much would we need to deviate from Python's APIs ?
    • Do we setup a unified async API like Tigerbeetle's ?
      • Can we keep mostly the same API as Python's but make it async ?
  • How do we deal with external C library dependencies like liburing if we decide to use it ?
    • Do we wait for everything to be implemented in Mojo ? (io_uring project)
  • How portable is an async completion model for WASI and microcontrollers (FreeRTOS & others) ?
    • How do we deal with other platforms without async IO in the kernel ?
      • Can an async API still wrap sync IO ?

Possible steps to approach this

1. Build scaffolding for most important platforms in an extensible manner.
2. Setup a Unified socket interface that all platforms adhere to but constraint on what is currently supported for each.
3. Make sync APIs first with async wrappers, progresively develop async infra.
   1. Make sync TCP work for Linux as a starting point.
   2. Develop sync TCP for other platforms.
   3. Start making things really async under the hood.
   4. Develop other protocols.

(repository where I'm sketching this package out, any help implementing platform specific sockets is welcome :) )