Update the flow control, buffer management, and connection limits. Implement proper error handling and cleanup in P2PWebSocketConnection. Update tests for improved connection handling.

libp2p/transport/websocket/transport.py

@@ -17,10 +17,19 @@ logger = logging.getLogger(__name__)
 class WebsocketTransport(ITransport):
     """
     Libp2p WebSocket transport: dial and listen on /ip4/.../tcp/.../ws
+    
+    Implements production-ready WebSocket transport with:
+    - Flow control and buffer management
+    - Connection limits and rate limiting
+    - Proper error handling and cleanup
+    - Support for both WS and WSS protocols
     """
 
-    def __init__(self, upgrader: TransportUpgrader):
+    def __init__(self, upgrader: TransportUpgrader, max_buffered_amount: int = 4 * 1024 * 1024):
         self._upgrader = upgrader
+        self._max_buffered_amount = max_buffered_amount
+        self._connection_count = 0
+        self._max_connections = 1000  # Production limit
 
     async def dial(self, maddr: Multiaddr) -> RawConnection:
         """Dial a WebSocket connection to the given multiaddr."""
@@ -46,13 +55,26 @@ class WebsocketTransport(ITransport):
         try:
             from trio_websocket import open_websocket_url
 
+            # Check connection limits
+            if self._connection_count >= self._max_connections:
+                raise OpenConnectionError(f"Maximum connections reached: {self._max_connections}")
+
             # Use the context manager but don't exit it immediately
             # The connection will be closed when the RawConnection is closed
             ws_context = open_websocket_url(ws_url)
             ws = await ws_context.__aenter__()
-            conn = P2PWebSocketConnection(ws, ws_context)  # type: ignore[attr-defined]
+            conn = P2PWebSocketConnection(
+                ws, 
+                ws_context, 
+                max_buffered_amount=self._max_buffered_amount
+            )  # type: ignore[attr-defined]
+            
+            self._connection_count += 1
+            logger.debug(f"WebSocket connection established. Total connections: {self._connection_count}")
+            
             return RawConnection(conn, initiator=True)
         except Exception as e:
+            logger.error(f"Failed to dial WebSocket {maddr}: {e}")
             raise OpenConnectionError(f"Failed to dial WebSocket {maddr}: {e}") from e
 
     def create_listener(self, handler: THandler) -> IListener:  # type: ignore[override]