muxio-tokio-rpc-server: WebSocket RPC Server

Loading…

muxio-tokio-rpc-server: WebSocket RPC Server

Relevant source files

• extensions/muxio-rpc-service-caller/src/caller_interface.rs
• extensions/muxio-tokio-rpc-server/Cargo.toml
• extensions/muxio-tokio-rpc-server/README.md
• extensions/muxio-tokio-rpc-server/src/lib.rs
• extensions/muxio-tokio-rpc-server/src/rpc_server.rs
• extensions/muxio-wasm-rpc-client/src/static_lib/static_transport_bridge.rs

The muxio-tokio-rpc-server crate provides a reference implementation of a Muxio RPC server using the Tokio runtime, Axum web framework, and WebSockets (via tokio-tungstenite). It enables high-performance, bidirectional RPC communication where the server can both handle requests from clients and initiate calls back to connected clients extensions/muxio-tokio-rpc-server/src/rpc_server.rs:1-6

Core Architecture

The server is built around the RpcServer struct, which manages an RpcServiceEndpoint for request dispatching and handles the lifecycle of WebSocket connections. Each connection is managed by two primary asynchronous tasks: a sender_task and a receiver_task extensions/muxio-tokio-rpc-server/src/rpc_server.rs:66-69 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:183-191

Server-to-Code Entity Mapping

The following diagram illustrates how high-level server concepts map to specific structs and functions in the codebase.

Entity Relationship Diagram

graph TD
    subgraph "muxio-tokio-rpc-server"
        RpcServer["RpcServer (struct)"]
ConnectionContext["ConnectionContext (struct)"]
ConnectionContextHandle["ConnectionContextHandle (newtype)"]
WsSenderContext["WsSenderContext (type alias)"]
RpcServer -- "owns" --> RpcServiceEndpoint["RpcServiceEndpoint (from muxio-rpc-service-endpoint)"]
RpcServer -- "spawns" --> receiver_task["RpcServer::receiver_task (fn)"]
ConnectionContext -- "contains" --> RpcDispatcher["RpcDispatcher (from muxio-core)"]
ConnectionContext -- "contains" --> WsSenderContext
        ConnectionContextHandle -- "wraps" --> ConnectionContext
    end

    subgraph "External Dependencies"
        AxumRouter["axum::Router"]
TcpListener["tokio::net::TcpListener"]
end

    RpcServer -- "configures" --> AxumRouter
    RpcServer -- "binds" --> TcpListener

Sources: extensions/muxio-tokio-rpc-server/src/rpc_server.rs:48-69 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:183-191

Key Components

RpcServer

The RpcServer is the entry point for starting the service. It allows for registering RPC handlers via its internal endpoint and provides methods for binding to network interfaces.

• new(event_tx) : Initializes the server. It can optionally take an mpsc::UnboundedSender<RpcServerEvent> to notify the application when clients connect or disconnect extensions/muxio-tokio-rpc-server/src/rpc_server.rs:77-82
• endpoint() : Provides access to the RpcServiceEndpoint, allowing the registration of RpcMethodPrebuffered handlers extensions/muxio-tokio-rpc-server/src/rpc_server.rs:85-87
• serve(addr) / serve_with_listener(listener): Starts the Axum-based web server. It routes WebSocket upgrades to the /ws endpoint extensions/muxio-tokio-rpc-server/src/rpc_server.rs:90-122

Connection Management

Every client connection is represented by a ConnectionContext. Because the server supports bidirectional calling, each connection maintains its own RpcDispatcher extensions/muxio-tokio-rpc-server/src/rpc_server.rs54

• ConnectionContextHandle : A wrapper around Arc<ConnectionContext> used to implement RpcServiceCallerInterface, enabling the server to call methods on the client extensions/muxio-tokio-rpc-server/src/rpc_server.rs:57-60
• RpcServerEvent : An enum used to signal connection lifecycle changes:
  • ClientConnected(ConnectionContextHandle) extensions/muxio-tokio-rpc-server/src/rpc_server.rs44
  • ClientDisconnected(SocketAddr) extensions/muxio-tokio-rpc-server/src/rpc_server.rs45

