Connection Lifecycle and State Management

Loading…

Connection Lifecycle and State Management

Relevant source files

• extensions/muxio-rpc-service-caller/src/lib.rs
• extensions/muxio-rpc-service-caller/tests/dynamic_channel_tests.rs
• extensions/muxio-tokio-rpc-client/src/lib.rs
• extensions/muxio-tokio-rpc-client/src/rpc_client.rs
• extensions/muxio-tokio-rpc-client/tests/transport_state_tests.rs
• extensions/muxio-wasm-rpc-client/src/lib.rs
• extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs
• extensions/muxio-wasm-rpc-client/src/static_lib/static_client.rs

Purpose and Scope

This document describes how muxio tracks connection state, manages lifecycle events, and performs automatic cleanup during disconnection across both Tokio-based native clients (5.2) and WASM browser clients (5.3). It covers the RpcTransportState representation, state change handlers, lifecycle transitions, and cleanup mechanisms that ensure pending requests are properly failed and resources are released when connections terminate.

Connection State Representation

RpcTransportState Enum

The RpcTransportState enum defines the two possible connection states:

This enum is shared across all client implementations and is used as the parameter type for state change handlers.

Sources: extensions/muxio-rpc-service-caller/src/transport_state.rs

Atomic Connection Tracking

Both RpcClient and RpcWasmClient maintain an is_connected field of type Arc<AtomicBool> to track the current connection state:

The use of AtomicBool enables lock-free reads from multiple concurrent tasks (Tokio) or callbacks (WASM). The Arc wrapper allows the flag to be shared with background tasks and closures without ownership transfer. State changes use SeqCst (sequentially consistent) ordering to ensure all threads observe changes in a consistent order.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs30 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs23

State Change Handlers

Applications can register a callback to be notified of connection state changes using the set_state_change_handler() method from the RpcServiceCallerInterface trait:

The handler is stored in an RpcTransportStateChangeHandler:

When set, the handler is immediately invoked with the current state if the client is connected. This ensures the application receives the initial Connected event without race conditions.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:22-334 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:13-180

Connection Lifecycle Diagram

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:79-221 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:38-134

Connection Establishment

Tokio Client Initialization

The Tokio RpcClient::new() method establishes the connection and sets initial state:

1. Connects to WebSocket endpoint via connect_async()
2. Splits the stream into sender and receiver
3. Creates MPSC channel for outbound messages
4. Sets is_connected to true atomically
5. Spawns three background tasks:
   • Heartbeat task (sends pings every 1 second)
   • Receive loop (processes incoming WebSocket messages)
   • Send loop (drains outbound MPSC channel)

The client uses Arc::new_cyclic() to allow background tasks to hold weak references, preventing reference cycles that would leak memory.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:111-271

WASM Client Initialization

The WASM RpcWasmClient::new() method creates the client structure but does not establish the connection immediately:

1. Creates a new RpcWasmClient with is_connected set to false
2. Stores the emit callback for sending data to JavaScript
3. Returns the client instance

Connection establishment is triggered later by JavaScript calling the exported handle_connect() method when the browser’s WebSocket onopen event fires:

1. Sets is_connected to true atomically
2. Invokes the state change handler with RpcTransportState::Connected

Sources: extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:27-44

Active Connection Management

Heartbeat Mechanism (Tokio Only)

The Tokio client spawns a dedicated heartbeat task that sends WebSocket ping frames every 1 second to keep the connection alive and detect failures:

If the channel is closed (indicating shutdown), the task exits cleanly. Pong responses from the server are handled by the receive loop.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:139-154

Connection Status Checks

Both implementations provide an is_connected() method via the RpcServiceCallerInterface trait:

This method is checked before initiating RPC calls or emitting data. If the client is disconnected, operations are aborted early to prevent sending data over a closed connection.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:284-296 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:137-166

Disconnection and Cleanup

Shutdown Paths

The Tokio RpcClient implements two shutdown methods:

shutdown_sync()

Used during Drop to ensure cleanup happens synchronously. This method:

1. Swaps is_connected to false using SeqCst ordering
2. Acquires the state change handler lock (best-effort)
3. Invokes the handler with RpcTransportState::Disconnected
4. Does NOT fail pending requests (dispatcher lock not acquired)

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:55-77

shutdown_async()

Used when errors are detected during background task execution. This method:

1. Swaps is_connected to false using SeqCst ordering
2. Acquires the state change handler lock (best-effort)
3. Invokes the handler with RpcTransportState::Disconnected
4. Acquires the dispatcher lock (async)
5. Calls fail_all_pending_requests() to cancel all in-flight RPCs

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:79-108

