mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
Stream rearchitecture (#126)
* Add generic protocol handler * Add generic protocol handler to stream muxing pipeline * Modify conn_handler to only deal with connections * mplex accept stream architecture changes * Add create generic protocol handler * Fix minor bugs * who would win 4 devs or one not * Debugging * rearch with handle_incoming infinite loop, seems to work, needs cleanup" * passing linting, still needs cleanup * fixing linting again; code still needs cleanup * fixing tests; code still needs cleanup * adding test cleanup and task cleanup, removing prints * linting, and cleanup complete * storing connections based on peer id * remove dead code * remove unnecessary peer_id
This commit is contained in:
Robert Zajac
@ -1,5 +1,7 @@
|
||||
import pytest
|
||||
import asyncio
|
||||
|
||||
from tests.utils import cleanup
|
||||
from libp2p import new_node
|
||||
from libp2p.peer.peerinfo import info_from_p2p_addr
|
||||
from libp2p.protocol_muxer.multiselect_client import MultiselectClientError
|
||||
@ -27,25 +29,28 @@ async def hello_world(host_a, host_b):
|
||||
|
||||
|
||||
async def connect_write(host_a, host_b):
|
||||
messages = [b'data %d' % i for i in range(5)]
|
||||
messages = ['data %d' % i for i in range(5)]
|
||||
received = []
|
||||
|
||||
async def stream_handler(stream):
|
||||
received = []
|
||||
while True:
|
||||
try:
|
||||
received.append((await stream.read()).decode())
|
||||
except Exception: # exception is raised when other side close the stream ?
|
||||
break
|
||||
await stream.close()
|
||||
assert received == messages
|
||||
host_a.set_stream_handler(PROTOCOL_ID, stream_handler)
|
||||
|
||||
# Start a stream with the destination.
|
||||
# Multiaddress of the destination peer is fetched from the peerstore using 'peerId'.
|
||||
stream = await host_b.new_stream(host_a.get_id(), [PROTOCOL_ID])
|
||||
for message in messages:
|
||||
await stream.write(message)
|
||||
await stream.write(message.encode())
|
||||
|
||||
# Reader needs time due to async reads
|
||||
await asyncio.sleep(2)
|
||||
|
||||
await stream.close()
|
||||
assert received == messages
|
||||
|
||||
|
||||
async def connect_read(host_a, host_b):
|
||||
@ -103,3 +108,5 @@ async def test_chat(test):
|
||||
await host_b.connect(info)
|
||||
|
||||
await test(host_a, host_b)
|
||||
|
||||
await cleanup()
|
||||
|
||||
@ -1,9 +1,12 @@
|
||||
import multiaddr
|
||||
import pytest
|
||||
|
||||
from tests.utils import cleanup
|
||||
from libp2p import new_node
|
||||
from libp2p.peer.peerinfo import info_from_p2p_addr
|
||||
|
||||
# pylint: disable=too-many-locals
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_messages():
|
||||
@ -13,10 +16,8 @@ async def test_simple_messages():
|
||||
async def stream_handler(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
print("host B received:" + read_string)
|
||||
|
||||
response = "ack:" + read_string
|
||||
print("sending response:" + response)
|
||||
await stream.write(response.encode())
|
||||
|
||||
node_b.set_stream_handler("/echo/1.0.0", stream_handler)
|
||||
@ -32,11 +33,10 @@ async def test_simple_messages():
|
||||
|
||||
response = (await stream.read()).decode()
|
||||
|
||||
print("res: " + response)
|
||||
assert response == ("ack:" + message)
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
return
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -47,38 +47,31 @@ async def test_double_response():
|
||||
async def stream_handler(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
print("host B received:" + read_string)
|
||||
|
||||
response = "ack1:" + read_string
|
||||
print("sending response:" + response)
|
||||
await stream.write(response.encode())
|
||||
|
||||
response = "ack2:" + read_string
|
||||
print("sending response:" + response)
|
||||
await stream.write(response.encode())
|
||||
|
||||
node_b.set_stream_handler("/echo/1.0.0", stream_handler)
|
||||
|
||||
# Associate the peer with local ip address (see default parameters of Libp2p())
|
||||
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
|
||||
print("node_a about to open stream")
|
||||
stream = await node_a.new_stream(node_b.get_id(), ["/echo/1.0.0"])
|
||||
|
||||
messages = ["hello" + str(x) for x in range(10)]
|
||||
for message in messages:
|
||||
await stream.write(message.encode())
|
||||
|
||||
response1 = (await stream.read()).decode()
|
||||
|
||||
print("res1: " + response1)
|
||||
assert response1 == ("ack1:" + message)
|
||||
|
||||
response2 = (await stream.read()).decode()
|
||||
|
||||
print("res2: " + response2)
|
||||
assert response2 == ("ack2:" + message)
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
return
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -127,7 +120,197 @@ async def test_multiple_streams():
|
||||
assert response_a == ("ack_b:" + a_message) and response_b == ("ack_a:" + b_message)
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
return
|
||||
await cleanup()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_streams_same_initiator_different_protocols():
|
||||
node_a = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
node_b = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
|
||||
async def stream_handler_a1(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_a1:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
async def stream_handler_a2(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_a2:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
async def stream_handler_a3(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_a3:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
node_b.set_stream_handler("/echo_a1/1.0.0", stream_handler_a1)
|
||||
node_b.set_stream_handler("/echo_a2/1.0.0", stream_handler_a2)
|
||||
node_b.set_stream_handler("/echo_a3/1.0.0", stream_handler_a3)
|
||||
|
||||
# Associate the peer with local ip address (see default parameters of Libp2p())
|
||||
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
|
||||
node_b.get_peerstore().add_addrs(node_a.get_id(), node_a.get_addrs(), 10)
|
||||
|
||||
# Open streams to node_b over echo_a1 echo_a2 echo_a3 protocols
|
||||
stream_a1 = await node_a.new_stream(node_b.get_id(), ["/echo_a1/1.0.0"])
|
||||
stream_a2 = await node_a.new_stream(node_b.get_id(), ["/echo_a2/1.0.0"])
|
||||
stream_a3 = await node_a.new_stream(node_b.get_id(), ["/echo_a3/1.0.0"])
|
||||
|
||||
messages = ["hello" + str(x) for x in range(10)]
|
||||
for message in messages:
|
||||
a1_message = message + "_a1"
|
||||
a2_message = message + "_a2"
|
||||
a3_message = message + "_a3"
|
||||
|
||||
await stream_a1.write(a1_message.encode())
|
||||
await stream_a2.write(a2_message.encode())
|
||||
await stream_a3.write(a3_message.encode())
|
||||
|
||||
response_a1 = (await stream_a1.read()).decode()
|
||||
response_a2 = (await stream_a2.read()).decode()
|
||||
response_a3 = (await stream_a3.read()).decode()
|
||||
|
||||
assert (response_a1 == ("ack_a1:" + a1_message)
|
||||
and response_a2 == ("ack_a2:" + a2_message)
|
||||
and response_a3 == ("ack_a3:" + a3_message))
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
await cleanup()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_streams_two_initiators():
|
||||
node_a = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
node_b = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
|
||||
async def stream_handler_a1(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_a1:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
async def stream_handler_a2(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_a2:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
async def stream_handler_b1(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_b1:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
async def stream_handler_b2(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack_b2:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
node_a.set_stream_handler("/echo_b1/1.0.0", stream_handler_b1)
|
||||
node_a.set_stream_handler("/echo_b2/1.0.0", stream_handler_b2)
|
||||
|
||||
node_b.set_stream_handler("/echo_a1/1.0.0", stream_handler_a1)
|
||||
node_b.set_stream_handler("/echo_a2/1.0.0", stream_handler_a2)
|
||||
|
||||
# Associate the peer with local ip address (see default parameters of Libp2p())
|
||||
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
|
||||
node_b.get_peerstore().add_addrs(node_a.get_id(), node_a.get_addrs(), 10)
|
||||
|
||||
stream_a1 = await node_a.new_stream(node_b.get_id(), ["/echo_a1/1.0.0"])
|
||||
stream_a2 = await node_a.new_stream(node_b.get_id(), ["/echo_a2/1.0.0"])
|
||||
|
||||
stream_b1 = await node_b.new_stream(node_a.get_id(), ["/echo_b1/1.0.0"])
|
||||
stream_b2 = await node_b.new_stream(node_a.get_id(), ["/echo_b2/1.0.0"])
|
||||
|
||||
# A writes to /echo_b via stream_a, and B writes to /echo_a via stream_b
|
||||
messages = ["hello" + str(x) for x in range(10)]
|
||||
for message in messages:
|
||||
a1_message = message + "_a1"
|
||||
a2_message = message + "_a2"
|
||||
|
||||
b1_message = message + "_b1"
|
||||
b2_message = message + "_b2"
|
||||
|
||||
await stream_a1.write(a1_message.encode())
|
||||
await stream_a2.write(a2_message.encode())
|
||||
|
||||
await stream_b1.write(b1_message.encode())
|
||||
await stream_b2.write(b2_message.encode())
|
||||
|
||||
response_a1 = (await stream_a1.read()).decode()
|
||||
response_a2 = (await stream_a2.read()).decode()
|
||||
|
||||
response_b1 = (await stream_b1.read()).decode()
|
||||
response_b2 = (await stream_b2.read()).decode()
|
||||
|
||||
assert (response_a1 == ("ack_a1:" + a1_message)
|
||||
and response_a2 == ("ack_a2:" + a2_message)
|
||||
and response_b1 == ("ack_b1:" + b1_message)
|
||||
and response_b2 == ("ack_b2:" + b2_message))
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
await cleanup()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_triangle_nodes_connection():
|
||||
node_a = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
node_b = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
node_c = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
|
||||
|
||||
async def stream_handler(stream):
|
||||
while True:
|
||||
read_string = (await stream.read()).decode()
|
||||
|
||||
response = "ack:" + read_string
|
||||
await stream.write(response.encode())
|
||||
|
||||
node_a.set_stream_handler("/echo/1.0.0", stream_handler)
|
||||
node_b.set_stream_handler("/echo/1.0.0", stream_handler)
|
||||
node_c.set_stream_handler("/echo/1.0.0", stream_handler)
|
||||
|
||||
# Associate the peer with local ip address (see default parameters of Libp2p())
|
||||
# Associate all permutations
|
||||
node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
|
||||
node_a.get_peerstore().add_addrs(node_c.get_id(), node_c.get_addrs(), 10)
|
||||
|
||||
node_b.get_peerstore().add_addrs(node_a.get_id(), node_a.get_addrs(), 10)
|
||||
node_b.get_peerstore().add_addrs(node_c.get_id(), node_c.get_addrs(), 10)
|
||||
|
||||
node_c.get_peerstore().add_addrs(node_a.get_id(), node_a.get_addrs(), 10)
|
||||
node_c.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
|
||||
|
||||
stream_a_to_b = await node_a.new_stream(node_b.get_id(), ["/echo/1.0.0"])
|
||||
stream_a_to_c = await node_a.new_stream(node_c.get_id(), ["/echo/1.0.0"])
|
||||
|
||||
stream_b_to_a = await node_b.new_stream(node_a.get_id(), ["/echo/1.0.0"])
|
||||
stream_b_to_c = await node_b.new_stream(node_c.get_id(), ["/echo/1.0.0"])
|
||||
|
||||
stream_c_to_a = await node_c.new_stream(node_a.get_id(), ["/echo/1.0.0"])
|
||||
stream_c_to_b = await node_c.new_stream(node_b.get_id(), ["/echo/1.0.0"])
|
||||
|
||||
messages = ["hello" + str(x) for x in range(5)]
|
||||
streams = [stream_a_to_b, stream_a_to_c, stream_b_to_a, stream_b_to_c,
|
||||
stream_c_to_a, stream_c_to_b]
|
||||
|
||||
for message in messages:
|
||||
for stream in streams:
|
||||
await stream.write(message.encode())
|
||||
|
||||
response = (await stream.read()).decode()
|
||||
|
||||
assert response == ("ack:" + message)
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -152,3 +335,6 @@ async def test_host_connect():
|
||||
ma_node_b = multiaddr.Multiaddr('/p2p/%s' % node_b.get_id().pretty())
|
||||
for addr in node_a.get_peerstore().addrs(node_b.get_id()):
|
||||
assert addr.encapsulate(ma_node_b) in node_b.get_addrs()
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
await cleanup()
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
import pytest
|
||||
|
||||
from tests.utils import cleanup
|
||||
from libp2p import new_node
|
||||
from libp2p.protocol_muxer.multiselect_client import MultiselectClientError
|
||||
|
||||
@ -49,7 +50,7 @@ async def perform_simple_test(expected_selected_protocol,
|
||||
assert expected_selected_protocol == stream.get_protocol()
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
return
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -65,6 +66,9 @@ async def test_single_protocol_fails():
|
||||
await perform_simple_test("", ["/echo/1.0.0"],
|
||||
["/potato/1.0.0"])
|
||||
|
||||
# Cleanup not reached on error
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_protocol_first_is_valid_succeeds():
|
||||
@ -91,3 +95,6 @@ async def test_multiple_protocol_fails():
|
||||
with pytest.raises(MultiselectClientError):
|
||||
await perform_simple_test("", protocols_for_client,
|
||||
protocols_for_listener)
|
||||
|
||||
# Cleanup not reached on error
|
||||
await cleanup()
|
||||
|
||||
14
tests/utils.py
Normal file
14
tests/utils.py
Normal file
@ -0,0 +1,14 @@
|
||||
import asyncio
|
||||
|
||||
from contextlib import suppress
|
||||
|
||||
|
||||
async def cleanup():
|
||||
pending = asyncio.all_tasks()
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
|
||||
# Now we should await task to execute it's cancellation.
|
||||
# Cancelled task raises asyncio.CancelledError that we can suppress:
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
Reference in New Issue
Block a user