nomos-specs/mixnet/test_node.py

125 lines
4.6 KiB
Python
Raw Normal View History

2024-01-25 09:04:55 +00:00
import asyncio
2024-01-23 01:29:14 +00:00
from datetime import datetime
from typing import Tuple
2024-01-25 09:04:55 +00:00
from unittest import IsolatedAsyncioTestCase
2024-01-23 01:29:14 +00:00
import numpy
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from pysphinx.sphinx import SphinxPacket
from mixnet.bls import generate_bls
from mixnet.mixnet import Mixnet, MixnetTopology
from mixnet.node import MixNode, NodeAddress, PacketPayloadQueue, PacketQueue
from mixnet.packet import PacketBuilder
from mixnet.poisson import poisson_interval_sec, poisson_mean_interval_sec
2024-01-25 09:04:55 +00:00
from mixnet.test_utils import with_test_timeout
2024-01-23 01:29:14 +00:00
from mixnet.utils import random_bytes
2024-01-25 09:04:55 +00:00
class TestMixNodeRunner(IsolatedAsyncioTestCase):
@with_test_timeout(180)
async def test_mixnode_runner_emission_rate(self):
2024-01-23 01:29:14 +00:00
"""
Test if MixNodeRunner works as a M/M/inf queue.
If inputs are arrived at Poisson rate `lambda`,
and if processing is delayed according to an exponential distribution with a rate `mu`,
the rate of outputs should be `lambda`.
"""
mixnet, topology = self.init()
2024-01-25 09:04:55 +00:00
inbound_socket: PacketQueue = asyncio.Queue()
outbound_socket: PacketPayloadQueue = asyncio.Queue()
2024-01-23 01:29:14 +00:00
packet, route = PacketBuilder.real(b"msg", mixnet, topology).next()
delay_rate_per_min = 30 # mu (= 2s delay on average)
# Start only the first mix node for testing
2024-01-25 09:04:55 +00:00
_ = route[0].start(delay_rate_per_min, inbound_socket, outbound_socket)
2024-01-23 01:29:14 +00:00
# Send packets to the first mix node in a Poisson distribution
packet_count = 100
emission_rate_per_min = 120 # lambda (= 2msg/sec)
2024-01-25 09:04:55 +00:00
# This queue is just for counting how many packets have been sent so far.
sent_packet_queue: PacketQueue = asyncio.Queue()
_ = asyncio.create_task(
self.send_packets(
2024-01-23 01:29:14 +00:00
inbound_socket,
packet,
route[0].addr,
packet_count,
emission_rate_per_min,
2024-01-25 09:04:55 +00:00
sent_packet_queue,
)
2024-01-23 01:29:14 +00:00
)
# Calculate intervals between outputs and gather num_jobs in the first mix node.
intervals = []
num_jobs = []
ts = datetime.now()
for _ in range(packet_count):
2024-01-25 09:04:55 +00:00
_ = await outbound_socket.get()
2024-01-23 01:29:14 +00:00
now = datetime.now()
intervals.append((now - ts).total_seconds())
2024-01-25 09:04:55 +00:00
# Calculate the current # of jobs staying in the mix node
num_packets_emitted_from_mixnode = len(intervals)
num_packets_sent_to_mixnode = sent_packet_queue.qsize()
num_jobs.append(
num_packets_sent_to_mixnode - num_packets_emitted_from_mixnode
)
2024-01-23 01:29:14 +00:00
ts = now
2024-01-25 09:04:55 +00:00
2024-01-23 01:29:14 +00:00
# Remove the first interval that would be much larger than other intervals,
# because of the delay in mix node.
intervals = intervals[1:]
num_jobs = num_jobs[1:]
# Check if the emission rate of the first mix node is the same as
# the emission rate of the message sender, but with a delay.
# If outputs follow the Poisson distribution with a rate `lambda`,
# a mean interval between outputs must be `1/lambda`.
self.assertAlmostEqual(
float(numpy.mean(intervals)),
poisson_mean_interval_sec(emission_rate_per_min),
delta=1.0,
)
# If runner is a M/M/inf queue,
# a mean number of jobs being processed/scheduled in the runner must be `lambda/mu`.
self.assertAlmostEqual(
float(numpy.mean(num_jobs)),
round(emission_rate_per_min / delay_rate_per_min),
2024-01-23 01:46:00 +00:00
delta=1.5,
2024-01-23 01:29:14 +00:00
)
@staticmethod
2024-01-25 09:04:55 +00:00
async def send_packets(
2024-01-23 01:29:14 +00:00
inbound_socket: PacketQueue,
packet: SphinxPacket,
node_addr: NodeAddress,
cnt: int,
rate_per_min: int,
2024-01-25 09:04:55 +00:00
# For testing purpose, to inform the caller how many packets have been sent to the inbound_socket
sent_packet_queue: PacketQueue,
2024-01-23 01:29:14 +00:00
):
for _ in range(cnt):
2024-01-25 09:04:55 +00:00
# Since the task is not heavy, just sleep for seconds instead of using emission_notifier
await asyncio.sleep(poisson_interval_sec(rate_per_min))
await inbound_socket.put((node_addr, packet))
await sent_packet_queue.put((node_addr, packet))
2024-01-23 01:29:14 +00:00
@staticmethod
def init() -> Tuple[Mixnet, MixnetTopology]:
mixnet = Mixnet(
[
MixNode(
generate_bls(),
X25519PrivateKey.generate(),
random_bytes(32),
)
for _ in range(12)
]
)
topology = mixnet.build_topology(b"entropy", 3, 3)
return mixnet, topology