Proposal: Add a Stream State Machine to py-libp2p

[seetadev]

Description

We propose introducing a formal Stream State mechanism in py-libp2p to track the lifecycle of a Stream. This would allow both internal components and external consumers to determine whether a stream is open, closed, reset, or has errored.

This mirrors behavior in go-libp2p and js-libp2p, and is aligned with good practices from transport protocols like TCP and QUIC.

Motivation

Currently, py-libp2p has no explicit notion of a stream’s state. As a result:

• Errors from operations on a closed or reset stream are only caught at runtime.
• There’s no unified API to check whether a stream is in a valid operational state.
• Bugs like double-close, write-after-close, or inconsistent state transitions are harder to catch.
• This makes reasoning about stream lifecycle harder for transport and application layers.

Introducing a structured StreamState will:

• Improve debuggability and error prevention
• Allow future stream muxers (e.g., Yamux, Mplex) to better coordinate state
• Enable uniform stream lifecycle handling across libp2p implementations

Current Implementation

• There is no state tracking today in the Stream class.
• Streams are managed directly using asyncio reader/writer objects.
• No guards exist for write/read/reset based on state.

Proposed Change

Step 1: Define a StreamState Enum

from enum import Enum, auto

class StreamState(Enum):
    INIT = auto()       # Created, not yet started
    OPEN = auto()       # Ready and transmitting data
    RESET = auto()      # Reset by remote or local peer
    CLOSED = auto()     # Gracefully closed
    ERRORED = auto()    # Unexpected failure

Step 2: Add State to Stream Class

class Stream:
    def __init__(...):
        self._state = StreamState.INIT

Update transitions in:

• start() → OPEN
• reset() → RESET
• close() → CLOSED
• on Exception → ERRORED

Expose via:

@property
def state(self) -> StreamState:
    return self._state

Step 3: Guard Operations

Add guards in critical methods:

async def read(self, n):
    if self._state != StreamState.OPEN:
        raise StreamClosedError("Cannot read from stream; not open")

Apply similar logic to write(), reset(), and close().

---

Example Usage

stream = await libp2p_host.new_stream(peer_id)
print(stream.state)  # Should be OPEN
await stream.close()
print(stream.state)  # Should be CLOSED

---

Design Notes

• Aligns with go-libp2p-core/network.Stream semantics.
• State machine supports future enhancements: e.g., idle, half-close, draining.
• Can help stream muxer implementations (like Yamux) enforce protocol correctness.

Testing Plan

• Unit tests for valid and invalid state transitions
• Tests for read/write/reset after close or reset
• Add integration test with a mock stream peer

Next Steps

• Gather feedback from maintainers/community
• Finalize StreamState enum and transition logic
• Implement in Stream class
• Add tests and documentation

Related Topics

• Feature parity with go-libp2p
• Supports ongoing work in stream muxer abstraction
• Future integration with yamux-py or multiplexers for clean lifecycle management

[acul71]

NetStream Code Analysis

Overview

The NetStream class in libp2p/network/stream/net_stream.py is a comprehensive implementation of network stream management with state tracking, thread safety, and robust error handling. This analysis provides detailed code breakdowns and explanations of every component and their functionality.

Core Architecture

1. Imports and Dependencies

from enum import Enum, auto
import logging
from typing import TYPE_CHECKING
import trio

Code Breakdown:

• Enum, auto: Used for creating the StreamState enum with automatic value assignment
• logging: Provides structured logging for debugging and monitoring stream state changes
• TYPE_CHECKING: Allows conditional imports to avoid circular dependencies
• trio: Async framework that provides trio.Lock() for thread-safe state management

Key Dependencies:

• trio: Async framework providing trio.Lock() for thread safety
• IMuxedStream, INetStream: Abstract interfaces for stream operations
• Stream exceptions: StreamClosed, StreamEOF, StreamError, StreamReset
• QUIC exceptions: QUICStreamClosedError, QUICStreamResetError

2. StreamState Enum

class StreamState(Enum):
    INIT = auto()        # Stream created but not established
    OPEN = auto()        # Stream ready for I/O operations
    CLOSE_READ = auto()  # Read side closed, write may be open
    CLOSE_WRITE = auto() # Write side closed, read may be open
    CLOSE_BOTH = auto()  # Both sides closed, stream terminated
    RESET = auto()       # Stream reset by peer or locally
    ERROR = auto()       # Unrecoverable error condition

