research/sharding_fork_choice_poc/test.py

57 lines
2.3 KiB
Python

from networksim import NetworkSimulator
from beacon_chain_node import Node, NOTARIES, SHARDS, BeaconBlock, MainChainBlock, ShardCollation, main_genesis, beacon_genesis
net = NetworkSimulator(latency=19)
notaries = [Node(i, net, sleepy=i % 5 == 9) for i in range(NOTARIES)]
net.agents = notaries
net.generate_peers()
for i in range(2000):
net.tick()
for n in notaries:
print("Beacon head: %d" % n.blocks[n.beacon_chain[-1]].number)
print("Main chain head: %d" % n.blocks[n.main_chain[-1]].number)
print("Shard heads: %r" % [n.blocks[x[-1]].number for x in n.shard_chains])
print("Total beacon blocks received: %d" % (len([b for b in n.blocks.values() if isinstance(b, BeaconBlock)]) - 1))
print("Total beacon blocks received and signed: %d" % (len([b for b in n.blocks.keys() if b in n.sigs and len(n.sigs[b]) >= n.blocks[b].notary_req]) - 1))
print("Total main chain blocks received: %d" % (len([b for b in n.blocks.values() if isinstance(b, MainChainBlock)]) - 1))
print("Total shard blocks received: %r" % [len([b for b in n.blocks.values() if isinstance(b, ShardCollation) and b.shard_id == i]) - 1 for i in range(SHARDS)])
import matplotlib.pyplot as plt
import networkx as nx
import random
G=nx.Graph()
#positions = {main_genesis.hash: 0, beacon_genesis.hash: 0}
#queue = [
for b in n.blocks.values():
if b.number > 0:
if isinstance(b, BeaconBlock):
G.add_edge(b.hash, b.main_chain_ref, color='g')
G.add_edge(b.hash, b.parent_hash, color='y')
elif isinstance(b, MainChainBlock):
G.add_edge(b.hash, b.parent_hash, color='b')
elif isinstance(b, ShardCollation):
G.add_edge(b.hash, b.beacon_ref, color='g')
G.add_edge(b.hash, b.parent_hash, color='r')
def mkoffset(b):
return random.randrange(5) + \
(5 if isinstance(b, MainChainBlock) else
0 if isinstance(b, BeaconBlock) else
-5 - 5 * b.shard_id if isinstance(b, ShardCollation) else
None)
pos={b.hash: (b.ts + mkoffset(b), b.ts) for b in n.blocks.values()}
edges = G.edges()
colors = [G[u][v]['color'] for u,v in edges]
nx.draw_networkx_nodes(G,pos,node_size=10,node_shape='o',node_color='0.75')
nx.draw_networkx_edges(G,pos,
width=2,edge_color=colors)
plt.axis('off')
# plt.savefig("degree.png", bbox_inches="tight")
plt.show()