WASM Disconnect Handling

The WASM client’s handle_disconnect() method performs similar cleanup:

1. Swaps is_connected to false using SeqCst ordering
2. Invokes the state change handler with RpcTransportState::Disconnected
3. Acquires the dispatcher lock
4. Calls fail_all_pending_requests() with FrameDecodeError::ReadAfterCancel

Sources: extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:124-134

Disconnect Sequence Diagram

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:42-108 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:124-134

Pending Request Cleanup

When a connection is lost, all pending RPC requests must be failed to prevent the application from waiting indefinitely. The RpcDispatcher::fail_all_pending_requests() method accomplishes this:

1. Iterates over all pending requests in the dispatcher’s request map
2. For each pending request, extracts the response callback
3. Invokes the callback with a FrameDecodeError (typically ReadAfterCancel)
4. Clears the pending request map

This ensures that all outstanding RPC calls return an error immediately, allowing application code to handle the failure and retry or report the error to the user.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs102 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:130-132

Drop Behavior and Resource Cleanup

Tokio Client Drop Implementation

The Tokio RpcClient implements the Drop trait to ensure proper cleanup when the client is no longer in use:

This implementation:

1. Aborts all background tasks (heartbeat, send loop, receive loop)
2. Calls shutdown_sync() to trigger state change handler
3. Does NOT fail pending requests (to avoid blocking in destructor)

The Drop implementation ensures that even if the application drops the client without explicit cleanup, resources are released and handlers are notified.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:42-52

WASM Client Lifecycle

The WASM RpcWasmClient does not implement Drop because:

1. WASM uses cooperative JavaScript event handling rather than background tasks
2. Connection state is explicitly managed by JavaScript callbacks (handle_connect(), handle_disconnect())
3. No OS-level resources need cleanup (WebSocket is owned by JavaScript)

The JavaScript glue code is responsible for calling handle_disconnect() when the WebSocket closes.

Sources: extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:16-35

Component Relationship Diagram

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:25-108 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:17-134

Platform-Specific Differences

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs

Ordering Guarantees

Both implementations use SeqCst (sequentially consistent) ordering when swapping the is_connected flag:

This ensures:

1. Only one thread/task performs cleanup (the swap returns true only once)
2. All subsequent reads of is_connected see false across all threads
3. No operations are reordered across the swap boundary

The use of Relaxed ordering for reads (is_connected() method) is acceptable because the flag only transitions from true to false, never the reverse, making eventual consistency sufficient for read operations.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:61-284 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:39-138

Testing Connection Lifecycle

The codebase includes comprehensive integration tests verifying lifecycle behavior:

Test: Connection Failure

Verifies that attempting to connect to a non-listening port returns a ConnectionRefused error.

Sources: extensions/muxio-tokio-rpc-client/tests/transport_state_tests.rs:15-31

Test: State Change Handler

Verifies that:

1. Handler is called with Connected when connection is established
2. Handler is called with Disconnected when server closes the connection
3. Events occur in the correct order

Sources: extensions/muxio-tokio-rpc-client/tests/transport_state_tests.rs:33-165

Test: Pending Requests Fail on Disconnect

Verifies that:

1. RPC calls can be initiated while connected
2. If the connection is lost before a response arrives, pending requests fail
3. The error message indicates cancellation or transport failure

This test uses a oneshot::channel to capture the result of a spawned RPC task and validates that it receives an error after the server closes the connection.

Sources: extensions/muxio-tokio-rpc-client/tests/transport_state_tests.rs:167-292

Best Practices

1. Always Set State Change Handler Early : Register the handler immediately after creating the client to avoid missing the initial Connected event.

2. Handle Disconnection Gracefully : Applications should assume connections can fail at any time and implement retry logic or user notifications.

3. Check Connection Before Critical Operations : Use is_connected() to avoid attempting operations that will fail immediately.

4. Avoid Long-Running Handlers : State change handlers should complete quickly to avoid blocking disconnect processing. Spawn separate tasks for expensive operations.

5. WASM: Callhandle_disconnect() Reliably: Ensure JavaScript glue code calls handle_disconnect() in both onclose and onerror WebSocket event handlers.

6. Testing: Use Adequate Delays : Integration tests should allow sufficient time for background tasks to register pending requests before triggering disconnects.

Sources: extensions/muxio-tokio-rpc-client/src/rpc_client.rs:316-334 extensions/muxio-wasm-rpc-client/src/rpc_wasm_client.rs:168-180 extensions/muxio-tokio-rpc-client/tests/transport_state_tests.rs:220-246