This documentation is part of the "Projects with Books" initiative at zenOSmosis.

The source code for this project is available on GitHub.

WebSocket Client and Service Definition

Loading…

WebSocket Client and Service Definition

Relevant source files

The WebSocket implementation in simd-r-drive provides a networked interface to the core storage engine. It leverages the muxio framework for multiplexed RPC over WebSockets, allowing multiple concurrent requests and responses over a single connection. This layer consists of a shared service definition, a high-performance server, and an asynchronous client that implements the standard data store traits.

Service Definition and Serialization

The communication between client and server is governed by a “prebuffered” service definition. This approach uses the bitcode serialization format for high-efficiency encoding of request and response parameters experiments/simd-r-drive-muxio-service-definition/src/prebuffered.rs:1-26

Each RPC method is defined as a struct implementing RpcMethodPrebuffered, which associates a unique METHOD_ID with specific input and output types via the rpc_method_id! macro experiments/simd-r-drive-muxio-service-definition/src/prebuffered/write.rs:15-19

RPC Method Mapping

Method	Request Type	Response Type	Method ID Constant
`write`	`WriteRequestParams`	`WriteResponseParams`	`Write::METHOD_ID`
`read`	`ReadRequestParams`	`ReadResponseParams`	`Read::METHOD_ID`
`batch_write`	`BatchWriteRequestParams`	`BatchWriteResponseParams`	`BatchWrite::METHOD_ID`
`batch_read`	`BatchReadRequestParams`	`BatchReadResponseParams`	`BatchRead::METHOD_ID`
`exists`	`ExistsRequestParams`	`ExistsResponseParams`	`Exists::METHOD_ID`
`len`	`LenRequestParams`	`LenResponseParams`	`Len::METHOD_ID`
`is_empty`	`IsEmptyRequestParams`	`IsEmptyResponseParams`	`IsEmpty::METHOD_ID`
`file_size`	`FileSizeRequestParams`	`FileSizeResponseParams`	`FileSize::METHOD_ID`
`delete`	`DeleteRequestParams`	`DeleteResponseParams`	`Delete::METHOD_ID`

Sources: experiments/simd-r-drive-ws-client/src/ws_client.rs:7-12 experiments/simd-r-drive-muxio-service-definition/src/prebuffered.rs:1-26 experiments/simd-r-drive-muxio-service-definition/src/prebuffered/write.rs:5-22 experiments/simd-r-drive-muxio-service-definition/src/prebuffered/batch_read.rs:5-22 experiments/simd-r-drive-muxio-service-definition/src/prebuffered/batch_write.rs:5-22

WsClient Implementation

The WsClient struct in the simd-r-drive-ws-client crate acts as the primary gateway for remote storage operations. It wraps an Arc<RpcClient> from the muxio-tokio-rpc-client crate experiments/simd-r-drive-ws-client/src/ws_client.rs:20-22

Async Trait Integration

WsClient implements the AsyncDataStoreReader and AsyncDataStoreWriter traits, allowing it to be used interchangeably with other asynchronous storage backends experiments/simd-r-drive-ws-client/src/ws_client.rs:42-127 experiments/simd-r-drive-ws-client/src/ws_client.rs:130-213

Key implementation details:

Data Handling : Because WebSocket transport requires copying data into buffers for serialization, the EntryHandleType for WsClient is defined as Vec<u8> rather than a zero-copy memory map handle experiments/simd-r-drive-ws-client/src/ws_client.rs132
Multiplexing : Requests are dispatched via METHOD::call (e.g., Write::call), which allows the muxio runtime to handle request IDs and response matching experiments/simd-r-drive-ws-client/src/ws_client.rs:56-63
Error Handling : RPC-specific errors are converted to standard std::io::Error via rpc_error_to_io to maintain trait compatibility experiments/simd-r-drive-ws-client/src/ws_client.rs:16-18

Connection Management

