Skip to content
Snippets Groups Projects
network.py 8.67 KiB
Newer Older
  • Learn to ignore specific revisions
  • Jarrod Pas's avatar
    Jarrod Pas committed
    from argparse import ArgumentParser
    from collections import defaultdict, OrderedDict
    from itertools import combinations
    import math
    import random
    import sys
    
    import simpy
    import yaml
    import networkx as nx
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    from .core import TickProcess
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    from .communities import types as communities
    from .routers import types as routers
    from .traces import types as traces
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
    class Network:
        ''''''
        def __init__(self, env,
                     packets=None,
                     node_factory=None,
                     community=None,
                     trace=None):
            ''''''
            self.env = env
    
            # contact trace
            if trace is None:
                trace = traces['random']()
            self.trace = trace
            self.trace.start(env, self)
    
            # community detection
            self.community = community
            if community is not None:
                self.community.start(env, self)
    
            # packet generation
            if packets is None:
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                packets = {}
            self.packets = PacketGenerator(**packets)
            self.packets.start(env, self)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
            # create node network
            if node_factory is None:
                node_factory = NodeFactory(tick_rate=1, router=routers['direct'])
            self.nodes = [
                node_factory(env, self, nid)
                for nid in range(self.trace.nodes)
            ]
            self.links = [
                (a, b)
                for a, b in combinations(self.nodes, 2)
            ]
    
            # set up networkx graph
            self.graph = nx.Graph()
            self.graph.add_nodes_from(self.nodes)
            self.graph.add_edges_from([
                (a, b, { 'state': False })
                for a, b in self.links
            ])
    
        def set_link(self, a, b, state):
            if isinstance(a, int):
                a = self.nodes[a]
    
            if isinstance(b, int):
                b = self.nodes[b]
    
            edge = self[a][b]
            if edge['state'] == state:
                return
    
            if state is None:
                state = not edge['state']
            edge['state'] = state
    
            if self.community:
                self.community.set_link(a, b, state, self.env.now)
    
        def toggle_link(self, a, b):
            self.set_link(a, b, None)
    
        def send_link(self, a, b, packet):
            ''''''
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            if self.graph[a][b]['state']:
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                # TODO: transfer delay
                b.recv(packet)
            else:
                raise Exception('Nodes {} and {} not connected'.format(a, b))
    
        def __getitem__(self, node):
            ''''''
            return self.graph[node]
    
    
    def NodeFactory(router, **kwargs):
        nid = -1
        def factory(env, network, nid):
            nid += 1
            return Node(env, network, nid, router=router, **kwargs)
        return factory
    
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    class Node(TickProcess):
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        ''''''
        def __init__(self, env, network, nid,
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                     buffer_size=None, tick_time=1, router=None):
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            ''''''
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            super().__init__(tick_time)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.env = env
    
            self.network = network
            self.id = nid
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.buffer = Buffer(self.env, capacity=buffer_size)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
            # bind router as a class method
            if router is None:
                router = routers['direct']
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.router = router #router.__get__(self, Node)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.router_state = {}
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.start(env)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        def route_packets(self):
            packets_to_delete = []
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            for packet in self.buffer:
                if self.router(self, packet, self.router_state):
                    packets_to_delete.append(packet)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            for packet in packets_to_delete:
                self.buffer.remove(packet)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        def process(self):
            ''''''
            while True:
                self.buffer.clean()
                self.route_packets()
                yield self.tick()
    
        def send(self, to, packet, reason=None):
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            # TODO: transfer delay
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            packet.send(self, to, reason=reason)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.network.send_link(self, to, packet)
    
        def recv(self, packet):
            if packet.destination == self:
                packet.recv()
            else:
                self.buffer.add(packet)
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        '''
        def __eq__(self, other):
            return self.id == other.id
        '''
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        @property
        def community(self):
            return self.network.community[self]
    
        @property
        def links(self):
            '''
            Returns a list of connected links.
            '''
            links = {
                met: data
                for met, data in self.network[self].items()
                if data['state']
            }
            return links
    
        def __repr__(self):
            return 'Node(id={})'.format(self.id)
    
    
    class Buffer:
        def __init__(self, env, capacity=0):
            self.env = env
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
            if capacity is None or capacity <= 0:
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                self.capacity = float('inf')
            else:
                self.capacity = capacity
    
            self.buffer = OrderedDict()
            self.used = 0
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        def clean(self):
            packets_to_drop = []
    
            for packet in self:
                if packet.ttl < self.env.now:
                    packets_to_drop.append(packet)
    
            for packet in packets_to_drop:
                self.remove(packet)
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        def add(self, packet):
            if self.used < self.capacity:
                self.used += 1
                self.buffer[packet] = None
            else:
                raise Exception('buffer full')
    
        def remove(self, packet):
            self.used -= 1
            del self.buffer[packet]
    
        def __contains__(self, packet):
            return packet in self.buffer
    
        def __iter__(self):
            return iter(self.buffer)
    
        def __len__(self):
            return len(self.buffer)
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    class PacketGenerator(TickProcess):
        def __init__(self, ticks_per_packet=1, start_delay=0, time_to_live=60):
            super().__init__(ticks_per_packet)
    
            if time_to_live is None:
                raise ValueError('time_to_live must be specified')
    
            self.start_delay = start_delay
            self.time_to_live = time_to_live
            self.packets = []
    
        def process(self, network):
            packet_id = 0
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            def create(packet_id):
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                source, dest = random.choice(network.links)
                packet = Packet(packet_id, source, dest,
                                self.env.now + self.time_to_live, None)
                self.packets.append(packet)
                source.recv(packet)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            yield self.env.timeout(self.start_delay)
            print('starting packets')
            #@DEBUG
            for i in range(217):
                pass
            #@DEBUG
            while True:
                create(packet_id)
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                yield self.tick()
                packet_id += 1
    
        @property
        def recieved(self):
            return len([
                packet
                for packet in self.packets
                if packet.recieved
            ])
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        @property
        def recieved_count(self):
            return sum([
                packet.recieved_count
                for packet in self.packets
            ])
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        @property
        def sent(self):
            return sum([
                packet.sent
                for packet in self.packets
            ])
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        @property
        def stats(self):
            stats = {}
            for packet in self.packets:
                for stat, value in packet.stats.items():
                    if stat not in stats:
                        stats[stat] = value
                    else:
                        stats[stat] += value
            return stats
    
        @property
        def paths(self):
            return {
                packet: packet.path
                for packet in self.packets
                if packet.recieved
            }
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        @property
        def total(self):
            return len(self.packets)
    
        def __str__(self):
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            ret = "recieved: {}/{}, delivery ratio: {}, packets: {}, sent: {}, stats: {}, path: {}".format(
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                self.recieved,
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                self.recieved_count,
                self.recieved / len(self.packets),
                len(self.packets),
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                self.sent,
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                self.stats,
                sum([ len(path) for path in self.paths.values() ]) / self.recieved
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            )
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            '''
            for packet, path in self.paths.items():
                ret += '\n  {}: {}'.format(packet, path)
            '''
            return ret
    
    Jarrod Pas's avatar
    Jarrod Pas committed
    
    class Packet:
        def __init__(self, id, source, destination, ttl, payload):
            self.id = id
            self.source = source
            self.destination = destination
            self.ttl = ttl
            self.payload = payload
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.path = []
    
            self.stats = dict()
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.sent = 0
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.recieved = False
            self.recieved_count = 0
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.dropped = False
            self.dropped_count = 0
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        def send(self, a, b, reason=None):
            if reason is None:
                self.path.append((a.id, b.id))
            else:
                self.path.append((a.id, b.id, reason))
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            self.sent += 1
    
        def drop(self):
            self.dropped = True
            self.dropped_count += 1
    
    
    Jarrod Pas's avatar
    Jarrod Pas committed
        def recv(self):
            self.recieved = True
            self.recieved_count += 1
    
        def __str__(self):
    
    Jarrod Pas's avatar
    Jarrod Pas committed
            return "Packet(id={}, src={}, dst={})".format(
    
    Jarrod Pas's avatar
    Jarrod Pas committed
                self.id,
                self.source,
                self.destination
            )