# Buggy code
def stream_protocol_messages(
    self,
    protocol_identifier: Union[ProtocolAPI, Type[ProtocolAPI]],
) -> AsyncIterator[CommandAPI[Any]]:
    """
    Stream the messages for the specified protocol.
    """
    if isinstance(protocol_identifier, ProtocolAPI):
        protocol_class = type(protocol_identifier)
    elif isinstance(protocol_identifier, type) and issubclass(
        protocol_identifier, ProtocolAPI
    ):
        protocol_class = protocol_identifier
    else:
        raise TypeError("Unknown protocol identifier: {protocol}")

    if not self.has_protocol(protocol_class):
        raise UnknownProtocol(f"Unknown protocol '{protocol_class}'")

    if self._protocol_locks.is_locked(protocol_class):
        raise Exception(f"Streaming lock for {protocol_class} is not free.")
    elif not self._multiplex_lock.locked():
        raise Exception("Not multiplexed.")

    # Mostly a sanity check but this ensures we do better than accidentally
    # raising an attribute error in whatever race conditions or edge cases
    # potentially make the `_multiplex_token` unavailable.
    if not hasattr(self, "_multiplex_token"):
        raise Exception("No cancel token found for multiplexing.")

    # We do the wait_iter here so that the call sites in the handshakers
    # that use this don't need to be aware of cancellation tokens.
    return self.wait_iter(
        self._stream_protocol_messages(protocol_class),
        token=self._multiplex_token,
    )


# Traceback (KeyError)
DEBUG  12-18 15:21:47  async_process_runner.py  b'\x1b[1m\x1b[31m   ERROR  2019-12-18 15:21:47,268           LESPeerPool  unexpected error during peer connection\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'Traceback (most recent call last):\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/p2p/peer_pool.py", line 209, in maybe_connect_more_peers\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    for backend in self.peer_backends\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/p2p/peer_pool.py", line 191, in _add_peers_from_backend\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    await self.connect_to_nodes(iter(candidates))\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/p2p/peer_pool.py", line 419, in connect_to_nodes\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    loop=self.get_event_loop(),\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/p2p/peer_pool.py", line 440, in connect_to_node\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    async with self.lock_node_for_handshake(node):\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/p2p/peer_pool.py", line 423, in lock_node_for_handshake\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    return self._handshake_locks.lock(node)\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/p2p/resource_lock.py", line 18, in lock\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    lock = self._locks[resource]\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'  File "/home/circleci/repo/.tox/py36-wheel-cli/lib/python3.6/weakref.py", line 394, in __getitem__\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b'    return self.data[ref(key)]\n'
DEBUG  12-18 15:21:47  async_process_runner.py  b"KeyError: <weakref at 0x7f48b9d9cf48; to 'Node' at 0x7f48b83175c0>\x1b[0m\n"

# Fixed code
def stream_protocol_messages(
    self,
    protocol_identifier: Union[ProtocolAPI, Type[ProtocolAPI]],
) -> AsyncIterator[CommandAPI[Any]]:
    """
    Stream the messages for the specified protocol.
    """
    if isinstance(protocol_identifier, ProtocolAPI):
        protocol_class = type(protocol_identifier)
    elif isinstance(protocol_identifier, type) and issubclass(
        protocol_identifier, ProtocolAPI
    ):
        protocol_class = protocol_identifier
    else:
        raise TypeError("Unknown protocol identifier: {protocol}")

    if not self.has_protocol(protocol_class):
        raise UnknownProtocol(f"Unknown protocol '{protocol_class}'")

    if self._protocol_locks[protocol_class].locked():
        raise Exception(f"Streaming lock for {protocol_class} is not free.")
    elif not self._multiplex_lock.locked():
        raise Exception("Not multiplexed.")

    # Mostly a sanity check but this ensures we do better than accidentally
    # raising an attribute error in whatever race conditions or edge cases
    # potentially make the `_multiplex_token` unavailable.
    if not hasattr(self, "_multiplex_token"):
        raise Exception("No cancel token found for multiplexing.")

    # We do the wait_iter here so that the call sites in the handshakers
    # that use this don't need to be aware of cancellation tokens.
    return self.wait_iter(
        self._stream_protocol_messages(protocol_class),
        token=self._multiplex_token,
    )