Code Breakdown:

• auto(): Automatically assigns unique integer values (0, 1, 2, 3, 4, 5, 6)
• Each state represents a specific phase in the stream lifecycle
• States are designed to match TCP and QUIC stream semantics

State Machine Design:

• Terminal States: RESET and ERROR - no further transitions possible
• Half-closed States: CLOSE_READ, CLOSE_WRITE - essential for QUIC compatibility
• Operational States: INIT, OPEN - allow I/O operations

State Relationships:

• INIT → OPEN: Stream establishment
• OPEN → CLOSE_READ/CLOSE_WRITE: Half-close operations
• CLOSE_READ + CLOSE_WRITE → CLOSE_BOTH: Complete closure
• Any state → RESET/ERROR: Terminal conditions

Class Implementation

3. Constructor and Initialization

def __init__(self, muxed_stream: IMuxedStream, swarm_conn: "SwarmConn | None") -> None:
    self.muxed_stream = muxed_stream
    self.muxed_conn = muxed_stream.muxed_conn
    self.protocol_id = None
    self._state = StreamState.INIT
    self.swarm_conn = swarm_conn
    self.logger = logging.getLogger(__name__)
    
    # Thread safety for state operations (following AkMo3's approach)
    self._state_lock = trio.Lock()

Code Breakdown:

• muxed_stream: IMuxedStream: The underlying stream that handles actual I/O operations
• swarm_conn: "SwarmConn | None": Optional connection reference for stream lifecycle management
• self.muxed_conn = muxed_stream.muxed_conn: Delegates to the underlying connection
• self.protocol_id = None: Protocol identifier, set later via set_protocol()
• self._state = StreamState.INIT: All streams start in INIT state
• self.logger = logging.getLogger(__name__): Creates logger with module name for structured logging
• self._state_lock = trio.Lock(): Creates a new lock instance for this stream's state operations

Key Features:

• Dependency Injection: SwarmConn passed for stream lifecycle management
• State Initialization: Starts in INIT state
• Thread Safety: trio.Lock() for atomic state operations
• Logging: Structured logging for debugging and monitoring

Thread Safety Design:

• Each NetStream instance has its own _state_lock
• This prevents lock contention between different streams
• Lock is only held during state checks and modifications, not during I/O

4. Protocol Management

def get_protocol(self) -> TProtocol | None:
    return self.protocol_id

def set_protocol(self, protocol_id: TProtocol) -> None:
    self.protocol_id = protocol_id

Code Breakdown:

• get_protocol(): Simple getter that returns the current protocol identifier
• set_protocol(protocol_id): Setter that assigns the protocol identifier
• TProtocol: Type alias for protocol identifier (typically a string)
• These methods are synchronous and don't require state locking since they don't affect stream state

Purpose: Manages the protocol identifier for the stream, allowing protocol-specific handling.

Usage Pattern:

• Protocol is typically set after stream establishment but before I/O operations
• Used by the swarm to identify which protocol handler should process the stream
• Enables protocol-specific stream management and routing

5. State Management

State Property (Async)

@property
async def state(self) -> StreamState:
    async with self._state_lock:
        return self._state

Code Breakdown:

• @property: Makes this an async property that can be accessed with await stream.state
• async with self._state_lock:: Acquires the lock before accessing state
• return self._state: Returns the current state value
• Lock is automatically released when exiting the async with block

Thread Safety: State access is protected by lock to prevent race conditions.

Usage Pattern:

current_state = await stream.state
if current_state == StreamState.OPEN:
    # Stream is ready for I/O

State Setting with Logging

async def set_state(self, state: StreamState) -> None:
    async with self._state_lock:
        old_state = self._state
        self._state = state
        
        # Log state transition for debugging and monitoring
        self.logger.debug(f"Stream state transition: {old_state.name} → {state.name}")
        
        # Log important state changes at info level
        if state in [StreamState.ERROR, StreamState.RESET, StreamState.CLOSE_BOTH]:
            self.logger.info(f"Stream entered {state.name} state (from {old_state.name})")

Code Breakdown:

