From 550c23f9f9f991ab942133edddf7646d327ae46f Mon Sep 17 00:00:00 2001 From: mhchia Date: Tue, 20 Aug 2019 15:27:07 +0800 Subject: [PATCH] PR feedback - Use the order in `MuxerMultistream` as the precedence in multiselect --- libp2p/security/security_multistream.py | 19 +++++++++++++++---- libp2p/stream_muxer/muxer_multistream.py | 23 ++++++++++++++++------- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/libp2p/security/security_multistream.py b/libp2p/security/security_multistream.py index 6e736c1..ec24b46 100644 --- a/libp2p/security/security_multistream.py +++ b/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) diff --git a/libp2p/stream_muxer/muxer_multistream.py b/libp2p/stream_muxer/muxer_multistream.py index c247e53..703c4e2 100644 --- a/libp2p/stream_muxer/muxer_multistream.py +++ b/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)