Feat: Adding Yamux as default multiplexer, keeping Mplex as fallback (#538)

* feat: Replace mplex with yamux as default multiplexer in py-libp2p

* Retain Mplex alongside Yamux in new_swarm with messaging that Yamux is preferred

* moved !BBHII to a constant YAMUX_HEADER_FORMAT at the top of yamux.py with a comment explaining its structure

* renamed the news fragment to 534.feature.rst and updated the description

* renamed the news fragment to 534.feature.rst and updated the description

* added a docstring to clarify that Yamux does not support deadlines natively

* Remove the __main__ block entirely from test_yamux.py

* Replaced the print statements in test_yamux.py with logging.debug

* Added a comment linking to the spec for clarity

* Raise NotImplementedError in YamuxStream.set_deadline per review

* Add muxed_conn to YamuxStream and test deadline NotImplementedError

* Fix Yamux implementation to meet libp2p spec

* Fix None handling in YamuxStream.read and Yamux.read_stream

* Fix test_connected_peers.py to correctly handle peer connections

* fix: Ensure StreamReset is raised on read after local reset in yamux

* fix: Map MuxedStreamError to StreamClosed in NetStream.write for Yamux

* fix: Raise MuxedStreamReset in Yamux.read_stream for closed streams

* fix: Correct Yamux stream read behavior for NetStream tests

Fixed 	est_net_stream_read_after_remote_closed by updating NetStream.read to raise StreamEOF when the stream is remotely closed and no data is available, aligning with test expectations and Fixed 	est_net_stream_read_until_eof by modifying YamuxStream.read to block until the stream is closed (
ecv_closed=True) for
=-1 reads, ensuring data is only returned after remote closure.

* fix: Correct Yamux stream read behavior for NetStream tests

Fixed 	est_net_stream_read_after_remote_closed by updating NetStream.read to raise StreamEOF when the stream is remotely closed and no data is available, aligning with test expectations and Fixed 	est_net_stream_read_until_eof by modifying YamuxStream.read to block until the stream is closed (
ecv_closed=True) for
=-1 reads, ensuring data is only returned after remote closure.

* fix: raise StreamEOF when reading from closed stream with empty buffer

* fix: prioritize returning buffered data even after stream reset

* fix: prioritize returning buffered data even after stream reset

* fix: Ensure test_net_stream_read_after_remote_closed_and_reset passes in full suite

* fix: Add __init__.py to yamux module to fix documentation build

* fix: Add __init__.py to yamux module to fix documentation build

* fix: Add libp2p.stream_muxer.yamux to libp2p.stream_muxer.rst toctree

* fix: Correct title underline length in libp2p.stream_muxer.yamux.rst

* fix: Add a = so that is matches the libp2p.stream\_muxer.yamux length

* fix(tests): Resolve race condition in network notification test

* fix: fixing failing tests and examples with yamux and noise

* refactor: remove debug logging and improve x25519 tests

* fix: Add functionality for users to choose between Yamux and Mplex

* fix: increased trio sleep to 0.1 sec for slow environment

* feat: Add test for switching between Yamux and mplex

* refactor: move host fixtures to interop tests

* chore: Update __init__.py removing unused import

removed unused
```python
import os
import logging
```

* lint: fix import order

* fix: Resolve conftest.py conflict by removing trio test support

* fix: Resolve test skipping by keeping trio test support

* Fix: add a newline at end of the file

---------

Co-authored-by: acul71 <[email protected]>
Co-authored-by: acul71 <[email protected]>

Author: paschal533

docs/libp2p.stream_muxer.rst

@@ -8,6 +8,7 @@ Subpackages
    :maxdepth: 4
 
    libp2p.stream_muxer.mplex
+   libp2p.stream_muxer.yamux
 
 Submodules
 ----------

docs/libp2p.stream_muxer.yamux.rst

@@ -0,0 +1,7 @@
+libp2p.stream\_muxer.yamux
+==========================
+
+.. automodule:: libp2p.stream_muxer.yamux
+   :members:
+   :undoc-members:
+   :show-inheritance:

examples/identify_push/identify_push_demo.py

@@ -85,6 +85,52 @@ async def main() -> None:
         logger.info("Host 2 connected to Host 1")
         print("Host 2 successfully connected to Host 1")
 
+        # Run the identify protocol from host_2 to host_1
+        # (so Host 1 learns Host 2's address)
+        from libp2p.identity.identify.identify import ID as IDENTIFY_PROTOCOL_ID
+
+        stream = await host_2.new_stream(host_1.get_id(), (IDENTIFY_PROTOCOL_ID,))
+        response = await stream.read()
+        await stream.close()
+
+        # Run the identify protocol from host_1 to host_2
+        # (so Host 2 learns Host 1's address)
+        stream = await host_1.new_stream(host_2.get_id(), (IDENTIFY_PROTOCOL_ID,))
+        response = await stream.read()
+        await stream.close()
+
+        # --- NEW CODE: Update Host 1's peerstore with Host 2's addresses ---
+        from libp2p.identity.identify.pb.identify_pb2 import (
+            Identify,
+        )
+
+        identify_msg = Identify()
+        identify_msg.ParseFromString(response)
+        peerstore_1 = host_1.get_peerstore()
+        peer_id_2 = host_2.get_id()
+        for addr_bytes in identify_msg.listen_addrs:
+            maddr = multiaddr.Multiaddr(addr_bytes)
+            # TTL can be any positive int
+            peerstore_1.add_addr(
+                peer_id_2,
+                maddr,
+                ttl=3600,
+            )
+        # --- END NEW CODE ---
+
+        # Now Host 1's peerstore should have Host 2's address
+        peerstore_1 = host_1.get_peerstore()
+        peer_id_2 = host_2.get_id()
+        addrs_1_for_2 = peerstore_1.addrs(peer_id_2)
+        logger.info(
+            f"[DEBUG] Host 1 peerstore addresses for Host 2 before push: "
+            f"{addrs_1_for_2}"
+        )
+        print(
+            f"[DEBUG] Host 1 peerstore addresses for Host 2 before push: "
+            f"{addrs_1_for_2}"
+        )
+
         # Push identify information from host_1 to host_2
         logger.info("Host 1 pushing identify information to Host 2")
         print("\nHost 1 pushing identify information to Host 2...")
@@ -104,6 +150,9 @@ async def main() -> None:
             logger.error(f"Error during identify push: {str(e)}")
             print(f"\nError during identify push: {str(e)}")
 
+        # Add this at the end of your async with block:
+        await trio.sleep(0.5)  # Give background tasks time to finish
+
 
 if __name__ == "__main__":
     trio.run(main)

libp2p/__init__.py

@@ -1,17 +1,29 @@
+from collections.abc import (
+    Mapping,
+)
 from importlib.metadata import version as __version
+from typing import (
+    Literal,
+    Optional,
+    Type,
+    cast,
+)
 
 from libp2p.abc import (
     IHost,
+    IMuxedConn,
     INetworkService,
     IPeerRouting,
     IPeerStore,
+    ISecureTransport,
 )
 from libp2p.crypto.keys import (
     KeyPair,
 )
 from libp2p.crypto.rsa import (
     create_new_key_pair,
 )
+from libp2p.crypto.x25519 import create_new_key_pair as create_new_x25519_key_pair
 from libp2p.custom_types import (
     TMuxerOptions,
     TProtocol,
@@ -36,11 +48,17 @@
     PLAINTEXT_PROTOCOL_ID,
     InsecureTransport,
 )
+from libp2p.security.noise.transport import PROTOCOL_ID as NOISE_PROTOCOL_ID
+from libp2p.security.noise.transport import Transport as NoiseTransport
 import libp2p.security.secio.transport as secio
 from libp2p.stream_muxer.mplex.mplex import (
     MPLEX_PROTOCOL_ID,
     Mplex,
 )
+from libp2p.stream_muxer.yamux.yamux import (
+    Yamux,
+)
+from libp2p.stream_muxer.yamux.yamux import PROTOCOL_ID as YAMUX_PROTOCOL_ID
 from libp2p.transport.tcp.tcp import (
     TCP,
 )
@@ -54,6 +72,60 @@
 # Initialize logging configuration
 setup_logging()
 
+# Default multiplexer choice
+DEFAULT_MUXER = "YAMUX"
+
+# Multiplexer options
+MUXER_YAMUX = "YAMUX"
+MUXER_MPLEX = "MPLEX"
+
+
+def set_default_muxer(muxer_name: Literal["YAMUX", "MPLEX"]) -> None:
+    """
+    Set the default multiplexer protocol to use.
+
+    :param muxer_name: Either "YAMUX" or "MPLEX"
+    :raise ValueError: If an unsupported muxer name is provided
+    """
+    global DEFAULT_MUXER
+    muxer_upper = muxer_name.upper()
+    if muxer_upper not in [MUXER_YAMUX, MUXER_MPLEX]:
+        raise ValueError(f"Unknown muxer: {muxer_name}. Use 'YAMUX' or 'MPLEX'.")
+    DEFAULT_MUXER = muxer_upper
+
+
+def get_default_muxer() -> str:
+    """
+    Returns the currently selected default muxer.
+
+    :return: Either "YAMUX" or "MPLEX"
+    """
+    return DEFAULT_MUXER
+
+
+def create_yamux_muxer_option() -> TMuxerOptions:
+    """
+    Returns muxer options with Yamux as the primary choice.
+
+    :return: Muxer options with Yamux first
+    """
+    return {
+        TProtocol(YAMUX_PROTOCOL_ID): Yamux,  # Primary choice
+        TProtocol(MPLEX_PROTOCOL_ID): Mplex,  # Fallback for compatibility
+    }
+
+
+def create_mplex_muxer_option() -> TMuxerOptions:
+    """
+    Returns muxer options with Mplex as the primary choice.
+
+    :return: Muxer options with Mplex first
+    """
+    return {
+        TProtocol(MPLEX_PROTOCOL_ID): Mplex,  # Primary choice
+        TProtocol(YAMUX_PROTOCOL_ID): Yamux,  # Fallback
+    }
+
 
 def generate_new_rsa_identity() -> KeyPair:
     return create_new_key_pair()
@@ -64,11 +136,24 @@ def generate_peer_id_from(key_pair: KeyPair) -> ID:
     return ID.from_pubkey(public_key)
 
 
+def get_default_muxer_options() -> TMuxerOptions:
+    """
+    Returns the default muxer options based on the current default muxer setting.
+
+    :return: Muxer options with the preferred muxer first
+    """
+    if DEFAULT_MUXER == "MPLEX":
+        return create_mplex_muxer_option()
+    else:  # YAMUX is default
+        return create_yamux_muxer_option()
+
+
 def new_swarm(
-    key_pair: KeyPair = None,
-    muxer_opt: TMuxerOptions = None,
-    sec_opt: TSecurityOptions = None,
-    peerstore_opt: IPeerStore = None,
+    key_pair: Optional[KeyPair] = None,
+    muxer_opt: Optional[TMuxerOptions] = None,
+    sec_opt: Optional[TSecurityOptions] = None,
+    peerstore_opt: Optional[IPeerStore] = None,
+    muxer_preference: Optional[Literal["YAMUX", "MPLEX"]] = None,
 ) -> INetworkService:
     """
     Create a swarm instance based on the parameters.
@@ -77,7 +162,13 @@ def new_swarm(
     :param muxer_opt: optional choice of stream muxer
     :param sec_opt: optional choice of security upgrade
     :param peerstore_opt: optional peerstore
+    :param muxer_preference: optional explicit muxer preference
     :return: return a default swarm instance
+
+    Note: Yamux (/yamux/1.0.0) is the preferred stream multiplexer
+          due to its improved performance and features.
+          Mplex (/mplex/6.7.0) is retained for backward compatibility
+          but may be deprecated in the future.
     """
     if key_pair is None:
         key_pair = generate_new_rsa_identity()
@@ -87,13 +178,41 @@ def new_swarm(
     # TODO: Parse `listen_addrs` to determine transport
     transport = TCP()
 
-    muxer_transports_by_protocol = muxer_opt or {MPLEX_PROTOCOL_ID: Mplex}
-    security_transports_by_protocol = sec_opt or {
-        TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair),
+    # Generate X25519 keypair for Noise
+    noise_key_pair = create_new_x25519_key_pair()
+
+    # Default security transports (using Noise as primary)
+    secure_transports_by_protocol: Mapping[TProtocol, ISecureTransport] = sec_opt or {
+        NOISE_PROTOCOL_ID: NoiseTransport(
+            key_pair, noise_privkey=noise_key_pair.private_key
+        ),
         TProtocol(secio.ID): secio.Transport(key_pair),
+        TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair),
     }
+
+    # Use given muxer preference if provided, otherwise use global default
+    if muxer_preference is not None:
+        temp_pref = muxer_preference.upper()
+        if temp_pref not in [MUXER_YAMUX, MUXER_MPLEX]:
+            raise ValueError(
+                f"Unknown muxer: {muxer_preference}. Use 'YAMUX' or 'MPLEX'."
+            )
+        active_preference = temp_pref
+    else:
+        active_preference = DEFAULT_MUXER
+
+    # Use provided muxer options if given, otherwise create based on preference
+    if muxer_opt is not None:
+        muxer_transports_by_protocol = muxer_opt
+    else:
+        if active_preference == MUXER_MPLEX:
+            muxer_transports_by_protocol = create_mplex_muxer_option()
+        else:  # YAMUX is default
+            muxer_transports_by_protocol = create_yamux_muxer_option()
+
     upgrader = TransportUpgrader(
-        security_transports_by_protocol, muxer_transports_by_protocol
+        secure_transports_by_protocol=secure_transports_by_protocol,
+        muxer_transports_by_protocol=muxer_transports_by_protocol,
     )
 
     peerstore = peerstore_opt or PeerStore()
@@ -104,11 +223,12 @@ def new_swarm(
 
 
 def new_host(
-    key_pair: KeyPair = None,
-    muxer_opt: TMuxerOptions = None,
-    sec_opt: TSecurityOptions = None,
-    peerstore_opt: IPeerStore = None,
-    disc_opt: IPeerRouting = None,
+    key_pair: Optional[KeyPair] = None,
+    muxer_opt: Optional[TMuxerOptions] = None,
+    sec_opt: Optional[TSecurityOptions] = None,
+    peerstore_opt: Optional[IPeerStore] = None,
+    disc_opt: Optional[IPeerRouting] = None,
+    muxer_preference: Optional[Literal["YAMUX", "MPLEX"]] = None,
 ) -> IHost:
     """
     Create a new libp2p host based on the given parameters.
@@ -118,20 +238,20 @@ def new_host(
     :param sec_opt: optional choice of security upgrade
     :param peerstore_opt: optional peerstore
     :param disc_opt: optional discovery
+    :param muxer_preference: optional explicit muxer preference
     :return: return a host instance
     """
     swarm = new_swarm(
         key_pair=key_pair,
         muxer_opt=muxer_opt,
         sec_opt=sec_opt,
         peerstore_opt=peerstore_opt,
+        muxer_preference=muxer_preference,
     )
-    host: IHost
-    if disc_opt:
-        host = RoutedHost(swarm, disc_opt)
-    else:
-        host = BasicHost(swarm)
-    return host
+
+    if disc_opt is not None:
+        return RoutedHost(swarm, disc_opt)
+    return BasicHost(swarm)
 
 
 __version__ = __version("libp2p")