272 lines
8.3 KiB
Python
272 lines
8.3 KiB
Python
import argparse
|
|
import sys
|
|
import time
|
|
from random import randint
|
|
|
|
import multiaddr
|
|
import trio
|
|
from constants import *
|
|
from executor import Executor
|
|
from libp2p.peer.peerinfo import info_from_p2p_addr
|
|
from network import DANetwork
|
|
from subnet import calculate_subnets
|
|
|
|
"""
|
|
Entry point for the poc.
|
|
Handles cli arguments, initiates the network
|
|
and waits for it to complete.
|
|
|
|
Also does some simple completion check.
|
|
"""
|
|
|
|
|
|
async def run_network(params):
|
|
"""
|
|
Create the network.
|
|
Run the run_subnets
|
|
"""
|
|
|
|
num_nodes = int(params.nodes)
|
|
net = DANetwork(num_nodes)
|
|
shutdown = trio.Event()
|
|
disperse_send, disperse_recv = trio.open_memory_channel(0)
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(net.build, nursery, shutdown, disperse_send)
|
|
nursery.start_soon(
|
|
run_subnets, net, params, nursery, shutdown, disperse_send, disperse_recv
|
|
)
|
|
|
|
|
|
async def run_subnets(net, params, nursery, shutdown, disperse_send, disperse_recv):
|
|
"""
|
|
Run the actual PoC logic.
|
|
Calculate the subnets.
|
|
-> Establish connections based on the subnets <-
|
|
Runs the executor.
|
|
Runs simulated sampling.
|
|
Runs simple completion check
|
|
"""
|
|
|
|
num_nodes = int(params.nodes)
|
|
num_subnets = int(params.subnets)
|
|
data_size = int(params.data_size)
|
|
sample_threshold = int(params.sample_threshold)
|
|
fault_rate = int(params.fault_rate)
|
|
replication_factor = int(params.replication_factor)
|
|
|
|
while len(net.get_nodes()) != num_nodes:
|
|
print("nodes not ready yet")
|
|
await trio.sleep(0.1)
|
|
|
|
print("Nodes ready")
|
|
nodes = net.get_nodes()
|
|
subnets = calculate_subnets(nodes, num_subnets, replication_factor)
|
|
await print_subnet_info(subnets)
|
|
|
|
print("Establishing connections...")
|
|
node_list = {}
|
|
all_node_instances = set()
|
|
await establish_connections(subnets, node_list, all_node_instances, fault_rate)
|
|
|
|
print("Starting executor...")
|
|
exe = Executor.new(EXECUTOR_PORT, node_list, num_subnets, data_size)
|
|
|
|
print("Start dispersal and wait to complete...")
|
|
print("depending on network and subnet size this may take a while...")
|
|
global TIMESTAMP
|
|
TIMESTAMP = time.time()
|
|
async with trio.open_nursery() as subnursery:
|
|
subnursery.start_soon(wait_disperse_finished, disperse_recv, num_subnets)
|
|
subnursery.start_soon(exe.disperse, nursery)
|
|
subnursery.start_soon(disperse_watcher, disperse_send.clone())
|
|
|
|
print()
|
|
print()
|
|
|
|
print("OK. Start sampling...")
|
|
checked = []
|
|
for _ in range(sample_threshold):
|
|
nursery.start_soon(sample_node, exe, subnets, checked)
|
|
|
|
print("Waiting for sampling to finish...")
|
|
await check_complete(checked, sample_threshold)
|
|
|
|
print_connections(all_node_instances)
|
|
|
|
print("Test completed")
|
|
shutdown.set()
|
|
|
|
|
|
TIMESTAMP = time.time()
|
|
|
|
|
|
def print_connections(node_list):
|
|
for n in node_list:
|
|
for p in n.net_iface().get_peerstore().peer_ids():
|
|
if p == n.net_iface().get_id():
|
|
continue
|
|
print("node {} is connected to {}".format(n.get_id(), p))
|
|
print()
|
|
|
|
|
|
async def disperse_watcher(disperse_send):
|
|
while time.time() - TIMESTAMP < 5:
|
|
await trio.sleep(1)
|
|
|
|
await disperse_send.send(9999)
|
|
print("canceled")
|
|
|
|
|
|
async def wait_disperse_finished(disperse_recv, num_subnets):
|
|
# run until there are no changes detected
|
|
async for value in disperse_recv:
|
|
if value == 9999:
|
|
print("dispersal finished")
|
|
return
|
|
|
|
print(".", end="")
|
|
|
|
global TIMESTAMP
|
|
TIMESTAMP = time.time()
|
|
|
|
|
|
async def print_subnet_info(subnets):
|
|
"""
|
|
Print which node is in what subnet
|
|
"""
|
|
|
|
print()
|
|
print("By subnets: ")
|
|
for subnet in subnets:
|
|
print("subnet: {} - ".format(subnet), end="")
|
|
for n in subnets[subnet]:
|
|
print(n.get_id().pretty()[:16], end=", ")
|
|
print()
|
|
|
|
print()
|
|
print()
|
|
|
|
|
|
async def establish_connections(subnets, node_list, all_node_instances, fault_rate=0):
|
|
"""
|
|
Each node in a subnet connects to the other ones in that subnet.
|
|
"""
|
|
for subnet in subnets:
|
|
# n is a DANode
|
|
for n in subnets[subnet]:
|
|
# while nodes connect to each other, they are **mutually** added
|
|
# to their peer lists. Hence, we don't need to establish connections
|
|
# again to peers we are already connected.
|
|
# So in each iteration we get the peer list for the current node
|
|
# to later check if we are already connected with the next peer
|
|
this_nodes_peers = n.net_iface().get_peerstore().peer_ids()
|
|
all_node_instances.add(n)
|
|
faults = []
|
|
for i in range(fault_rate):
|
|
faults.append(randint(0, len(subnets[subnet])))
|
|
for i, nn in enumerate(subnets[subnet]):
|
|
# don't connect to self
|
|
if nn.get_id() == n.get_id():
|
|
continue
|
|
if i in faults:
|
|
continue
|
|
remote_id = nn.get_id().pretty()
|
|
remote_port = nn.get_port()
|
|
# this script only works on localhost!
|
|
addr = "/ip4/127.0.0.1/tcp/{}/p2p/{}/".format(remote_port, remote_id)
|
|
remote_addr = multiaddr.Multiaddr(addr)
|
|
remote = info_from_p2p_addr(remote_addr)
|
|
if subnet not in node_list:
|
|
node_list[subnet] = []
|
|
node_list[subnet].append(remote)
|
|
# check if we are already connected with this peer. If yes, skip connecting
|
|
if nn.get_id() in this_nodes_peers:
|
|
continue
|
|
if DEBUG:
|
|
print("{} connecting to {}...".format(n.get_id(), addr))
|
|
await n.net_iface().connect(remote)
|
|
|
|
print()
|
|
|
|
|
|
async def check_complete(checked, sample_threshold):
|
|
"""
|
|
Simple completion check:
|
|
Check how many nodes have already been "sampled"
|
|
"""
|
|
|
|
while len(checked) < sample_threshold:
|
|
await trio.sleep(0.5)
|
|
print("check_complete exiting")
|
|
return
|
|
|
|
|
|
async def sample_node(exe, subnets, checked):
|
|
"""
|
|
Pick a random subnet.
|
|
Pick a random node in that subnet.
|
|
As the executor has a list of hashes per subnet,
|
|
we can ask that node if it has that hash.
|
|
"""
|
|
|
|
# s: subnet
|
|
s = randint(0, len(subnets) - 1)
|
|
# n: node (index)
|
|
n = randint(0, len(subnets[s]) - 1)
|
|
# actual node
|
|
node = subnets[s][n]
|
|
# pick the hash to check
|
|
hashstr = exe.get_hash(s)
|
|
# run the "sampling"
|
|
has = await node.has_hash(hashstr)
|
|
if has:
|
|
print("node {} has hash {}".format(node.get_id().pretty(), hashstr))
|
|
else:
|
|
print("node {} does NOT HAVE hash {}".format(node.get_id().pretty(), hashstr))
|
|
print("TEST FAILED")
|
|
# signal we "sampled" another node
|
|
checked.append(1)
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-s", "--subnets", help="Number of subnets [default: 256]")
|
|
parser.add_argument("-n", "--nodes", help="Number of nodes [default: 32]")
|
|
parser.add_argument(
|
|
"-t",
|
|
"--sample-threshold",
|
|
help="Threshold for sampling request attempts [default: 12]",
|
|
)
|
|
parser.add_argument("-d", "--data-size", help="Size of packages [default: 1024]")
|
|
parser.add_argument("-f", "--fault_rate", help="Fault rate [default: 0]")
|
|
parser.add_argument(
|
|
"-r", "--replication_factor", help="Replication factor [default: 4]"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if not args.subnets:
|
|
args.subnets = DEFAULT_SUBNETS
|
|
if not args.nodes:
|
|
args.nodes = DEFAULT_NODES
|
|
if not args.sample_threshold:
|
|
args.sample_threshold = DEFAULT_SAMPLE_THRESHOLD
|
|
if not args.data_size:
|
|
args.data_size = DEFAULT_DATA_SIZE
|
|
if not args.replication_factor:
|
|
args.replication_factor = DEFAULT_REPLICATION_FACTOR
|
|
if not args.fault_rate:
|
|
args.fault_rate = 0
|
|
|
|
print("Number of subnets will be: {}".format(args.subnets))
|
|
print("Number of nodes will be: {}".format(args.nodes))
|
|
print("Size of data package will be: {}".format(args.data_size))
|
|
print("Threshold for sampling attempts will be: {}".format(args.sample_threshold))
|
|
print("Fault rate will be: {}".format(args.fault_rate))
|
|
|
|
print()
|
|
print("*******************")
|
|
print("Starting network...")
|
|
|
|
trio.run(run_network, args)
|