From 96e3228c9d8abc654a6f1572edb248a3b52a80e1 Mon Sep 17 00:00:00 2001
From: Aaron Tan <jangsutsr@gmail.com>
Date: Sat, 24 Sep 2016 00:46:42 -0400
Subject: [PATCH] Change global settings structure and fix potential celery
 issue.

---
 awx/main/tasks.py                             | 41 +++++++++++--------
 awx/settings/local_settings.py.docker_compose |  7 ++++
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/awx/main/tasks.py b/awx/main/tasks.py
index e2cb50c737..2e25ee39b4 100644
--- a/awx/main/tasks.py
+++ b/awx/main/tasks.py
@@ -501,7 +501,7 @@ class BaseTask(Task):
         return OrderedDict()
 
     def run_pexpect(self, instance, args, cwd, env, passwords, stdout_handle,
-                    output_replacements=None):
+                    output_replacements=None, runtime_flags={}):
         '''
         Run the given command using pexpect to capture output and provide
         passwords when requested.
@@ -517,13 +517,17 @@ class BaseTask(Task):
         if pexpect_sleep is not None:
             logger.info("Suspending Job Execution for QA Work")
             time.sleep(pexpect_sleep)
-        local_timeout = getattr(instance, 'timeout', 0)
-        global_timeout = getattr(settings, 'JOB_TIMEOUT', 0)
-        job_timeout = global_timeout if local_timeout == 0 else local_timeout
+        global_timeout = getattr(settings, 'DEFAULT_TIMEOUT', {})
+        cls_name = instance.__class__.__name__
+        if cls_name in global_timeout:
+            local_timeout = getattr(instance, 'timeout', 0)
+            job_timeout = global_timeout[cls_name] if local_timeout == 0 else local_timeout
+        else:
+            job_timeout = 0
         child = pexpect.spawnu(args[0], args[1:], cwd=cwd, env=env)
         child.logfile_read = logfile
-        canceled = False
-        self.timed_out = False
+        runtime_flags['canceled'] = False
+        runtime_flags['timed_out'] = False
         last_stdout_update = time.time()
         idle_timeout = self.get_idle_timeout()
         expect_list = []
@@ -546,22 +550,23 @@ class BaseTask(Task):
             # Refresh model instance from the database (to check cancel flag).
             instance = self.update_model(instance.pk)
             if instance.cancel_flag:
-                canceled = self.__handle_termination(instance, child)
+                runtime_flags['canceled'] = True
             elif job_timeout != 0 and (time.time() - job_start) > job_timeout:
-                self.__handle_termination(instance, child, is_cancel=False)
+                runtime_flags['timed_out'] = True
+            if any(list(runtime_flags.values())):
+                self._handle_termination(instance, child, is_cancel=runtime_flags['canceled'])
             if idle_timeout and (time.time() - last_stdout_update) > idle_timeout:
                 child.close(True)
-                canceled = True
-        if canceled:
+                runtime_flags['canceled'] = True
+        if runtime_flags['canceled']:
             return 'canceled', child.exitstatus
-        elif child.exitstatus == 0 and not self.timed_out:
+        elif child.exitstatus == 0 and not runtime_flags['timed_out']:
             return 'successful', child.exitstatus
         else:
             return 'failed', child.exitstatus
 
-    def __handle_termination(self, instance, job, is_cancel=True):
-        '''Helper function to properly terminate specified job and return correct
-        flags.
+    def _handle_termination(self, instance, job, is_cancel=True):
+        '''Helper function to properly terminate specified job.
 
         Args:
             instance: The corresponding model instance of this task.
@@ -570,7 +575,7 @@ class BaseTask(Task):
                 cancel_flag.
 
         Return:
-            True if is_cancel is set or None.
+            None.
         '''
         try:
             if tower_settings.AWX_PROOT_ENABLED and self.should_use_proot(instance):
@@ -618,6 +623,7 @@ class BaseTask(Task):
         instance.websocket_emit_status("running")
         status, rc, tb = 'error', None, ''
         output_replacements = []
+        runtime_flags = {}
         try:
             self.pre_run_hook(instance, **kwargs)
             if instance.cancel_flag:
@@ -662,7 +668,8 @@ class BaseTask(Task):
                 safe_args = self.wrap_args_with_ssh_agent(safe_args, ssh_key_path, ssh_auth_sock)
             instance = self.update_model(pk, job_args=json.dumps(safe_args),
                                          job_cwd=cwd, job_env=safe_env, result_stdout_file=stdout_filename)
-            status, rc = self.run_pexpect(instance, args, cwd, env, kwargs['passwords'], stdout_handle)
+            status, rc = self.run_pexpect(instance, args, cwd, env, kwargs['passwords'], stdout_handle,
+                                          runtime_flags=runtime_flags)
         except Exception:
             if status != 'canceled':
                 tb = traceback.format_exc()
@@ -680,7 +687,7 @@ class BaseTask(Task):
             try:
                 stdout_handle.flush()
                 stdout_handle.close()
-                if getattr(self, 'timed_out', False):
+                if runtime_flags.get('timed_out', False):
                     job_explanation = "Job terminated due to timeout"
             except Exception:
                 pass
diff --git a/awx/settings/local_settings.py.docker_compose b/awx/settings/local_settings.py.docker_compose
index 088f651e77..5a2306bc17 100644
--- a/awx/settings/local_settings.py.docker_compose
+++ b/awx/settings/local_settings.py.docker_compose
@@ -276,3 +276,10 @@ TEST_OPENSTACK_PROJECT = ''
 # Azure credentials.
 TEST_AZURE_USERNAME = ''
 TEST_AZURE_KEY_DATA = ''
+
+# Exemplary job timeout settings
+# DEFAULT_TIMEOUT = {
+#   'Job': 10,
+#   'InventoryUpdate': 15,
+#   'ProjectUpdate': 20,
+# }