Yamux Implementation Details and Interop Guide

[acul71]

Yamux Implementation Details and Interop Guide

Overview

This discussion thread serves as a living document for Yamux stream multiplexer implementation details, design decisions, and interop considerations with the Noise module. This is particularly important as Yamux will be used in various interop efforts.

Recent Fixes (PR #1014)

Issue #930: Connection Closure Handling

Problem: accept_stream() would hang indefinitely when the connection was closed.

Solution:

• Ensured new_stream_send_channel is closed in all shutdown scenarios (start(), close(), _cleanup_on_error())
• accept_stream() now raises MuxedConnUnavailable when the channel closes, matching Mplex and QUIC behavior
• Added proper exception handling for trio.ClosedResourceError and trio.BrokenResourceError

Key Implementation Points:

• Channel closure must happen in multiple places to handle both clean shutdowns and error scenarios
• The start() method closes the channel after the nursery exits, with graceful handling if already closed
• Error handling in handle_incoming() triggers cleanup which closes the channel

Issue #931: FIN/RST Flags on WINDOW_UPDATE Frames

Problem: FIN and RST flags were only processed on DATA frames, but the Yamux spec allows them on WINDOW_UPDATE frames.

Solution:

• Added flag checking in the TYPE_WINDOW_UPDATE handler
• Stream state is now properly updated when FIN/RST arrives on window updates
• Ensures protocol compliance with Yamux specification

Key Implementation Points:

• Flags are checked after processing the window increment
• Stream events are set to wake up waiting readers
• Both FIN (half-close) and RST (reset) are handled correctly

Connection Lifecycle Management

Startup Sequence

1. Yamux.start() is called, which:
   • Creates a nursery for handle_incoming()
   • Sets event_started to signal readiness
   • On exit, closes new_stream_send_channel (with error handling)

Stream Acceptance

• accept_stream() waits on new_stream_receive_channel
• When channel closes (connection termination), raises MuxedConnUnavailable
• This allows callers (e.g., SwarmConn) to properly clean up

Error Handling

• handle_incoming() catches connection errors (RawConnError, OSError, IncompleteReadError, trio.ClosedResourceError, trio.BrokenResourceError)
• Calls _cleanup_on_error() which:
  • Sets shutdown flag
  • Closes the new stream channel (unblocks accept_stream())
  • Marks all streams as closed/reset
  • Wakes up waiting tasks
  • Clears all data structures

Clean Shutdown

• close() sends GO_AWAY frame
• Sets shutdown flag
• Closes channel to unblock accept_stream()
• Cleans up all streams and resources

Protocol Compliance

Flag Handling

Per the Yamux spec, FIN and RST flags can appear on:

• DATA frames (already handled)
• WINDOW_UPDATE frames (now handled in PR Fix: Handle Yamux connection closure and FIN/RST flags #1014)

Stream State Management

• recv_closed: Set when FIN is received (half-close)
• send_closed: Set when local end sends FIN
• closed: Set when both directions are closed
• reset_received: Set when RST is received

Window Management

• Window updates increment send_window
• Flags on WINDOW_UPDATE are processed after window increment
• Events are set to wake up waiting writers/readers

Interop Considerations with Noise Module

Connection Lifecycle

• Noise handshake completes before Yamux starts
• Yamux operates on the secured connection from Noise
• Connection errors propagate through both layers

Error Propagation

• MuxedConnUnavailable is raised when connection closes
• This allows upper layers to detect connection loss
• Matches behavior of Mplex and QUIC for consistency

Resource Cleanup

• Both Noise and Yamux must clean up properly
• Yamux cleanup happens in _cleanup_on_error() and close()
• Ensures no resource leaks during interop scenarios

Testing Strategies

Unit Tests

• Test connection closure scenarios (clean and error)
• Test FIN/RST flag handling on both DATA and WINDOW_UPDATE frames
• Test stream state transitions
• Test error propagation

Platform-Specific Considerations

• Windows: Direct connection closure during active read can cause worker crashes
• Solution: Skip problematic tests on Windows, use alternative test approaches
• Linux: All tests pass reliably

Test Coverage

• test_yamux_accept_stream_unblocks_on_close: Tests clean closure
• test_yamux_accept_stream_unblocks_on_error: Tests error scenarios (skipped on Windows)
• test_yamux_fin_on_window_update: Tests FIN on WINDOW_UPDATE
• test_yamux_rst_on_window_update: Tests RST on WINDOW_UPDATE

Common Pitfalls and Troubleshooting

Hanging accept_stream()

Symptom: accept_stream() never returns after connection closes.

Cause: Channel not closed on connection termination.

Solution: Ensure channel is closed in all shutdown paths (start(), close(), _cleanup_on_error()).

Streams Not Closing

Symptom: Streams remain open after remote sends FIN/RST.

Cause: Flags only checked on DATA frames, not WINDOW_UPDATE.

Solution: Check flags on WINDOW_UPDATE frames per Yamux spec.

Resource Leaks

Symptom: Memory/connection leaks during interop.

Cause: Incomplete cleanup on errors.

Solution: Ensure _cleanup_on_error() clears all data structures and wakes waiting tasks.

Race Conditions

Symptom: Intermittent failures during connection closure.

Cause: Channel closure timing issues.

Solution: Use proper exception handling and ensure idempotent operations (e.g., try/except around channel.close()).

Design Decisions

Exception Type: MuxedConnUnavailable

• Chosen to match Mplex and QUIC behavior
• Allows SwarmConn to handle connection loss consistently
• More specific than generic MuxedStreamError

Channel Closure in Multiple Places

• start(): Handles normal exit
• close(): Handles explicit close
• _cleanup_on_error(): Handles error scenarios
• All use try/except to handle already-closed channels gracefully

Flag Processing Order

• Window increment processed first
• Then flags are checked and stream state updated
• Events set to wake waiting tasks

Future Considerations

Performance Optimization

• Consider batching window updates
• Optimize stream lookup in hot paths
• Profile interop scenarios with Noise

Additional Protocol Features

• Consider implementing additional Yamux features as needed
• Monitor spec updates and compliance requirements

Interop Testing

• Expand interop test coverage with Go implementations
• Test with various Noise configurations
• Validate behavior under network conditions

References

• Yamux Specification
• PR #1014 - Recent fixes
• Issue #930 - Connection closure
• Issue #931 - FIN/RST flags

---

Note: This is a living document. Please update with additional findings, interop results, and implementation details as they emerge.

[acul71]

Yamux Stream State Diagram

stateDiagram-v2
    [*] --> Open: Stream Created\n(SYN sent/received)
    
    Open --> HalfClosedSend: Local sends FIN\n(send_closed=True)
    Open --> HalfClosedRecv: Remote sends FIN\n(recv_closed=True)
    Open --> Reset: RST received/sent\n(reset_received=True)
    Open --> FullyClosed: Both send FIN\nsimultaneously
    
    HalfClosedSend --> FullyClosed: Remote sends FIN\n(recv_closed=True)
    HalfClosedRecv --> FullyClosed: Local sends FIN\n(send_closed=True)
    
    Reset --> [*]: Stream terminated\n(closed=True)
    FullyClosed --> [*]: Stream terminated\n(closed=True)
    
    note right of Open
        State: closed=False
        send_closed=False
        recv_closed=False
        reset_received=False
        Can read and write
    end note
    
    note right of HalfClosedSend
        State: closed=False
        send_closed=True
        recv_closed=False
        Can only read
    end note
    
    note right of HalfClosedRecv
        State: closed=False
        send_closed=False
        recv_closed=True
        Can only write
    end note
    
    note right of FullyClosed
        State: closed=True
        send_closed=True
        recv_closed=True
        No operations allowed
    end note
    
    note right of Reset
        State: closed=True
        send_closed=True
        recv_closed=True
        reset_received=True
        Immediate termination
    end note

Loading

State Descriptions

Open

• Initial state after stream creation (SYN sent/received)
• Both directions are open for reading and writing
• Flags: closed=False, send_closed=False, recv_closed=False, reset_received=False

Half-Closed (Send)

• Local end has sent FIN (half-close)
• Can still receive data from remote
• Flags: send_closed=True, recv_closed=False, closed=False
• Transition: Local calls stream.close() → sends FIN flag

Half-Closed (Receive)

• Remote end has sent FIN (half-close)
• Can still send data to remote
• Flags: send_closed=False, recv_closed=True, closed=False
• Transition: Remote sends FIN flag (on DATA or WINDOW_UPDATE frame)

Fully Closed

• Both ends have sent FIN
• Stream is completely closed
• Flags: closed=True, send_closed=True, recv_closed=True
• Transition: When both send_closed and recv_closed become True

Reset

• Stream was reset (RST flag received or sent)
• Immediate termination, no graceful shutdown
• Flags: closed=True, send_closed=True, recv_closed=True, reset_received=True
• Transition: RST flag received (on DATA or WINDOW_UPDATE frame) or stream.reset() called

State Transitions

From State	Event	To State	Description
Open	Local sends FIN	HalfClosedSend	Local half-close
Open	Remote sends FIN	HalfClosedRecv	Remote half-close
Open	RST received/sent	Reset	Immediate reset
Open	Both send FIN	FullyClosed	Simultaneous close
HalfClosedSend	Remote sends FIN	FullyClosed	Complete close
HalfClosedRecv	Local sends FIN	FullyClosed	Complete close
Any	Connection error	Reset	Error cleanup

Key Implementation Details

1. FIN Flag Handling: FIN can arrive on:

   • DATA frames (with or without data)
   • WINDOW_UPDATE frames (per Yamux spec)

2. RST Flag Handling: RST can arrive on:

   • DATA frames
   • WINDOW_UPDATE frames (per Yamux spec)

3. State Checks:

   • write() checks send_closed - raises error if closed for sending
   • read() checks recv_closed and reset_received - raises EOF or Reset accordingly
   • closed=True only when both directions are closed

4. Error Handling:

   • Connection errors trigger cleanup → all streams set to Reset state
   • Partial data is returned before raising exceptions
   • Stream events are set to wake waiting tasks

Code References

• Stream initialization: YamuxStream.__init__() (lines 73-86)
• Half-close: YamuxStream.close() (lines 302-321)
• Reset: YamuxStream.reset() (lines 323-338)
• FIN handling: Yamux.handle_incoming() (lines 740-790)
• RST handling: Yamux.handle_incoming() (lines 751-757, 799-801)