Networking

"""
Networking for VRED 2018 and later.
Sync Node Transformations and send remote precdure calls over TCP and UDP
"""

import os, sys
import socket
import timeit
import threading
import ConfigParser
import json
from collections import deque

libs_dir = os.path.expanduser("~/Autodesk/")
sys.path.insert(0, os.path.abspath(libs_dir))

import pyuv
import msgpack


# To be inserted into VRED code editor
_editor_code = """

### Networking Configuration
selectVariantSet("VRTools_Networking")

networking = Networking()

print "Networking set up"

net_update_key = vrKey(Key_U, AltButton)
net_update_key.connect(networking.update_synced)

net_connect_key = vrKey(Key_C, AltButton)
net_connect_key.connect(networking.connect)

"""

# OSB file to load on code injection. Not used in networking
_file = None

class Networking():

    def __init__(self, ip = None, tcp_port = 40305, udp_port = 40306, send_rate = 0.05, heartbeat_rate = 10, config_file = ""):
        # Try getting settings from config file, might not be present
        global libs_dir
        self.config = ConfigParser.SafeConfigParser({'ip': "127.0.0.1", 'name':  "Unnamed", 'netid': '0', 'create_cam': 'False', 'passive': 'False', 'passive_nodes': "[]"})
        self.config.read([os.path.abspath(libs_dir + 'Networking.cfg'), os.path.abspath(libs_dir + config_file)])
 
        if not "Networking" in self.config.sections():
            self.config.add_section("Networking")
        
        print "Config", self.config.items("Networking")

        if not ip:
            ip = self.config.get("Networking", "ip")

        self.passive = self.config.getboolean("Networking", "passive")
        self.passive_nodes = json.loads(self.config.get("Networking", "passive_nodes"))
        
        # How often to sync transformations, in "every send_rate seconds". Defaults to every 100ms
        self.send_rate = send_rate
        self.send_timer = vrTimer(self.send_rate)
        self.send_timer.connect(self.send_synced)

        self.heartbeat_rate = heartbeat_rate
        self.heartbeat_timer = vrTimer(self.heartbeat_rate)
        self.heartbeat_timer.connect(self.send_heartbeat)

        # Tags by which to identify nodes to sync. You probably shouldn't change these.
        self.pos_tag = "_Network sync position"
        self.rot_tag = "_Network sync rotation"
        self.scale_tag = "_Network sync scale"
        self.state_tag = "_Network sync active state"
        addNodeTags(getRootNode(), [self.pos_tag, self.rot_tag, self.scale_tag, self.state_tag])
        
        # Updated and used by find_networked_objects(). Is used to identify nodes across the network
        # Uses ints instead of random strings (eg UUID) to avoid network overhead
        self.highest_net_id = 0
        
        # Maps network ID to node, what to sync, last value, seq counter
        self.synchronized = self.find_networked_objects()

        self.loop = pyuv.Loop.default_loop()
        self.uv_thread = threading.Thread(target=self.run_uv_loop)
        self.udp = pyuv.UDP(self.loop)
        self.tcp = pyuv.TCP(self.loop)
        self.tcp.nodelay(True)
        self.connected = False
        
        # Server address
        self.udp_addr = socket.getaddrinfo(ip, udp_port)[0][4]
        self.tcp_addr = socket.getaddrinfo(ip, tcp_port)[0][4]
        
        # maxlen denotes the number of ping timings used to calculate the average
        self.ping_times = deque(maxlen = 5)

     
    def run_uv_loop(self):
        """ Runs the uv loop. """
        print "Running background UV thread"
        self.loop.run()


    def update_synced(self):
        self.synchronized = self.find_networked_objects()


    def connect(self):
        if self.connected:
            print "Already connected"
            return

        try:
            self.udp.bind(("0.0.0.0", 0))
            self.udp.start_recv(self.udp_recv)

            self.tcp.bind(("0.0.0.0", 0))
            self.tcp.connect(self.tcp_addr, self.on_tcp_connect)

            if not self.passive:
                self.send_timer.setActive(True)
            self.heartbeat_timer.setActive(True)

            self.uv_thread.start()
        except Exception as e:
            raise e


    def on_tcp_connect(self, handle, error):
        if error is None:
            handle.start_read(self.tcp_recv)
            self.connected = True
            
            # Establish connection to server and grab initial data
            self.tcp.write(msgpack.packb(["hey"]) + "\n")


    def shutdown(self):
        try:
            self.udp.close()
        except Exception as e:
            print e

        try:
            self.tcp.close()
        except Exception as e:
            print e
        
        self.send_timer.setActive(False)
        self.heartbeat_timer.setActive(False)

        self.connected = False


    def reset_net_ids(self):
        """ Reset all net ids and start counting from zero again """
        sync_position = getNodesWithTag(self.pos_tag)
        sync_rotation = getNodesWithTag(self.rot_tag)
        sync_scale = getNodesWithTag(self.scale_tag)
        sync_state = getNodesWithTag(self.state_tag)
        for node in sync_position + sync_rotation + sync_scale + sync_state:
            if node.hasAttachment("ValuePair"):
                node.subAttachment(node.getAttachment("ValuePair"))
        self.highest_net_id = 0
        self.synchronized = self.find_networked_objects()


    def find_networked_objects(self):
        """ Create the map of ojects to sync. Needed after loading or adding objects to sync

        Returns a structure which maps network ids to an dict with node pointer plus
        should sync (boolean), sequence number and last value for
        position, scale, rotation and active state.
        """
        # TODO: This does not handle duplicate network ids (Created for example by duplicating networked nodes)

        sync_map = {}

        # Need to find highest Network ID first, so we can later add more valid ids to node which don't have one already
        sync_position = getNodesWithTag(self.pos_tag)
        sync_rotation = getNodesWithTag(self.rot_tag)
        sync_scale = getNodesWithTag(self.scale_tag)
        sync_state = getNodesWithTag(self.state_tag)

        # This assumes no other scripts use Value Pairs. Could be improved?
        for node in sync_position + sync_rotation + sync_scale + sync_state:
            if node.hasAttachment("ValuePair"):
                net_id = vrFieldAccess(node.getAttachment("ValuePair")).getMString("value")[0]
                self.highest_net_id = max(self.highest_net_id, int(net_id))
                
        for node in sync_position + sync_rotation + sync_scale + sync_state:
            node_data = {"node": node}

            node_data["pos"] = hasNodeTag(node, self.pos_tag)
            node_data["pos_last"] = node.getTranslation()
            node_data["pos_seq"] = 0

            node_data["rot"] = hasNodeTag(node, self.rot_tag)
            node_data["rot_last"] = node.getRotation()
            node_data["rot_seq"] = 0

            node_data["scale"] = hasNodeTag(node, self.scale_tag)
            node_data["scale_last"] = node.getScale()
            node_data["scale_seq"] = 0

            node_data["state"] = hasNodeTag(node, self.state_tag)
            node_data["state_last"] = node.getActive()
            node_data["state_seq"] = 0

            if node.hasAttachment("ValuePair"):
                net_id = int(vrFieldAccess(node.getAttachment("ValuePair")).getMString("value")[0])
            else:
                self.highest_net_id += 1
                net_id = self.highest_net_id
                attachment = createAttachment("ValuePair")
                vrFieldAccess(attachment).setMString("value", [str(net_id)])
                node.addAttachment(attachment)

            sync_map[net_id] = node_data

        print "Node sync update complete"
        return sync_map


    def udp_recv(self, handle, ip_port, flags, data, error):
        """ Receive udp packet """
        if data is not None:
            try:
                self.recv(msgpack.unpackb(data, use_list=False))
            except Exception as e:
                print "Receiving UDP message failed!"
                print e


    def tcp_recv(self, tcp_handle, data, error):
        if data is None:
            tcp_handle.close()
            return

        split = data.split(b"\n")
        for message in split:
            if len(message) > 0:
                try:
                    self.recv(msgpack.unpackb(message, use_list=False))
                except Exception as e:
                    print "Receiving TCP data failed!"
                    print e

    def recv(self, message):
        """ Handle different types of messages """
        tpe = message[0]
        if tpe == "pong":
            ping = timeit.default_timer() - message[1]
            self.ping_times.append(ping)
            avg = sum(self.ping_times) / len(self.ping_times)
            print "Current ping: {:.2f}, avg: {:.2f}".format(ping*1000, avg*1000)

        elif tpe == "pos" or tpe == "rot" or tpe == "scale" or tpe == "state":
            for pack in message[1]:
                self.apply_pack(tpe, pack[0], pack[1], pack[2])

        elif message[0] == "rpc":
            globals()[message[1]](*message[2])

        elif message[0] == "ho":
            print "Initial state", message
            for tpe, pack in message[1].items():
                for net_id, data in pack.items():
                    self.apply_pack(tpe, net_id, data[0], data[1])


        else:
            print "Unknown message type " + message[0]


    def apply_pack(self, tpe, net_id, seq, data):
        """ Apply pack data from sync or initial state message """
        if not net_id in self.synchronized:
            print "Unknown network ID {} found!".format(net_id)
            return

        id_data = self.synchronized[net_id]
        if id_data[tpe + "_seq"] < seq:
            id_data[tpe + "_seq"] = seq
            node = id_data["node"]
            if tpe == "pos":
                node.setTranslation(*data)
                id_data["pos_last"] = list(data)
            elif tpe == "rot":
                node.setRotation(*data)
                id_data["rot_last"] = list(data)
            elif tpe == "scale":
                node.setScale(*data)
                id_data["scale_last"] = list(data)
            else:
                node.setActive(data)
                id_data["state_last"] = data
        else:
            print "Got old data for {}: seq {}, current {}".format(id_data["node"].getName(), seq, id_data[tpe + "_seq"])

    def send_heartbeat(self):
        """ Send heartbeat packet to server. Include current time to measure rount trip time """
        if not self.udp:
            return
        self.udp.send(self.udp_addr, msgpack.packb(("ping", timeit.default_timer())))


    def rpc(self, fname, *args):
        """ Need to make sure args only contains vars msgpack can pack for transmission """
        if self.connected:
            self.tcp.write(msgpack.packb(("rpc", fname, args)) + "\n")
        else:
            print("Can't use rpc while not connected")


    def send_synced(self):
        """ Pack and send data (Position/Rotation/Scale/ActiveState) from network synchronized nodes
            This uses data from nodes in self.synchronized.
            Only data that has changed since last send will be transmitted.
            Sends data from up to 10 nodes in a single UDP packet for best speed/compatability.

            TODO: Could decouple this method from UDP, and instead send on a supplied channel
        """
        pos_packs = []
        rot_packs = []
        scale_packs = []
        state_packs = []

        # TODO: This could be written more concisely...
        for net_id, data in self.synchronized.items():
            if data["node"].getName() in self.passive_nodes:
                continue

            if data["pos"]:
                pos = data["node"].getTranslation()
                if data["pos_last"] != pos:
                    data["pos_last"] = pos
                    seq = data["pos_seq"] + 1
                    data["pos_seq"] = seq
                    pos_packs.append((net_id, seq, pos))

            if data["rot"]:
                rot = data["node"].getRotation()
                if data["rot_last"] != rot:
                    data["rot_last"] = rot
                    seq = data["rot_seq"] + 1
                    data["rot_seq"] = seq
                    rot_packs.append((net_id, seq, rot))

            if data["scale"]:
                scale = data["node"].getScale()
                if data["scale_last"] != scale:
                    data["scale_last"] = scale
                    seq = data["scale_seq"] + 1
                    data["scale_seq"] = seq
                    scale_packs.append((net_id, seq, scale))

            if data["state"]:
                state = data["node"].getActive()
                if data["state_last"] != state:
                    data["state_last"] = state
                    seq = data["state_seq"] + 1
                    data["state_seq"] = seq
                    state_packs.append((net_id, seq, state))


        # Send data packs in chunks of 10. This results in a packet of roughly 350 bytes,
        # which should not be fragmented or otherwise delayed on the wire.
        # These data chucks are not well compressible, so no additional compression is used
        # TODO: Quite some useless slicing and packing in here. Measure execution time / optimize?
        # (50 nodes position ~0.3ms)
        for i in range(0, max(len(pos_packs), len(rot_packs), len(scale_packs), len(state_packs)), 10):
            pospack = ("pos", pos_packs[i:i+10])
            rotpack = ("rot", rot_packs[i:i+10])
            scalepack = ("scale", scale_packs[i:i+10])
            statepack = ("state", state_packs[i:i+10])
            
            for pack in (pospack, rotpack, scalepack, statepack):
                if len(pack[1]) > 0:
                    self.udp.send(self.udp_addr, msgpack.packb(pack))