2024-07-15 18:30:27 +09:00

116 lines
4.3 KiB
Python

import usim
from matplotlib import pyplot
import mixnet.framework.usim as usimfw
from mixnet.config import GlobalConfig, MixMembership, NodeInfo
from mixnet.framework import Framework
from mixnet.node import Node, PeeringDegreeReached
from mixnet.sim.config import Config
from mixnet.sim.connection import MeteredRemoteSimplexConnection
from mixnet.sim.state import NodeState, NodeStateTable
from mixnet.sim.stats import ConnectionStats
class Simulation:
"""
Manages the entire cycle of simulation: initialization, running, and analysis.
"""
def __init__(self, config: Config):
self.config = config
async def run(self):
# Run the simulation
conn_stats, node_state_table = await self.__run()
# Analyze the simulation results
conn_stats.analyze()
node_state_table.analyze()
# Show plots
if self.config.simulation.show_plots:
pyplot.show()
async def __run(self) -> tuple[ConnectionStats, NodeStateTable]:
# Initialize analysis tools
node_state_table = NodeStateTable(
self.config.network.num_nodes, self.config.simulation.duration_sec
)
conn_stats = ConnectionStats()
# Create a μSim scope and run the simulation
async with usim.until(usim.time + self.config.simulation.duration_sec) as scope:
self.framework = usimfw.Framework(scope)
nodes, conn_stats, node_state_table = self.__init_nodes(
node_state_table, conn_stats
)
for node in nodes:
self.framework.spawn(self.__run_node_logic(node))
# Return analysis tools once the μSim scope is done
return conn_stats, node_state_table
def __init_nodes(
self, node_state_table: NodeStateTable, conn_stats: ConnectionStats
) -> tuple[list[Node], ConnectionStats, NodeStateTable]:
# Initialize node/global configurations
node_configs = self.config.node_configs()
global_config = GlobalConfig(
MixMembership(
[
NodeInfo(node_config.private_key.public_key())
for node_config in node_configs
],
self.config.mix.mix_path.seed,
),
self.config.mix.transmission_rate_per_sec,
self.config.mix.max_message_size,
self.config.mix.mix_path.max_length,
)
# Initialize Node instances
nodes = [
Node(self.framework, node_config, global_config)
for node_config in node_configs
]
# Connect nodes to each other
for i, node in enumerate(nodes):
# For now, we only consider a simple ring topology for simplicity.
peer_idx = (i + 1) % len(nodes)
peer = nodes[peer_idx]
node_states = node_state_table[i]
peer_states = node_state_table[peer_idx]
# Create simplex inbound/outbound connections
# and use them to connect node and peer.
inbound_conn, outbound_conn = (
self.__create_conn(peer_states, node_states),
self.__create_conn(node_states, peer_states),
)
node.connect(peer, inbound_conn, outbound_conn)
# Register the connections to the connection statistics
conn_stats.register(node, inbound_conn, outbound_conn)
conn_stats.register(peer, outbound_conn, inbound_conn)
return nodes, conn_stats, node_state_table
def __create_conn(
self, sender_states: list[NodeState], receiver_states: list[NodeState]
) -> MeteredRemoteSimplexConnection:
return MeteredRemoteSimplexConnection(
self.config.network,
self.framework,
sender_states,
receiver_states,
)
async def __run_node_logic(self, node: Node):
"""
Runs the lottery periodically to check if the node is selected to send a block.
If the node is selected, creates a block and sends it through mix nodes.
"""
lottery_config = self.config.logic.sender_lottery
while True:
await (usim.time + lottery_config.interval_sec)
if lottery_config.seed.random() < lottery_config.probability:
await node.send_message(b"selected block")