The client supports monitoring the underlying WebSocket connection state through set_state_change_handler, which accepts a callback receiving RpcTransportState updates experiments/simd-r-drive-ws-client/src/ws_client.rs:33-38 This is utilized by higher-level bindings to track connection health, such as the Python BaseDataStoreWsClient which updates an AtomicBool on disconnect experiments/bindings/python-ws-client/src/base_ws_client_py.rs:48-53

Data Flow: Client to Server

The following diagram illustrates how a write call is transformed from a trait method into a serialized RPC frame.

Logic to Code Entity Map: Client Request Path

graph TD
    subgraph "ApplicationLayer"
        User["User Code"] -- "ws_client.write(key, payload)" --> WsClient["WsClient (ws_client.rs)"]
end

    subgraph "TraitImplementation"
        WsClient -- "implements" --> ADSW["AsyncDataStoreWriter (traits.rs)"]
WsClient -- "constructs" --> WRP["WriteRequestParams (prebuffered/write.rs)"]
end

    subgraph "ServiceDefinition"
        WRP -- "passed to" --> WriteCall["Write::call (RpcCallPrebuffered)"]
WriteCall -- "uses" --> Encode["Write::encode_request (bitcode)"]
end

    subgraph "TransportLayer"
        Encode -- "binary payload" --> RpcClient["RpcClient (muxio-tokio-rpc-client)"]
RpcClient -- "WebSocket Frame" --> Network["TCP/IP"]
end

Sources: experiments/simd-r-drive-ws-client/src/ws_client.rs:55-67 experiments/simd-r-drive-ws-client/src/ws_client.rs:1-12 experiments/simd-r-drive-ws-client/src/ws_client.rs:20-29 experiments/simd-r-drive-muxio-service-definition/src/prebuffered/write.rs:24-26

Server-Side Execution

The server registers handlers for each METHOD_ID defined in the service definition. When a message arrives, the server decodes the request, performs the storage operation, and encodes the response back to the client.

Request Dispatching and Threading

To support high concurrency, the server typically utilizes tokio::task::spawn_blocking for storage operations. This prevents long-running disk I/O or heavy SIMD computations from blocking the asynchronous executor’s reactor threads. The DataStore is typically wrapped in a tokio::sync::RwLock to allow multiple concurrent readers.

Logic to Code Entity Map: Server Execution Path

graph LR
    subgraph "Network"
        WS["WebSocket Stream"]
end

    subgraph "RPCDispatch"
        Endpoint["RpcServiceEndpoint (muxio-tokio-rpc-server)"] -- "matches METHOD_ID" --> Handler["Registered Closure"]
end

    subgraph "ExecutionContext"
        Handler -- "spawn_blocking" --> Task["Tokio Blocking Task"]
Task -- "decode" --> Decode["Write::decode_request"]
end

    subgraph "StorageEngine"
        Decode -- "params" --> Lock["store.blocking_write()"]
Lock -- "access" --> DS["DataStore::write (storage_engine)"]
end

 
   WS --> Endpoint
    DS -- "tail_offset" --> Task
    Task -- "encode" --> WS

Sources: experiments/simd-r-drive-muxio-service-definition/src/prebuffered/write.rs:28-32 experiments/simd-r-drive-ws-client/src/ws_client.rs:56-64

Summary of Data Flow

The system maintains a strict separation between the transport (WebSocket), the protocol (muxio RPC), and the storage logic (DataStore).

Component	Responsibility	Key Files
Service Definition	Defines serializable params and `METHOD_ID`.	`experiments/simd-r-drive-muxio-service-definition/src/prebuffered.rs`
WebSocket Client	Maps `AsyncDataStore` traits to RPC calls.	`experiments/simd-r-drive-ws-client/src/ws_client.rs`
Python Binding	Wraps Rust `WsClient` for Python access.	`experiments/bindings/python-ws-client/src/base_ws_client_py.rs`

Sources: experiments/simd-r-drive-ws-client/src/ws_client.rs:1-130 experiments/bindings/python-ws-client/src/base_ws_client_py.rs:21-26 experiments/simd-r-drive-muxio-service-definition/src/prebuffered/batch_read.rs:1-44

Keyboard shortcuts

rust-simd-r-drive Documentation