• async with self._state_lock:: Acquires lock for atomic state modification
• old_state = self._state: Captures current state before change for logging
• self._state = state: Updates the state to the new value
• self.logger.debug(...): Logs all state transitions at debug level for debugging
• if state in [StreamState.ERROR, StreamState.RESET, StreamState.CLOSE_BOTH]:: Checks for critical states
• self.logger.info(...): Logs critical state changes at info level for monitoring

Features:

• Atomic Operations: State changes are atomic and thread-safe
• Transition Logging: All state changes are logged for debugging
• Important State Alerts: Critical states logged at info level

Logging Strategy:

• Debug Level: All state transitions (e.g., "OPEN → CLOSE_READ")
• Info Level: Critical states that require attention (ERROR, RESET, CLOSE_BOTH)
• Structured Format: Consistent logging format for easy parsing and analysis

I/O Operations

6. Read Operation

async def read(self, n: int | None = None) -> bytes:
    # Check state atomically to prevent race conditions
    async with self._state_lock:
        if self._state == StreamState.ERROR:
            raise StreamError("Cannot read from stream; stream is in error state")
        elif self._state == StreamState.RESET:
            raise StreamReset("Cannot read from stream; stream is reset")
        elif self._state == StreamState.CLOSE_READ:
            # Allow reads from streams that might still have buffered data
            pass
        # Note: Allow reading from CLOSE_BOTH state - buffered data may exist

    # Perform I/O operation without holding the lock to prevent deadlocks
    try:
        return await self.muxed_stream.read(n)
    except MuxedStreamEOF as error:
        # Handle state transitions when EOF is encountered
        async with self._state_lock:
            if self._state == StreamState.CLOSE_WRITE:
                self._state = StreamState.CLOSE_BOTH
            elif self._state == StreamState.OPEN:
                self._state = StreamState.CLOSE_READ
        raise StreamEOF() from error
    # ... additional exception handling

Code Breakdown:

State Validation Phase:

• async with self._state_lock:: Acquires lock for atomic state checking
• if self._state == StreamState.ERROR:: Prevents operations on error streams
• elif self._state == StreamState.RESET:: Prevents operations on reset streams
• elif self._state == StreamState.CLOSE_READ:: Allows reads from half-closed streams (buffered data)
• Lock is released after state validation, before I/O operation

I/O Operation Phase:

• return await self.muxed_stream.read(n): Delegates to underlying stream
• No lock held during I/O to prevent deadlocks
• This is the critical design decision that prevents deadlocks

EOF Handling:

• except MuxedStreamEOF as error:: Catches end-of-file from underlying stream
• async with self._state_lock:: Re-acquires lock for state transition
• if self._state == StreamState.CLOSE_WRITE:: If write side was closed, now both sides are closed
• elif self._state == StreamState.OPEN:: If stream was open, now read side is closed
• raise StreamEOF() from error:: Re-raises as our exception type

Key Design Decisions:

• Lock Release During I/O: Prevents deadlocks by releasing lock during actual I/O
• State Validation: Atomic state checks before I/O operations
• Graceful EOF Handling: Proper state transitions on EOF
• Exception Hierarchy: Distinguishes between expected and unexpected errors

7. Write Operation

async def write(self, data: bytes) -> None:
    # Check state atomically to prevent race conditions
    async with self._state_lock:
        if self._state == StreamState.ERROR:
            raise StreamError("Cannot write to stream; stream is in error state")
        elif self._state == StreamState.RESET:
            raise StreamReset("Cannot write to stream; stream is reset")
        elif self._state in [StreamState.CLOSE_WRITE, StreamState.CLOSE_BOTH]:
            raise StreamClosed("Cannot write to stream; closed for writing")

    # Perform I/O operation without holding the lock to prevent deadlocks
    try:
        await self.muxed_stream.write(data)
    except (MuxedStreamClosed, QUICStreamClosedError, QUICStreamResetError) as error:
        async with self._state_lock:
            if self._state == StreamState.CLOSE_READ:
                self._state = StreamState.CLOSE_BOTH
            elif self._state == StreamState.OPEN:
                self._state = StreamState.CLOSE_WRITE
        raise StreamClosed() from error
    # ... additional exception handling

Code Breakdown:

State Validation Phase:

