Refactor P2PWebSocketConnection and WebsocketTransport constructors for improved readability. Clean up whitespace and enhance logging for connection management.

This commit is contained in:
yashksaini-coder
2025-09-12 03:11:43 +05:30
parent 0271a36316
commit 4fdfdae9fb
3 changed files with 27 additions and 21 deletions

View File

@ -13,12 +13,17 @@ class P2PWebSocketConnection(ReadWriteCloser):
""" """
Wraps a WebSocketConnection to provide the raw stream interface Wraps a WebSocketConnection to provide the raw stream interface
that libp2p protocols expect. that libp2p protocols expect.
Implements production-ready buffer management and flow control Implements production-ready buffer management and flow control
as recommended in the libp2p WebSocket specification. as recommended in the libp2p WebSocket specification.
""" """
def __init__(self, ws_connection: Any, ws_context: Any = None, max_buffered_amount: int = 4 * 1024 * 1024) -> None: def __init__(
self,
ws_connection: Any,
ws_context: Any = None,
max_buffered_amount: int = 4 * 1024 * 1024,
) -> None:
self._ws_connection = ws_connection self._ws_connection = ws_connection
self._ws_context = ws_context self._ws_context = ws_context
self._read_buffer = b"" self._read_buffer = b""
@ -31,23 +36,24 @@ class P2PWebSocketConnection(ReadWriteCloser):
"""Write data with flow control and buffer management""" """Write data with flow control and buffer management"""
if self._closed: if self._closed:
raise IOException("Connection is closed") raise IOException("Connection is closed")
async with self._write_lock: async with self._write_lock:
try: try:
logger.debug(f"WebSocket writing {len(data)} bytes") logger.debug(f"WebSocket writing {len(data)} bytes")
# Check buffer amount for flow control # Check buffer amount for flow control
if hasattr(self._ws_connection, 'bufferedAmount'): if hasattr(self._ws_connection, "bufferedAmount"):
buffered = self._ws_connection.bufferedAmount buffered = self._ws_connection.bufferedAmount
if buffered > self._max_buffered_amount: if buffered > self._max_buffered_amount:
logger.warning(f"WebSocket buffer full: {buffered} bytes") logger.warning(f"WebSocket buffer full: {buffered} bytes")
# In production, you might want to wait or implement backpressure # In production, you might want to
# wait or implement backpressure
# For now, we'll continue but log the warning # For now, we'll continue but log the warning
# Send as a binary WebSocket message # Send as a binary WebSocket message
await self._ws_connection.send_message(data) await self._ws_connection.send_message(data)
logger.debug(f"WebSocket wrote {len(data)} bytes successfully") logger.debug(f"WebSocket wrote {len(data)} bytes successfully")
except Exception as e: except Exception as e:
logger.error(f"WebSocket write failed: {e}") logger.error(f"WebSocket write failed: {e}")
self._closed = True self._closed = True
@ -147,7 +153,7 @@ class P2PWebSocketConnection(ReadWriteCloser):
"""Close the WebSocket connection with proper cleanup""" """Close the WebSocket connection with proper cleanup"""
if self._closed: if self._closed:
return return
self._closed = True self._closed = True
try: try:
# Close the WebSocket connection # Close the WebSocket connection
@ -157,7 +163,7 @@ class P2PWebSocketConnection(ReadWriteCloser):
await self._ws_context.__aexit__(None, None, None) await self._ws_context.__aexit__(None, None, None)
except Exception as e: except Exception as e:
logger.error(f"Error closing WebSocket connection: {e}") logger.error(f"Error closing WebSocket connection: {e}")
def is_closed(self) -> bool: def is_closed(self) -> bool:
"""Check if the connection is closed""" """Check if the connection is closed"""
return self._closed return self._closed

View File

@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
class WebsocketTransport(ITransport): class WebsocketTransport(ITransport):
""" """
Libp2p WebSocket transport: dial and listen on /ip4/.../tcp/.../ws Libp2p WebSocket transport: dial and listen on /ip4/.../tcp/.../ws
Implements production-ready WebSocket transport with: Implements production-ready WebSocket transport with:
- Flow control and buffer management - Flow control and buffer management
- Connection limits and rate limiting - Connection limits and rate limiting
@ -25,7 +25,9 @@ class WebsocketTransport(ITransport):
- Support for both WS and WSS protocols - Support for both WS and WSS protocols
""" """
def __init__(self, upgrader: TransportUpgrader, max_buffered_amount: int = 4 * 1024 * 1024): def __init__(
self, upgrader: TransportUpgrader, max_buffered_amount: int = 4 * 1024 * 1024
):
self._upgrader = upgrader self._upgrader = upgrader
self._max_buffered_amount = max_buffered_amount self._max_buffered_amount = max_buffered_amount
self._connection_count = 0 self._connection_count = 0
@ -57,21 +59,21 @@ class WebsocketTransport(ITransport):
# Check connection limits # Check connection limits
if self._connection_count >= self._max_connections: if self._connection_count >= self._max_connections:
raise OpenConnectionError(f"Maximum connections reached: {self._max_connections}") raise OpenConnectionError(
f"Maximum connections reached: {self._max_connections}"
)
# Use the context manager but don't exit it immediately # Use the context manager but don't exit it immediately
# The connection will be closed when the RawConnection is closed # The connection will be closed when the RawConnection is closed
ws_context = open_websocket_url(ws_url) ws_context = open_websocket_url(ws_url)
ws = await ws_context.__aenter__() ws = await ws_context.__aenter__()
conn = P2PWebSocketConnection( conn = P2PWebSocketConnection(
ws, ws, ws_context, max_buffered_amount=self._max_buffered_amount
ws_context,
max_buffered_amount=self._max_buffered_amount
) # type: ignore[attr-defined] ) # type: ignore[attr-defined]
self._connection_count += 1 self._connection_count += 1
logger.debug(f"WebSocket connection established. Total connections: {self._connection_count}") logger.debug(f"Total connections: {self._connection_count}")
return RawConnection(conn, initiator=True) return RawConnection(conn, initiator=True)
except Exception as e: except Exception as e:
logger.error(f"Failed to dial WebSocket {maddr}: {e}") logger.error(f"Failed to dial WebSocket {maddr}: {e}")

View File

@ -16,8 +16,6 @@ from libp2p.peer.id import ID
from libp2p.peer.peerinfo import PeerInfo from libp2p.peer.peerinfo import PeerInfo
from libp2p.peer.peerstore import PeerStore from libp2p.peer.peerstore import PeerStore
from libp2p.security.insecure.transport import InsecureTransport from libp2p.security.insecure.transport import InsecureTransport
from libp2p.security.noise.transport import Transport as NoiseTransport
from libp2p.crypto.ed25519 import create_new_key_pair as create_ed25519_key_pair
from libp2p.stream_muxer.yamux.yamux import Yamux from libp2p.stream_muxer.yamux.yamux import Yamux
from libp2p.transport.upgrader import TransportUpgrader from libp2p.transport.upgrader import TransportUpgrader
from libp2p.transport.websocket.transport import WebsocketTransport from libp2p.transport.websocket.transport import WebsocketTransport