Data Flow & Task Architecture

The server utilizes a split-task architecture for every WebSocket connection to ensure full-duplex communication and independent heartbeat monitoring.

Connection Data Flow

sequenceDiagram
    participant Peer as "Remote Client"
    participant RX as "receiver_task"
    participant EP as "RpcServiceEndpoint"
    participant TX as "spawn_write_loop"
    participant MPSC as "Internal MPSC Channel"

    Note over RX, TX: Per-Connection Tasks
    
    Peer->>RX: WebSocket Message (Binary)
    RX->>EP: read_bytes(dispatcher, context, bytes, on_emit)
    EP->>EP: Stage 1: Decode & Identify
    EP->>EP: Stage 2: Concurrent Execute Handlers
    EP->>EP: Stage 3: Encode & Emit
    EP-->>MPSC: Response Frames
    MPSC->>TX: Internal Message
    TX->>Peer: WebSocket Message (Binary)

    loop Heartbeat
        RX->>Peer: Ping (Every 5s)
        Peer-->>RX: Pong
    end

Sources: extensions/muxio-tokio-rpc-server/src/rpc_server.rs:183-191 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:35-40 extensions/muxio-rpc-service-caller/src/caller_interface.rs:123-163

sender_task (via spawn_write_loop)

The server uses muxio_rpc_service_caller::write_channel::spawn_write_loop to manage outgoing data extensions/muxio-tokio-rpc-server/src/rpc_server.rs:157-164

• Egress : It listens to an internal mpsc channel for messages generated by RPC handlers or server-initiated calls and forwards them to the WsSenderContext extensions/muxio-tokio-rpc-server/src/rpc_server.rs:49-51 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:157-167

receiver_task & Heartbeats

The receiver_task processes incoming data and manages connection liveness:

• Heartbeat Mechanism : It monitors the connection and expects activity. While the receiver_task logic includes a CLIENT_TIMEOUT (15 seconds), it typically works in tandem with pings sent at HEARTBEAT_INTERVAL (5 seconds) extensions/muxio-tokio-rpc-server/src/rpc_server.rs:35-40
• Message Dispatch : It receives WebSocket messages from the client. Binary data is passed to the RpcServiceEndpointInterface::read_bytes method extensions/muxio-tokio-rpc-server/src/rpc_server.rs23 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:183-191
• Three-Stage Pipeline : read_bytes implements a three-stage pipeline:
  1. Decode : Incoming frames are decoded and finalized requests identified.
  2. Execute : Handlers are executed concurrently for all finalized requests.
  3. Emit : Responses are encoded and emitted back to the transport.
• Timeouts : If no message (including Pong) is received within CLIENT_TIMEOUT, the connection is considered dead and closed extensions/muxio-tokio-rpc-server/src/rpc_server.rs:38-40

Utility Functions

The crate includes utilities for managing TCP listeners, primarily used in testing or dynamic environment setups:

Sources: extensions/muxio-tokio-rpc-server/src/lib.rs5 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:1-200

Implementation Notes

• No Auth : The reference implementation does not include built-in authentication. Users should implement a custom AuthHook within the ws_handler before calling ws.on_upgrade extensions/muxio-tokio-rpc-server/src/rpc_server.rs:1-6 extensions/muxio-tokio-rpc-server/src/rpc_server.rs:130-147
• Concurrency : Request handling is concurrent. The RpcServiceEndpoint processes incoming RPC frames and can execute multiple handlers in parallel, piping results back through the connection’s sender channel extensions/muxio-tokio-rpc-server/src/rpc_server.rs:183-191
• Bidirectional Capability : By implementing RpcServiceCallerInterface on ConnectionContextHandle, the server can use the dispatcher inside the ConnectionContext to initiate calls to the client extensions/muxio-tokio-rpc-server/src/rpc_server.rs:54-60

Sources: extensions/muxio-tokio-rpc-server/src/rpc_server.rs:1-200 extensions/muxio-rpc-service-caller/src/caller_interface.rs:25-53