• async with self._state_lock:: Acquires lock for atomic state checking
• if self._state == StreamState.ERROR:: Prevents writes to error streams
• elif self._state == StreamState.RESET:: Prevents writes to reset streams
• elif self._state in [StreamState.CLOSE_WRITE, StreamState.CLOSE_BOTH]:: Prevents writes to closed streams
• Lock is released after validation, before I/O operation

I/O Operation Phase:

• await self.muxed_stream.write(data): Delegates to underlying stream
• No lock held during I/O to prevent deadlocks
• Same critical design as read operation

Stream Closure Handling:

• except (MuxedStreamClosed, QUICStreamClosedError, QUICStreamResetError) as error:: Catches stream closure exceptions
• async with self._state_lock:: Re-acquires lock for state transition
• if self._state == StreamState.CLOSE_READ:: If read side was closed, now both sides are closed
• elif self._state == StreamState.OPEN:: If stream was open, now write side is closed
• raise StreamClosed() from error:: Re-raises as our exception type

Features:

• State-based Validation: Prevents operations on invalid streams
• QUIC Compatibility: Handles QUIC-specific exceptions
• Atomic State Transitions: State changes are atomic and consistent

Stream Lifecycle Management

8. Close Operations

Complete Close

async def close(self) -> None:
    """Close stream completely and clean up resources."""
    await self.muxed_stream.close()
    await self.set_state(StreamState.CLOSE_BOTH)
    await self.remove()

Half-close Operations

async def close_read(self) -> None:
    """Close the stream for reading only."""
    async with self._state_lock:
        if self._state == StreamState.ERROR:
            raise StreamError("Cannot close read on stream; stream is in error state")
        elif self._state == StreamState.OPEN:
            self._state = StreamState.CLOSE_READ
        elif self._state == StreamState.CLOSE_WRITE:
            self._state = StreamState.CLOSE_BOTH
            await self.remove()

async def close_write(self) -> None:
    """Close the stream for writing only."""
    async with self._state_lock:
        if self._state == StreamState.ERROR:
            raise StreamError("Cannot close write on stream; stream is in error state")

    await self.muxed_stream.close()
    
    async with self._state_lock:
        if self._state == StreamState.OPEN:
            self._state = StreamState.CLOSE_WRITE
        elif self._state == StreamState.CLOSE_READ:
            self._state = StreamState.CLOSE_BOTH
            await self.remove()

Half-close Support: Essential for QUIC transport where streams can independently close read/write sides.

9. Reset Operation

async def reset(self) -> None:
    """Reset stream."""
    # Allow reset even from ERROR state for cleanup purposes
    try:
        await self.muxed_stream.reset()
    except Exception:
        # If reset fails, we still want to mark the stream as reset
        # This allows cleanup to proceed even if the underlying stream is broken
        pass
    async with self._state_lock:
        self._state = StreamState.RESET
    await self.remove()

Code Breakdown:

• try: await self.muxed_stream.reset(): Attempts to reset the underlying stream
• except Exception:: Catches any exception from the underlying reset operation
• pass: Ignores the exception - we still want to mark the stream as reset
• async with self._state_lock:: Acquires lock for state modification
• self._state = StreamState.RESET:: Sets stream to reset state (terminal)
• await self.remove(): Removes stream from connection and notifies swarm

Robust Cleanup Design:

• Reset operation is designed to succeed even if the underlying stream is broken
• This ensures proper cleanup can proceed even from broken streams
• The pass in the except block is intentional - we want to mark the stream as reset regardless
• This prevents resource leaks when streams are in an unrecoverable state

State Transition Logic:

• Reset can be called from any state, including ERROR state
• This allows cleanup to proceed even from terminal states
• The stream is marked as RESET (terminal state) regardless of underlying stream state

10. Stream Removal

async def remove(self) -> None:
    """Remove the stream from the connection and notify swarm that stream was closed."""
    if self.swarm_conn is not None:
        self.swarm_conn.remove_stream(self)
        await self.swarm_conn.swarm.notify_closed_stream(self)

Code Breakdown:

• if self.swarm_conn is not None:: Checks if we have a swarm connection reference
• self.swarm_conn.remove_stream(self): Removes stream from the connection's stream registry
• await self.swarm_conn.swarm.notify_closed_stream(self): Notifies the swarm about stream closure

Event Notification Design:

