[autotest] Merge site_server_job

BUG=chromium:684311
TEST=unittests, moblab trybot

Change-Id: I2528c2166c7753040d5799fad030e94f757e5590
Reviewed-on: https://chromium-review.googlesource.com/562220
Commit-Ready: Prathmesh Prabhu <[email protected]>
Tested-by: Prathmesh Prabhu <[email protected]>
Reviewed-by: Prathmesh Prabhu <[email protected]>
diff --git a/server/server_job.py b/server/server_job.py
index 8a04bd0..23359b5 100644
--- a/server/server_job.py
+++ b/server/server_job.py
@@ -16,6 +16,7 @@
 import getpass
 import itertools
 import logging
+import multiprocessing
 import os
 import pickle
 import platform
@@ -37,6 +38,8 @@
 from autotest_lib.client.common_lib import packages
 from autotest_lib.client.common_lib import utils
 from autotest_lib.server import profilers
+from autotest_lib.server import site_gtest_runner
+from autotest_lib.server import site_server_job_utils
 from autotest_lib.server import subcommand
 from autotest_lib.server import test
 from autotest_lib.server import utils as server_utils
@@ -197,7 +200,7 @@
             job._parse_status(rendered_entry)
 
 
-class base_server_job(base_job.base_job):
+class server_job(base_job.base_job):
     """The server-side concrete implementation of base_job.
 
     Optional properties provided by this implementation:
@@ -259,8 +262,8 @@
         @param in_lab: Boolean that indicates if this is running in the lab
                        environment.
         """
-        super(base_server_job, self).__init__(resultdir=resultdir,
-                                              test_retry=test_retry)
+        super(server_job, self).__init__(resultdir=resultdir,
+                                         test_retry=test_retry)
         self.test_retry = test_retry
         self.control = control
         self._uncollected_log_file = os.path.join(self.resultdir,
@@ -580,9 +583,9 @@
         """Wrap function as appropriate for calling by parallel_simple."""
         # machines could be a list of dictionaries, e.g.,
         # [{'host_attributes': {}, 'hostname': '100.96.51.226'}]
-        # The dictionary is generated in base_server_job.__init__, refer to
+        # The dictionary is generated in server_job.__init__, refer to
         # variable machine_dict_list, then passed in with namespace, see method
-        # base_server_job._make_namespace.
+        # server_job._make_namespace.
         # To compare the machinese to self.machines, which is a list of machine
         # hostname, we need to convert machines back to a list of hostnames.
         # Note that the order of hostnames doesn't matter, as is_forking will be
@@ -662,6 +665,94 @@
         return success_machines
 
 
+    def distribute_across_machines(self, tests, machines,
+                                   continuous_parsing=False):
+        """Run each test in tests once using machines.
+
+        Instead of running each test on each machine like parallel_on_machines,
+        run each test once across all machines. Put another way, the total
+        number of tests run by parallel_on_machines is len(tests) *
+        len(machines). The number of tests run by distribute_across_machines is
+        len(tests).
+
+        Args:
+            tests: List of tests to run.
+            machines: List of machines to use.
+            continuous_parsing: Bool, if true parse job while running.
+        """
+        # The Queue is thread safe, but since a machine may have to search
+        # through the queue to find a valid test the lock provides exclusive
+        # queue access for more than just the get call.
+        test_queue = multiprocessing.JoinableQueue()
+        test_queue_lock = multiprocessing.Lock()
+
+        unique_machine_attributes = []
+        sub_commands = []
+        work_dir = self.resultdir
+
+        for machine in machines:
+            if 'group' in self.resultdir:
+                work_dir = os.path.join(self.resultdir, machine)
+
+            mw = site_server_job_utils.machine_worker(self,
+                                                      machine,
+                                                      work_dir,
+                                                      test_queue,
+                                                      test_queue_lock,
+                                                      continuous_parsing)
+
+            # Create the subcommand instance to run this machine worker.
+            sub_commands.append(subcommand.subcommand(mw.run,
+                                                      [],
+                                                      work_dir))
+
+            # To (potentially) speed up searching for valid tests create a list
+            # of unique attribute sets present in the machines for this job. If
+            # sets were hashable we could just use a dictionary for fast
+            # verification. This at least reduces the search space from the
+            # number of machines to the number of unique machines.
+            if not mw.attribute_set in unique_machine_attributes:
+                unique_machine_attributes.append(mw.attribute_set)
+
+        # Only queue tests which are valid on at least one machine.  Record
+        # skipped tests in the status.log file using record_skipped_test().
+        for test_entry in tests:
+            # Check if it's an old style test entry.
+            if len(test_entry) > 2 and not isinstance(test_entry[2], dict):
+                test_attribs = {'include': test_entry[2]}
+                if len(test_entry) > 3:
+                    test_attribs['exclude'] = test_entry[3]
+                if len(test_entry) > 4:
+                    test_attribs['attributes'] = test_entry[4]
+
+                test_entry = list(test_entry[:2])
+                test_entry.append(test_attribs)
+
+            ti = site_server_job_utils.test_item(*test_entry)
+            machine_found = False
+            for ma in unique_machine_attributes:
+                if ti.validate(ma):
+                    test_queue.put(ti)
+                    machine_found = True
+                    break
+            if not machine_found:
+                self.record_skipped_test(ti)
+
+        # Run valid tests and wait for completion.
+        subcommand.parallel(sub_commands)
+
+
+    def record_skipped_test(self, skipped_test, message=None):
+        """Insert a failure record into status.log for this test."""
+        msg = message
+        if msg is None:
+            msg = 'No valid machines found for test %s.' % skipped_test
+        logging.info(msg)
+        self.record('START', None, skipped_test.test_name)
+        self.record('INFO', None, skipped_test.test_name, msg)
+        self.record('END TEST_NA', None, skipped_test.test_name, msg)
+
+
     def _has_failed_tests(self):
         """Parse status log for failed tests.
 
@@ -762,8 +853,11 @@
 
         self.aborted = False
         namespace.update(self._make_namespace())
-        namespace.update({'args' : self.args,
-                          'job_labels' : job_labels})
+        namespace.update({
+                'args': self.args,
+                'job_labels': job_labels,
+                'gtest_runner': site_gtest_runner.gtest_runner(),
+        })
         test_start_time = int(time.time())
 
         if self.resultdir:
@@ -1500,12 +1594,3 @@
         raise error.AutoservError('Could not obtain HostInfo for hostname %s' %
                                   hostname)
     return host_info_store
-
-
-site_server_job = utils.import_site_class(
-    __file__, "autotest_lib.server.site_server_job", "site_server_job",
-    base_server_job)
-
-
-class server_job(site_server_job):
-    pass