2024-07-08 16:46:17 +09:00

90 lines
3.0 KiB
Python

import math
from collections import Counter
from typing import Awaitable
import pandas
from mixnet.connection import SimplexConnection
from mixnet.framework.framework import Framework, Queue
from mixnet.sim.config import NetworkConfig
from mixnet.sim.state import NodeState
class MeteredRemoteSimplexConnection(SimplexConnection):
framework: Framework
latency: float
send_queue: Queue
mid_queue: Queue
recv_queue: Queue
send_task: Awaitable
send_meters: list[int]
recv_task: Awaitable
recv_meters: list[int]
send_node_states: list[NodeState]
recv_node_states: list[NodeState]
msg_sizes: Counter[int]
def __init__(
self,
config: NetworkConfig,
framework: Framework,
send_node_states: list[NodeState],
recv_node_states: list[NodeState],
):
self.framework = framework
self.latency = config.random_latency()
self.send_queue = framework.queue()
self.mid_queue = framework.queue()
self.recv_queue = framework.queue()
self.send_meters = []
self.send_task = framework.spawn(self.__run_send_task())
self.recv_meters = []
self.recv_task = framework.spawn(self.__run_recv_task())
self.send_node_states = send_node_states
self.recv_node_states = recv_node_states
self.msg_sizes = Counter()
async def send(self, data: bytes) -> None:
await self.send_queue.put(data)
self.msg_sizes.update([len(data)])
ms = math.floor(self.framework.now() * 1000)
self.send_node_states[ms] = NodeState.SENDING
async def recv(self) -> bytes:
data = await self.recv_queue.get()
ms = math.floor(self.framework.now() * 1000)
self.send_node_states[ms] = NodeState.RECEIVING
return data
async def __run_send_task(self):
start_time = self.framework.now()
while True:
data = await self.send_queue.get()
self.__update_meter(self.send_meters, len(data), start_time)
await self.mid_queue.put(data)
async def __run_recv_task(self):
start_time = self.framework.now()
while True:
data = await self.mid_queue.get()
if data is None:
break
await self.framework.sleep(self.latency)
self.__update_meter(self.recv_meters, len(data), start_time)
await self.recv_queue.put(data)
def __update_meter(self, meters: list[int], size: int, start_time: float):
slot = math.floor(self.framework.now() - start_time)
assert slot >= len(meters) - 1
meters.extend([0] * (slot - len(meters) + 1))
meters[-1] += size
def sending_bandwidths(self) -> pandas.Series:
return self.__bandwidths(self.send_meters)
def receiving_bandwidths(self) -> pandas.Series:
return self.__bandwidths(self.recv_meters)
def __bandwidths(self, meters: list[int]) -> pandas.Series:
return pandas.Series(meters, name="bandwidth")