• Notifies the swarm about stream closure for proper resource management
• This allows the swarm to clean up any stream-specific resources
• The swarm can update its internal state and notify other components
• This is part of the event-driven architecture for stream lifecycle management

Resource Cleanup:

• Removes stream from connection's active stream list
• Prevents memory leaks by ensuring proper cleanup
• Allows swarm to update connection statistics and state

Utility Methods

11. Remote Address

def get_remote_address(self) -> tuple[str, int] | None:
    """Delegate to the underlying muxed stream."""
    return self.muxed_stream.get_remote_address()

Delegation Pattern: Simple delegation to underlying muxed stream.

12. Operational Check

async def is_operational(self) -> bool:
    """Check if stream is in an operational state."""
    async with self._state_lock:
        return self._state not in [
            StreamState.ERROR,
            StreamState.RESET,
            StreamState.CLOSE_BOTH,
        ]

Code Breakdown:

• async with self._state_lock:: Acquires lock for atomic state checking
• return self._state not in [...]: Returns True if stream is NOT in a terminal/non-operational state
• StreamState.ERROR: Stream is in error state - cannot perform I/O
• StreamState.RESET: Stream is reset - cannot perform I/O
• StreamState.CLOSE_BOTH: Stream is completely closed - cannot perform I/O

State-based Validation: Determines if stream can perform I/O operations based on current state.

Operational States:

• INIT: Stream is created but not yet established (can transition to OPEN)
• OPEN: Stream is ready for I/O operations
• CLOSE_READ: Read side closed, but write side may still be open
• CLOSE_WRITE: Write side closed, but read side may still be open

Non-operational States:

• ERROR: Stream encountered an unrecoverable error
• RESET: Stream was reset by peer or locally
• CLOSE_BOTH: Both sides are closed, stream is terminated

Usage Pattern:

if await stream.is_operational():
    # Stream can perform I/O operations
    data = await stream.read(1024)
else:
    # Stream is in terminal state, cannot perform I/O
    # Handle accordingly

Thread Safety and Concurrency

13. Lock Strategy

Key Principles:

• Atomic State Operations: All state checks and modifications are atomic
• Lock Release During I/O: Prevents deadlocks by releasing locks during I/O operations
• Re-acquisition for State Transitions: Locks are re-acquired for state changes after I/O

Race Condition Prevention:

# Example: Read operation
async with self._state_lock:
    # Atomic state validation
    if self._state == StreamState.ERROR:
        raise StreamError(...)
    
# Lock released during I/O
result = await self.muxed_stream.read(n)

# Lock re-acquired for state transition
async with self._state_lock:
    if self._state == StreamState.OPEN:
        self._state = StreamState.CLOSE_READ

Error Handling Strategy

14. Exception Hierarchy

Known Exceptions (Handled Gracefully):

• MuxedStreamEOF → StreamEOF
• MuxedStreamReset → StreamReset
• MuxedStreamClosed → StreamClosed
• QUICStreamClosedError → StreamClosed
• QUICStreamResetError → StreamReset

Unknown Exceptions (Trigger ERROR State):

• Any unexpected exception → StreamError + ERROR state

15. Detailed Exception Handling in Read Operation

except MuxedStreamEOF as error:
    # Handle state transitions when EOF is encountered
    async with self._state_lock:
        if self._state == StreamState.CLOSE_WRITE:
            self._state = StreamState.CLOSE_BOTH
        elif self._state == StreamState.OPEN:
            self._state = StreamState.CLOSE_READ
    raise StreamEOF() from error
except (
    MuxedStreamReset,
    QUICStreamClosedError,
    QUICStreamResetError,
) as error:
    async with self._state_lock:
        self._state = StreamState.RESET
    raise StreamReset() from error
except Exception as error:
    # Only set ERROR state for truly unexpected errors
    # Known exceptions (MuxedStreamEOF, MuxedStreamReset, etc.)
    # are handled above
    if not isinstance(
        error,
        (
            MuxedStreamEOF,
            MuxedStreamReset,
            QUICStreamClosedError,
            QUICStreamResetError,
            StreamEOF,
            StreamReset,
            StreamClosed,
            ValueError,  # QUIC stream errors
        ),
    ):
        async with self._state_lock:
            self._state = StreamState.ERROR
        raise StreamError(f"Read operation failed: {error}") from error
    # Re-raise known exceptions as-is
    raise

