From 3131bb9306bf8293cab494885b9fdfe8939fad1b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Robert=20Chac=C3=B3n?= Date: Wed, 20 Aug 2025 12:10:20 -0600 Subject: [PATCH] Update integrationCommon.py --- src/integrationCommon.py | 441 ++++++++++++++++++++++++--------------- 1 file changed, 269 insertions(+), 172 deletions(-) diff --git a/src/integrationCommon.py b/src/integrationCommon.py index f57f574e..31513817 100644 --- a/src/integrationCommon.py +++ b/src/integrationCommon.py @@ -1,7 +1,7 @@ -# Provides common functionality shared between -# integrations. +# Optimized version of integrationCommon.py for better performance with large datasets +# Provides common functionality shared between integrations. -from typing import List, Any +from typing import List, Any, Dict, Set from liblqos_python import allowed_subnets, ignore_subnets, generated_pn_download_mbps, generated_pn_upload_mbps, \ circuit_name_use_address, upstream_bandwidth_capacity_download_mbps, upstream_bandwidth_capacity_upload_mbps, \ find_ipv6_using_mikrotik, exclude_sites, bandwidth_overhead_factor, committed_bandwidth_multiplier, \ @@ -120,16 +120,19 @@ class NetworkNode: class NetworkGraph: - # Defines a network as a graph topology - # allowing any integration to build the - # graph via a common API, emitting - # ShapedDevices and network.json files - # via a common interface. - + # OPTIMIZED VERSION: Uses dictionaries and caching for O(1) lookups + # instead of O(n) linear searches + nodes: List ipv4ToIPv6: Any excludeSites: List # Copied to allow easy in-test patching exceptionCPEs: Any + + # Performance optimization: Index structures for O(1) lookups + _id_to_index: Dict[str, int] # Maps node ID to index in nodes list + _name_to_index: Dict[str, int] # Maps display name to index + _children_cache: Dict[int, List[int]] # Cache parent->children mapping + _cache_valid: bool # Flag to invalidate cache when nodes change def __init__(self) -> None: self.nodes = [ @@ -138,6 +141,13 @@ class NetworkGraph: ] self.excludeSites = exclude_sites() self.exceptionCPEs = exception_cpes() + + # Initialize optimization structures + self._id_to_index = {"FakeRoot": 0} + self._name_to_index = {"Shaper Root": 0} + self._children_cache = {} + self._cache_valid = False + if find_ipv6_using_mikrotik(): from mikrotikFindIPv6 import pullMikrotikIPv6 self.ipv4ToIPv6 = pullMikrotikIPv6() @@ -149,18 +159,34 @@ class NetworkGraph: # If a site is excluded (via excludedSites in lqos.conf) # it won't be added if not node.displayName in self.excludeSites: - # TODO: Fixup exceptionCPE handling - #print(self.excludeSites) - #if node.displayName in self.exceptionCPEs.keys(): - # node.parentId = self.exceptionCPEs[node.displayName] + # Add to main list + idx = len(self.nodes) self.nodes.append(node) + + # Update indexes for O(1) lookup + self._id_to_index[node.id] = idx + self._name_to_index[node.displayName] = idx + + # Invalidate cache + self._cache_valid = False def replaceRootNode(self, node: NetworkNode) -> None: # Replaces the automatically generated root node # with a new node. Useful when you have a top-level # node specified (e.g. "uispSite" in the UISP # integration) + old_node = self.nodes[0] self.nodes[0] = node + + # Update indexes + if old_node.id in self._id_to_index: + del self._id_to_index[old_node.id] + if old_node.displayName in self._name_to_index: + del self._name_to_index[old_node.displayName] + + self._id_to_index[node.id] = 0 + self._name_to_index[node.displayName] = 0 + self._cache_valid = False def addNodeAsChild(self, parent: str, node: NetworkNode) -> None: # Searches the existing graph for a named parent, @@ -171,68 +197,89 @@ class NetworkGraph: # specifying the parent - we're assuming you really # meant it. if node.displayName in self.excludeSites: return - parentIdx = 0 - for (i, node) in enumerate(self.nodes): - if node.id == parent: - parentIdx = i + + # O(1) lookup instead of O(n) search + parentIdx = self._id_to_index.get(parent, 0) node.parentIndex = parentIdx + + idx = len(self.nodes) self.nodes.append(node) + + # Update indexes + self._id_to_index[node.id] = idx + self._name_to_index[node.displayName] = idx + self._cache_valid = False def __reparentById(self) -> None: - # Scans the entire node tree, searching for parents - # by name. Entries are re-mapped to match the named - # parents. You can use this to build a tree from a - # blob of raw data. + # OPTIMIZED: Uses dictionary lookups instead of nested O(n²) loops cached_root_list = promote_to_root_list() + + # Build a set for O(1) membership testing + root_list_set = set(cached_root_list) + for child in self.nodes: if child.parentId != "": - for (i, node) in enumerate(self.nodes): - # If the node is in `promote_to_root_list`, the parent for - # the node is set to 0, which is the root node. Otherwise, - # the parent is set to the node with the same id as the - # parentId of the child node. - if self.nodes[self.findNodeIndexById(child.parentId)].displayName in cached_root_list: + # O(1) lookup instead of O(n) search + parent_idx = self._id_to_index.get(child.parentId, -1) + + if parent_idx != -1: + # Check if parent should be promoted to root + parent_node = self.nodes[parent_idx] + if parent_node.displayName in root_list_set: child.parentIndex = 0 - elif node.id == child.parentId: - child.parentIndex = i + else: + child.parentIndex = parent_idx + + # Invalidate children cache since parent relationships changed + self._cache_valid = False def findNodeIndexById(self, id: str) -> int: - # Finds a single node by identity(id) - # Return -1 if not found - for (i, node) in enumerate(self.nodes): - if node.id == id: - return i - return -1 + # OPTIMIZED: O(1) dictionary lookup instead of O(n) search + return self._id_to_index.get(id, -1) def findNodeIndexByName(self, name: str) -> int: - # Finds a single node by identity(name) - # Return -1 if not found - for (i, node) in enumerate(self.nodes): - if node.displayName == name: - return i - return -1 + # OPTIMIZED: O(1) dictionary lookup instead of O(n) search + return self._name_to_index.get(name, -1) + + def _buildChildrenCache(self) -> None: + # Build parent->children mapping cache + self._children_cache = {} + for i, node in enumerate(self.nodes): + parent_idx = node.parentIndex + if parent_idx not in self._children_cache: + self._children_cache[parent_idx] = [] + self._children_cache[parent_idx].append(i) + self._cache_valid = True def findChildIndices(self, parentIndex: int) -> List: - # Returns the indices of all nodes with a - # parentIndex equal to the specified parameter - result = [] - for (i, node) in enumerate(self.nodes): - if node.parentIndex == parentIndex: - result.append(i) - return result + # OPTIMIZED: O(1) cached lookup instead of O(n) search + if not self._cache_valid: + self._buildChildrenCache() + return self._children_cache.get(parentIndex, []) def __promoteClientsWithChildren(self) -> None: - # Searches for client sites that have children, - # and changes their node type to clientWithChildren - for (i, node) in enumerate(self.nodes): + # OPTIMIZED: Single pass with cached children lookup + if not self._cache_valid: + self._buildChildrenCache() + + for i, node in enumerate(self.nodes): if node.type == NodeType.client: - for child in self.findChildIndices(i): - if self.nodes[child].type != NodeType.device: + # O(1) lookup of children + children = self._children_cache.get(i, []) + for child_idx in children: + if self.nodes[child_idx].type != NodeType.device: node.type = NodeType.clientWithChildren + break # No need to check other children def __clientsWithChildrenToSites(self) -> None: + # OPTIMIZED: Batch processing with deferred reparenting toAdd = [] - for (i, node) in enumerate(self.nodes): + reparent_map = {} # Store reparenting operations + + if not self._cache_valid: + self._buildChildrenCache() + + for i, node in enumerate(self.nodes): if node.type == NodeType.clientWithChildren: siteNode = NetworkNode( id=node.id + "_gen", @@ -243,60 +290,88 @@ class NetworkGraph: node.parentId = siteNode.id if node.type == NodeType.clientWithChildren: node.type = NodeType.client - for child in self.findChildIndices(i): - if self.nodes[child].type == NodeType.client or self.nodes[child].type == NodeType.clientWithChildren or self.nodes[child].type == NodeType.site: - self.nodes[child].parentId = siteNode.id + + # Store reparenting operations for batch processing + children = self._children_cache.get(i, []) + for child_idx in children: + child = self.nodes[child_idx] + if child.type in (NodeType.client, NodeType.clientWithChildren, NodeType.site): + reparent_map[child_idx] = siteNode.id + toAdd.append(siteNode) - + + # Batch add new nodes for n in toAdd: self.addRawNode(n) - + + # Batch reparent + for child_idx, new_parent_id in reparent_map.items(): + self.nodes[child_idx].parentId = new_parent_id + self.__reparentById() def __findUnconnectedNodes(self) -> List: - # Performs a tree-traversal and finds any nodes that - # aren't connected to the root. This is a "sanity check", - # and also an easy way to handle "flat" topologies and - # ensure that the unconnected nodes are re-connected to - # the root. - visited = [] + # OPTIMIZED: Uses set for O(1) membership testing + visited = set() next = [0] - - while len(next) > 0: - nextTraversal = next.pop() - visited.append(nextTraversal) - for idx in self.findChildIndices(nextTraversal): - if idx not in visited: - next.append(idx) - - result = [] - for i, n in enumerate(self.nodes): - if i not in visited: - result.append(i) - return result + + if not self._cache_valid: + self._buildChildrenCache() + + while next: + current = next.pop() + if current in visited: + continue + visited.add(current) + + # O(1) lookup of children + children = self._children_cache.get(current, []) + for child_idx in children: + if child_idx not in visited: + next.append(child_idx) + + # Find unvisited nodes + all_indices = set(range(len(self.nodes))) + unconnected = list(all_indices - visited) + return unconnected def __reconnectUnconnected(self): - # Finds any unconnected nodes and reconnects - # them to the root - for idx in self.__findUnconnectedNodes(): - if self.nodes[idx].type == NodeType.site: - self.nodes[idx].parentIndex = 0 - for idx in self.__findUnconnectedNodes(): - if self.nodes[idx].type == NodeType.clientWithChildren: - self.nodes[idx].parentIndex = 0 - for idx in self.__findUnconnectedNodes(): - if self.nodes[idx].type == NodeType.client: + # OPTIMIZED: Single pass with type grouping + unconnected = self.__findUnconnectedNodes() + + # Group by type for batch processing + by_type = { + NodeType.site: [], + NodeType.clientWithChildren: [], + NodeType.client: [] + } + + for idx in unconnected: + node_type = self.nodes[idx].type + if node_type in by_type: + by_type[node_type].append(idx) + + # Reconnect in priority order + for node_type in [NodeType.site, NodeType.clientWithChildren, NodeType.client]: + for idx in by_type[node_type]: self.nodes[idx].parentIndex = 0 + + self._cache_valid = False def prepareTree(self) -> None: # Helper function that calls all the cleanup and mapping # functions in the right order. Unless you are doing # something special, you can use this instead of # calling the functions individually + print("PrepareTree: Starting reparenting...") self.__reparentById() + print("PrepareTree: Promoting clients with children...") self.__promoteClientsWithChildren() + print("PrepareTree: Converting clients with children to sites...") self.__clientsWithChildrenToSites() + print("PrepareTree: Reconnecting unconnected nodes...") self.__reconnectUnconnected() + print("PrepareTree: Complete") def doesNetworkJsonExist(self): # Returns true if "network.json" exists, false otherwise @@ -311,10 +386,14 @@ class NetworkGraph: topLevelNode = {} self.__visited = [] # Protection against loops - never visit twice - for child in self.findChildIndices(0): + if not self._cache_valid: + self._buildChildrenCache() + + # O(1) lookup of root's children + root_children = self._children_cache.get(0, []) + for child in root_children: if child > 0 and self.__isSite(child): - topLevelNode[self.nodes[child].displayName] = self.__buildNetworkObject( - child) + topLevelNode[self.nodes[child].displayName] = self.__buildNetworkObject(child) del self.__visited @@ -342,11 +421,14 @@ class NetworkGraph: } children = {} hasChildren = False - for child in self.findChildIndices(idx): + + # O(1) lookup of children + child_indices = self._children_cache.get(idx, []) + for child in child_indices: if child > 0 and self.__isSite(child) and child not in self.__visited: - children[self.nodes[child].displayName] = self.__buildNetworkObject( - child) + children[self.nodes[child].displayName] = self.__buildNetworkObject(child) hasChildren = True + if hasChildren: node["children"] = children return node @@ -365,82 +447,93 @@ class NetworkGraph: ipv6.append(self.ipv4ToIPv6[ip]) def createShapedDevices(self): - import csv - # Builds ShapedDevices.csv from the network tree. - circuits = [] - for (i, node) in enumerate(self.nodes): - if node.type == NodeType.client: - parent = self.nodes[node.parentIndex].displayName - if parent == "Shaper Root": parent = "" - - if circuit_name_use_address(): - displayNameToUse = node.address - else: - if node.type == NodeType.client: - displayNameToUse = node.displayName - else: - displayNameToUse = node.customerName + " (" + nodeTypeToString(node.type) + ")" - circuit = { - "id": node.id, - "name": displayNameToUse, - "parent": parent, - "download": node.downloadMbps, - "upload": node.uploadMbps, - "devices": [] - } - for child in self.findChildIndices(i): - if self.nodes[child].type == NodeType.device and (len(self.nodes[child].ipv4)+len(self.nodes[child].ipv6)>0): - ipv4 = self.nodes[child].ipv4 - ipv6 = self.nodes[child].ipv6 - self.__addIpv6FromMap(ipv4, ipv6) - device = { - "id": self.nodes[child].id, - "name": self.nodes[child].displayName, - "mac": self.nodes[child].mac, - "ipv4": ipv4, - "ipv6": ipv6, - } - circuit["devices"].append(device) - if len(circuit["devices"]) > 0: - circuits.append(circuit) - with open('ShapedDevices.csv', 'w', newline='') as csvfile: - wr = csv.writer(csvfile, quoting=csv.QUOTE_ALL) - wr.writerow(['Circuit ID', 'Circuit Name', 'Device ID', 'Device Name', 'Parent Node', 'MAC', - 'IPv4', 'IPv6', 'Download Min', 'Upload Min', 'Download Max', 'Upload Max', 'Comment']) - for circuit in circuits: - for device in circuit["devices"]: - #Remove brackets and quotes of list so LibreQoS.py can parse it - device["ipv4"] = str(device["ipv4"]).replace('[','').replace(']','').replace("'",'') - device["ipv6"] = str(device["ipv6"]).replace('[','').replace(']','').replace("'",'') - if circuit["upload"] is None: - circuit["upload"] = 0.0 - if circuit["download"] is None: - circuit["download"] = 0.0 - row = [ - circuit["id"], - circuit["name"], - device["id"], - device["name"], - circuit["parent"], - device["mac"], - device["ipv4"], - device["ipv6"], - int(1), - int(1), - int(float(circuit["download"]) * client_bandwidth_multiplier()), - int(float(circuit["upload"]) * client_bandwidth_multiplier()), - "" - ] - wr.writerow(row) + import csv + # Builds ShapedDevices.csv from the network tree. + circuits = [] + + if not self._cache_valid: + self._buildChildrenCache() + + for i, node in enumerate(self.nodes): + if node.type == NodeType.client: + parent = self.nodes[node.parentIndex].displayName + if parent == "Shaper Root": parent = "" - # If we have an "appendToShapedDevices.csv" file, it gets appended to the end of the file. - # This is useful for adding devices that are not in the network tree, such as a - # "default" device that gets all the traffic that doesn't match any other device. - if os.path.isfile('appendToShapedDevices.csv'): - with open('appendToShapedDevices.csv', 'r') as f: - reader = csv.reader(f) - for row in reader: - wr.writerow(row) + if circuit_name_use_address(): + displayNameToUse = node.address + else: + if node.type == NodeType.client: + displayNameToUse = node.displayName + else: + displayNameToUse = node.customerName + " (" + nodeTypeToString(node.type) + ")" + + circuit = { + "id": node.id, + "name": displayNameToUse, + "parent": parent, + "download": node.downloadMbps, + "upload": node.uploadMbps, + "devices": [] + } + + # O(1) lookup of children + child_indices = self._children_cache.get(i, []) + for child_idx in child_indices: + child = self.nodes[child_idx] + if child.type == NodeType.device and (len(child.ipv4) + len(child.ipv6) > 0): + ipv4 = child.ipv4 + ipv6 = child.ipv6 + self.__addIpv6FromMap(ipv4, ipv6) + device = { + "id": child.id, + "name": child.displayName, + "mac": child.mac, + "ipv4": ipv4, + "ipv6": ipv6, + } + circuit["devices"].append(device) + + if len(circuit["devices"]) > 0: + circuits.append(circuit) + + with open('ShapedDevices.csv', 'w', newline='') as csvfile: + wr = csv.writer(csvfile, quoting=csv.QUOTE_ALL) + wr.writerow(['Circuit ID', 'Circuit Name', 'Device ID', 'Device Name', 'Parent Node', 'MAC', + 'IPv4', 'IPv6', 'Download Min', 'Upload Min', 'Download Max', 'Upload Max', 'Comment']) + for circuit in circuits: + for device in circuit["devices"]: + #Remove brackets and quotes of list so LibreQoS.py can parse it + device["ipv4"] = str(device["ipv4"]).replace('[','').replace(']','').replace("'",'') + device["ipv6"] = str(device["ipv6"]).replace('[','').replace(']','').replace("'",'') + if circuit["upload"] is None: + circuit["upload"] = 0.0 + if circuit["download"] is None: + circuit["download"] = 0.0 + row = [ + circuit["id"], + circuit["name"], + device["id"], + device["name"], + circuit["parent"], + device["mac"], + device["ipv4"], + device["ipv6"], + int(1), + int(1), + int(float(circuit["download"]) * client_bandwidth_multiplier()), + int(float(circuit["upload"]) * client_bandwidth_multiplier()), + "" + ] + wr.writerow(row) + + # If we have an "appendToShapedDevices.csv" file, it gets appended to the end of the file. + # This is useful for adding devices that are not in the network tree, such as a + # "default" device that gets all the traffic that doesn't match any other device. + if os.path.isfile('appendToShapedDevices.csv'): + with open('appendToShapedDevices.csv', 'r') as f: + reader = csv.reader(f) + for row in reader: + wr.writerow(row) def plotNetworkGraph(self, showClients=False): # Requires `pip install graphviz` to function. @@ -458,7 +551,10 @@ class NetworkGraph: dot = graphviz.Digraph( 'network', comment="Network Graph", engine="dot", graph_attr={'rankdir':'LR'}) - for (i, node) in enumerate(self.nodes): + if not self._cache_valid: + self._buildChildrenCache() + + for i, node in enumerate(self.nodes): if ((node.type != NodeType.client and node.type != NodeType.device) or showClients): color = "white" match node.type: @@ -469,11 +565,12 @@ class NetworkGraph: case NodeType.device: color = "white" case default: color = "grey" dot.node("N" + str(i), node.displayName, color=color) - children = self.findChildIndices(i) + + # O(1) lookup of children + children = self._children_cache.get(i, []) for child in children: if child != i: if (self.nodes[child].type != NodeType.client and self.nodes[child].type != NodeType.device) or showClients: dot.edge("N" + str(i), "N" + str(child)) dot = dot.unflatten(stagger=3)#, fanout=True) dot.render("network") -