mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-02-10 07:00:54 +00:00
Refactor P2PWebSocketConnection and WebsocketTransport constructors for improved readability. Clean up whitespace and enhance logging for connection management.
This commit is contained in:
@ -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
|
||||||
|
|||||||
@ -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}")
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user