Merge pull request #99 from rqueensen/support/rqueensen/execution_assignment

Execution Server Assignment Logic Change

Author: James Bell

setup.py

@@ -4,7 +4,7 @@
 
 setup(
     name='tgt_grease',
-    version='2.5.0',
+    version='2.6.0',
     license="MIT",
     description='Modern distributed automation engine built with love by Target',
     long_description="""
@@ -41,8 +41,8 @@
         'elasticsearch',
         'kafka-python'
     ] + (
-         ["pywin32"] if "nt" == os.name else []
-        ),
+        ["pywin32"] if "nt" == os.name else []
+    ),
     include_package_data=True,
     zip_safe=False,
     scripts=[

tgt_grease/enterprise/Model/CentralScheduling.py

@@ -3,6 +3,7 @@
 from .DeDuplication import Deduplication
 import pymongo
 import datetime
+import sys
 
 
 class Scheduling(object):
@@ -38,18 +39,21 @@ def scheduleDetection(self, source, configName, data):
         """
         if len(data) is 0 or not isinstance(data, list):
             self.ioc.getLogger().trace(
-                "Data provided empty or is not type list type: [{0}] len: [{1}]".format(str(type(data)), len(data)),
+                "Data provided empty or is not type list type: [{0}] len: [{1}]".format(
+                    str(type(data)), len(data)),
                 trace=True
             )
             return False
-        self.ioc.getLogger().trace("Preparing to schedule [{0}] source objects".format(len(data)), trace=True)
+        self.ioc.getLogger().trace(
+            "Preparing to schedule [{0}] source objects".format(len(data)), trace=True)
         sourceCollect = self.ioc.getCollection('SourceData')
         jServerCollect = self.ioc.getCollection('JobServer')
         # begin scheduling loop of each block
         for elem in data:
             if not isinstance(elem, dict):
                 self.ioc.getLogger().warning(
-                    "Element from data not of type dict! Got [{0}] DROPPED".format(str(type(elem))),
+                    "Element from data not of type dict! Got [{0}] DROPPED".format(
+                        str(type(elem))),
                     notify=False
                 )
                 continue
@@ -93,7 +97,8 @@ def scheduleDetection(self, source, configName, data):
                 )
             else:
                 self.ioc.getLogger().warning(
-                    "Failed to find detection server for data object from source [{0}]; DROPPED".format(source),
+                    "Failed to find detection server for data object from source [{0}]; DROPPED".format(
+                        source),
                     notify=False
                 )
                 self.ioc.getLogger().warning(
@@ -144,15 +149,34 @@ def determineDetectionServer(self):
             tuple: MongoDB Object ID of server & current job count
 
         """
-        result = self.ioc.getCollection('JobServer').find({
-            'active': True,
-            'prototypes': 'detect'
-        }).sort('jobs', pymongo.ASCENDING).limit(1)
-        if result.count():
-            return str(result[0]['_id']), int(result[0]['jobs'])
-        else:
+        servers = [
+            (server.get('_id'), server.get('jobs')) for server in self.ioc.getCollection('JobServer').find(
+                {
+                    'active': True,
+                    'prototypes': 'detect',
+                }
+            )
+        ]
+
+        best_server = {}
+        for (server, total_jobs) in servers:
+            active_jobs = self.ioc.getCollection('SourceData').find(
+                {
+                    'grease_data.detection.server': server,
+                    'grease_data.detection.end': None,
+                }
+            ).count()
+            if active_jobs < best_server.get('active_jobs', sys.maxsize):
+                best_server['_id'] = server
+                best_server['total_jobs'] = total_jobs
+                best_server['active_jobs'] = active_jobs
+
+        if not best_server.get('_id'):
+            self.ioc.getLogger().error("No active detection server found!")
             return "", 0
 
+        return best_server.get('_id'), best_server.get('total_jobs')
+
     def determineSchedulingServer(self):
         """Determines scheduling server to use
 
@@ -162,15 +186,34 @@ def determineSchedulingServer(self):
             tuple: MongoDB Object ID of server & current job count
 
         """
-        result = self.ioc.getCollection('JobServer').find({
-            'active': True,
-            'prototypes': 'schedule'
-        }).sort('jobs', pymongo.DESCENDING).limit(1)
-        if result.count():
-            return str(result[0]['_id']), int(result[0]['jobs'])
-        else:
+        servers = [
+            (server.get('_id'), server.get('jobs')) for server in self.ioc.getCollection('JobServer').find(
+                {
+                    'active': True,
+                    'prototypes': 'schedule',
+                }
+            )
+        ]
+
+        best_server = {}
+        for (server, total_jobs) in servers:
+            active_jobs = self.ioc.getCollection('SourceData').find(
+                {
+                    'grease_data.scheduling.server': server,
+                    'grease_data.scheduling.end': None,
+                }
+            ).count()
+            if active_jobs < best_server.get('active_jobs', sys.maxsize):
+                best_server['_id'] = server
+                best_server['total_jobs'] = total_jobs
+                best_server['active_jobs'] = active_jobs
+
+        if not best_server.get('_id'):
+            self.ioc.getLogger().error("No active scheduling server found!")
             return "", 0
 
+        return best_server.get('_id'), best_server.get('total_jobs')
+
     def determineExecutionServer(self, role):
         """Determines execution server to use
 
@@ -180,11 +223,31 @@ def determineExecutionServer(self, role):
             str: MongoDB Object ID of server; if one cannot be found then string will be empty
 
         """
-        result = self.ioc.getCollection('JobServer').find({
-            'active': True,
-            'roles': str(role)
-        }).sort('jobs', pymongo.DESCENDING).limit(1)
-        if result.count():
-            return str(result[0]['_id']), int(result[0]['jobs'])
-        else:
+        servers = [
+            (server.get('_id'), server.get('jobs')) for server in self.ioc.getCollection('JobServer').find(
+                {
+                    'active': True,
+                    'roles': str(role),
+                }
+            )
+        ]
+
+        best_server = {}
+        for (server, total_jobs) in servers:
+            active_jobs = self.ioc.getCollection('SourceData').find(
+                {
+                    'grease_data.execution.server': server,
+                    'grease_data.execution.completeTime': None,
+                    'grease_data.execution.failures': {'$lt': 6},
+                }
+            ).count()
+            if active_jobs < best_server.get('active_jobs', sys.maxsize):
+                best_server['_id'] = server
+                best_server['total_jobs'] = total_jobs
+                best_server['active_jobs'] = active_jobs
+
+        if not best_server.get('_id'):
+            self.ioc.getLogger().error("No active job server found with role {}!".format(role))
             return "", 0
+
+        return best_server.get('_id'), best_server.get('total_jobs')