PR feedback

- Use the order in `MuxerMultistream` as the precedence in multiselect

libp2p/security/security_multistream.py

@@ -1,5 +1,6 @@
 from abc import ABC
-from typing import Dict, Mapping
+from collections import OrderedDict
+from typing import Mapping
 
 from libp2p.network.connection.raw_connection_interface import IRawConnection
 from libp2p.peer.id import ID
@@ -25,14 +26,15 @@ class SecurityMultistream(ABC):
     Go implementation: github.com/libp2p/go-conn-security-multistream/ssms.go
     """
 
-    transports: Dict[TProtocol, ISecureTransport]
+    # NOTE: Can be changed to `typing.OrderedDict` since Python 3.7.2.
+    transports: "OrderedDict[TProtocol, ISecureTransport]"
     multiselect: Multiselect
     multiselect_client: MultiselectClient
 
     def __init__(
         self, secure_transports_by_protocol: Mapping[TProtocol, ISecureTransport]
     ) -> None:
-        self.transports = {}
+        self.transports = OrderedDict()
         self.multiselect = Multiselect()
         self.multiselect_client = MultiselectClient()
 
@@ -40,8 +42,17 @@ class SecurityMultistream(ABC):
             self.add_transport(protocol, transport)
 
     def add_transport(self, protocol: TProtocol, transport: ISecureTransport) -> None:
+        """
+        Add a protocol and its corresponding transport to multistream-select(multiselect).
+        The order that a protocol is added is exactly the precedence it is negotiated in
+        multiselect.
+        :param protocol: the protocol name, which is negotiated in multiselect.
+        :param transport: the corresponding transportation to the ``protocol``.
+        """
+        # If protocol is already added before, remove it and add it again.
+        if protocol in self.transports:
+            del self.transports[protocol]
         self.transports[protocol] = transport
-
         # Note: None is added as the handler for the given protocol since
         # we only care about selecting the protocol, not any handler function
         self.multiselect.add_handler(protocol, None)

libp2p/stream_muxer/muxer_multistream.py

@@ -1,4 +1,5 @@
-from typing import Dict, List, Mapping, Type
+from collections import OrderedDict
+from typing import Mapping, Type
 
 from libp2p.network.connection.raw_connection_interface import IRawConnection
 from libp2p.network.typing import GenericProtocolHandlerFn
@@ -23,25 +24,33 @@ class MuxerMultistream:
     go implementation: github.com/libp2p/go-stream-muxer-multistream/multistream.go
     """
 
-    transports: Dict[TProtocol, MuxerClassType]
+    # NOTE: Can be changed to `typing.OrderedDict` since Python 3.7.2.
+    transports: "OrderedDict[TProtocol, MuxerClassType]"
     multiselect: Multiselect
     multiselect_client: MultiselectClient
-    order_preference: List[TProtocol]
 
     def __init__(
         self, muxer_transports_by_protocol: Mapping[TProtocol, MuxerClassType]
     ) -> None:
-        self.transports = {}
+        self.transports = OrderedDict()
         self.multiselect = Multiselect()
         self.multiselect_client = MultiselectClient()
-        self.order_preference = []
         for protocol, transport in muxer_transports_by_protocol.items():
             self.add_transport(protocol, transport)
 
     def add_transport(self, protocol: TProtocol, transport: MuxerClassType) -> None:
+        """
+        Add a protocol and its corresponding transport to multistream-select(multiselect).
+        The order that a protocol is added is exactly the precedence it is negotiated in
+        multiselect.
+        :param protocol: the protocol name, which is negotiated in multiselect.
+        :param transport: the corresponding transportation to the ``protocol``.
+        """
+        # If protocol is already added before, remove it and add it again.
+        if protocol in self.transports:
+            del self.transports[protocol]
         self.transports[protocol] = transport
         self.multiselect.add_handler(protocol, None)
-        self.order_preference.append(protocol)
 
     async def select_transport(self, conn: IRawConnection) -> MuxerClassType:
         """
@@ -54,7 +63,7 @@ class MuxerMultistream:
         communicator = RawConnectionCommunicator(conn)
         if conn.initiator:
             protocol = await self.multiselect_client.select_one_of(
-                self.order_preference, communicator
+                tuple(self.transports.keys()), communicator
             )
         else:
             protocol, _ = await self.multiselect.negotiate(communicator)