def test_add_edge(self): dag = DAG() dag.add_edges(sources=['b', 'c'], target='a') self.assertCountEqual(dag.heads(), ['a']) self.assertTrue(dag.has('a')) self.assertTrue(dag.has('b')) self.assertTrue(dag.has('c')) self.assertFalse(dag.has('d')) self.assertTrue(dag.has_path('b', 'a')) self.assertFalse(dag.has_path('a', 'b'))
class Node: prune_depth = 2 # Set configs as otherwise configparses takes too long to 'get' them again message_processing_time = int(config["simulation"]["message_processing_time"]) share_period = int(config["shares"]["period"]) randomise_share_period = config.getboolean("shares", "randomise") block_probability = float(config["shares"]["block_probability"]) shares_limit = int(config["shares"]["limit"]) run_time = int(config["simulation"]["run_time"]) def __init__(self, *, env, name): self.seq_no = 0 self.env = env self.name = name self.out_pipe = BroadcastPipe(env=env, sender=self.name) self.in_pipe = simpy.Store(self.env) self.neighbours = set() self.dag = DAG() # required to help prune dag for large simulations self.blocks_received = [] self.shares_sent = [] # Track num_shares_sent separately as shares_sent is pruned when block is received self.num_shares_sent = 0 self.shares_not_rewarded = {} self.num_blocks = 0 def heads(self): return self.dag.heads() def add_neighbour(self, neighbour, reversed=True): self.neighbours.add(neighbour) self.out_pipe.add_receiver(neighbour.in_pipe) if reversed and self not in neighbour.neighbours: neighbour.add_neighbour(self) def add_neighbours(self, neighbours: list, reversed=True): for neighbour in neighbours: self.add_neighbour(neighbour, reversed) def get_next_share_time(self): period = self.share_period if self.randomise_share_period: return get_random(period=period) else: return period def generate_shares(self): """Process to generate shares at random intervals.""" block_probability = self.block_probability while True: # wait for next share limit = self.shares_limit if limit != -1 and self.seq_no >= limit: yield self.env.timeout(self.run_time) else: yield self.env.timeout(self.get_next_share_time()) share = Share( source=self.name, heads=self.heads(), env=self.env, seq_no=self.seq_no, is_block=get_random(period=1) < block_probability, ) self.seq_no += 1 if share.is_block: self.num_blocks += 1 msg = ShareMessage(share=share) self.add_to_dag(msg.share.hash, msg.share.heads) self.send(msg) self.shares_sent.append(msg.share.hash) self.num_shares_sent += 1 self.handle_block_found(msg) def _log_send(self, msg, forward): _type = "s" if not forward else "f" if _type == "s": logging.info(f"{int(self.env.now)} {_type} n: {self.name} {msg}") else: logging.debug(f"{int(self.env.now)} {_type} n: {self.name} {msg}") def send(self, msg, forward=False): self._log_send(msg, forward) self.out_pipe.put(msg) def receive(self): """A process which consumes messages.""" while True: msg = yield self.in_pipe.get() logging.info(f"{int(self.env.now)} r n: {self.name} {msg}") self.env.process(self.handle_receive(msg.copy())) def forward(self, msg): msg.decrement_count() self.send(msg, forward=True) def add_to_dag(self, hash, heads): self.dag.add_edges(sources=heads, target=hash) def handle_receive(self, msg): yield self.env.timeout(self.message_processing_time) msg_in_dag = self.dag.has(msg.share.hash) if not msg_in_dag: if msg.should_forward(): self.forward(msg) self.add_to_dag(msg.share.hash, msg.share.heads) self.handle_received_block_found(msg) def handle_received_block_found(self, msg): if not msg.share.is_block: return if msg.share.source == self.name: return self.handle_block_found(msg) def handle_block_found(self, msg): # find delta between latest share this node already sent and # the last one referenced in the receieved share. not_rewarded = self.dag.find_not_reachable(self.shares_sent, msg.share.hash) if not_rewarded: self.shares_not_rewarded[msg.share.hash] = not_rewarded self.blocks_received.append(msg.share.hash) self._prune() def _prune(self): """Prune dag, blocks_received and shares_sent up to last two blocks received""" if len(self.blocks_received) <= self.prune_depth: return prune_upto = self.blocks_received[-self.prune_depth :][0] pruned_hashes = self.dag.prune_upto(prune_upto) del self.blocks_received[: -self.prune_depth] self._prune_shares_sent(pruned_hashes) def _prune_shares_sent(self, pruned_hashes): # Possible improvement - the implementation here goes over shares_sent twice. for hash in pruned_hashes: if hash in self.shares_sent: self.shares_sent.remove(hash) def start(self): logging.info(f"{self.name} starting...") self.env.process(self.receive()) self.env.process(self.generate_shares())