PR feedback

- Use the order in `MuxerMultistream` as the precedence in multiselect
This commit is contained in:
mhchia
2019-08-20 15:27:07 +08:00
parent 8596f7390f
commit 550c23f9f9
2 changed files with 31 additions and 11 deletions

View File

@ -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)