1  from distributions import transform, normal_distribution
 2  import random
 3  
 4  
 5  class NetworkSimulator():
 6  
 7      def __init__(self, latency=50):
 8          self.agents = []
 9          self.latency_distribution_sample = transform(normal_distribution(latency, (latency * 2) // 5), lambda x: max(x, 0))
10          self.time = 0
11          self.objqueue = {}
12          self.peers = {}
13          self.reliability = 0.9
14  
15      def generate_peers(self, num_peers=5):
16          self.peers = {}
17          for a in self.agents:
18              p = []
19              while len(p) <= num_peers // 2:
20                  p.append(random.choice(self.agents))
21                  if p[-1] == a:
22                      p.pop()
23              self.peers[a.id] = self.peers.get(a.id, []) + p
24              for peer in p:
25                  self.peers[peer.id] = self.peers.get(peer.id, []) + [a]
26  
27      def tick(self):
28          if self.time in self.objqueue:
29              for recipient, obj in self.objqueue[self.time]:
30                  if random.random() < self.reliability:
31                      recipient.on_receive(obj)
32              del self.objqueue[self.time]
33          for a in self.agents:
34              a.tick()
35          self.time += 1
36  
37      def run(self, steps):
38          for i in range(steps):
39              self.tick()
40  
41      def broadcast(self, sender, obj):
42          for p in self.peers[sender.id]:
43              recv_time = self.time + self.latency_distribution_sample()
44              if recv_time not in self.objqueue:
45                  self.objqueue[recv_time] = []
46              self.objqueue[recv_time].append((p, obj))
47  
48      def direct_send(self, to_id, obj):
49          for a in self.agents:
50              if a.id == to_id:
51                  recv_time = self.time + self.latency_distribution_sample()
52                  if recv_time not in self.objqueue:
53                      self.objqueue[recv_time] = []
54                  self.objqueue[recv_time].append((a, obj))
55  
56      def knock_offline_random(self, n):
57          ko = {}
58          while len(ko) < n:
59              c = random.choice(self.agents)
60              ko[c.id] = c
61          for c in ko.values():
62              self.peers[c.id] = []
63          for a in self.agents:
64              self.peers[a.id] = [x for x in self.peers[a.id] if x.id not in ko]
65  
66      def partition(self):
67          a = {}
68          while len(a) < len(self.agents) / 2:
69              c = random.choice(self.agents)
70              a[c.id] = c
71          for c in self.agents:
72              if c.id in a:
73                  self.peers[c.id] = [x for x in self.peers[c.id] if x.id in a]
74              else:
75                  self.peers[c.id] = [x for x in self.peers[c.id] if x.id not in a]