Code Breakdown:

EOF Exception Handler:

• except MuxedStreamEOF as error:: Catches end-of-file from underlying stream
• async with self._state_lock:: Re-acquires lock for state transition
• if self._state == StreamState.CLOSE_WRITE:: If write side was already closed, now both sides are closed
• elif self._state == StreamState.OPEN:: If stream was open, now read side is closed
• raise StreamEOF() from error:: Re-raises as our exception type with original error context

Reset Exception Handler:

• except (MuxedStreamReset, QUICStreamClosedError, QUICStreamResetError) as error:: Catches reset/closure exceptions
• async with self._state_lock:: Re-acquires lock for state transition
• self._state = StreamState.RESET:: Sets stream to reset state (terminal)
• raise StreamReset() from error:: Re-raises as our exception type

Catch-All Exception Handler:

• except Exception as error:: This is the final catch-all for any exceptions not handled above
• if not isinstance(error, (...)):: Only set ERROR state for truly unexpected errors
• Known network exceptions are handled in the previous except blocks
• async with self._state_lock:: Re-acquires lock for state transition
• self._state = StreamState.ERROR:: Sets stream to error state (terminal)
• raise StreamError(f"Read operation failed: {error}") from error:: Re-raises with context
• raise: Re-raises known exceptions as-is without modification

Exception Filtering Logic:

• Known Exceptions: MuxedStreamEOF, MuxedStreamReset, QUICStreamClosedError, QUICStreamResetError
• Our Exceptions: StreamEOF, StreamReset, StreamClosed (already processed)
• QUIC Errors: ValueError (QUIC-specific error handling)
• Unknown Exceptions: Any other exception triggers ERROR state

15. State Transition Logic

EOF Handling:

except MuxedStreamEOF as error:
    async with self._state_lock:
        if self._state == StreamState.CLOSE_WRITE:
            self._state = StreamState.CLOSE_BOTH
        elif self._state == StreamState.OPEN:
            self._state = StreamState.CLOSE_READ
    raise StreamEOF() from error

Reset Handling:

except (MuxedStreamReset, QUICStreamResetError) as error:
    async with self._state_lock:
        self._state = StreamState.RESET
    raise StreamReset() from error

Documentation and Monitoring

16. Comprehensive Docstring

The class includes extensive documentation covering:

• State Machine Diagram: Visual representation of state transitions
• Operation Validity Matrix: Which operations are valid in each state
• Thread Safety Guarantees: Concurrency safety documentation
• QUIC Compatibility: Half-close state support
• Error Handling: Exception handling strategy

17. Logging Strategy

Debug Level: All state transitions
Info Level: Critical state changes (ERROR, RESET, CLOSE_BOTH)

Key Design Decisions

18. AkMo3's Feedback Implementation

Concurrency Safety:

• Implemented trio.Lock() for all state operations
• Atomic state checks and modifications
• Lock release during I/O to prevent deadlocks

Documentation:

• Moved documentation into class docstring
• Removed separate documentation methods
• Comprehensive state machine documentation

Terminal States:

• RESET and ERROR are terminal states
• Removed recover_from_error method (conceptually flawed)
• Proper cleanup even from terminal states

19. QUIC Compatibility

Half-close Support:

• CLOSE_READ and CLOSE_WRITE states
• Independent read/write closure
• Proper state transitions for QUIC streams

20. Performance Considerations

Lock Optimization:

• Minimal lock hold time
• Lock release during I/O operations
• Atomic state operations

Exception Handling:

• Fast path for known exceptions
• Minimal overhead for state transitions
• Efficient error state detection

Summary

The NetStream implementation represents a comprehensive solution for network stream management with:

• Thread Safety: Full concurrency safety with trio.Lock()
• State Management: Complete state machine with proper transitions
• Error Handling: Robust exception handling with state tracking
• QUIC Support: Half-close states for QUIC compatibility
• Documentation: Comprehensive inline documentation
• Performance: Optimized for minimal lock contention
• Monitoring: Extensive logging for debugging and monitoring

This implementation addresses all concerns raised by AkMo3 while maintaining compatibility with existing functionality and providing robust stream management capabilities.