Merge pull request #621 from ryanpetrello/set_stat_workflow_race_condition

don't process artifacts from custom `set_stat` calls asynchronously

awx/lib/awx_display_callback/module.py

@@ -18,7 +18,11 @@
 from __future__ import (absolute_import, division, print_function)
 
 # Python
+import codecs
 import contextlib
+import json
+import os
+import stat
 import sys
 import uuid
 from copy import copy
@@ -292,10 +296,22 @@ class BaseCallbackModule(CallbackBase):
             failures=stats.failures,
             ok=stats.ok,
             processed=stats.processed,
-            skipped=stats.skipped,
-            artifact_data=stats.custom.get('_run', {}) if hasattr(stats, 'custom') else {}
+            skipped=stats.skipped
         )
 
+        # write custom set_stat artifact data to the local disk so that it can
+        # be persisted by awx after the process exits
+        custom_artifact_data = stats.custom.get('_run', {}) if hasattr(stats, 'custom') else {}
+        if custom_artifact_data:
+            # create the directory for custom stats artifacts to live in (if it doesn't exist)
+            custom_artifacts_dir = os.path.join(os.getenv('AWX_PRIVATE_DATA_DIR'), 'artifacts')
+            os.makedirs(custom_artifacts_dir, mode=stat.S_IXUSR + stat.S_IWUSR + stat.S_IRUSR)
+
+            custom_artifacts_path = os.path.join(custom_artifacts_dir, 'custom')
+            with codecs.open(custom_artifacts_path, 'w', encoding='utf-8') as f:
+                os.chmod(custom_artifacts_path, stat.S_IRUSR | stat.S_IWUSR)
+                json.dump(custom_artifact_data, f)
+
         with self.capture_event_data('playbook_on_stats', **event_data):
             super(BaseCallbackModule, self).v2_playbook_on_stats(stats)
 

awx/lib/tests/test_display_callback.py

@@ -2,7 +2,9 @@ from collections import OrderedDict
 import json
 import mock
 import os
+import shutil
 import sys
+import tempfile
 
 import pytest
 
@@ -254,3 +256,26 @@ def test_callback_plugin_strips_task_environ_variables(executor, cache, playbook
     assert len(cache)
     for event in cache.values():
         assert os.environ['PATH'] not in json.dumps(event)
+
+
+@pytest.mark.parametrize('playbook', [
+{'custom_set_stat.yml': '''
+- name: custom set_stat calls should persist to the local disk so awx can save them
+  connection: local
+  hosts: all
+  tasks:
+    - set_stats:
+        data:
+          foo: "bar"
+'''},  # noqa
+])
+def test_callback_plugin_saves_custom_stats(executor, cache, playbook):
+    try:
+        private_data_dir = tempfile.mkdtemp()
+        with mock.patch.dict(os.environ, {'AWX_PRIVATE_DATA_DIR': private_data_dir}):
+            executor.run()
+            artifacts_path = os.path.join(private_data_dir, 'artifacts', 'custom')
+            with open(artifacts_path, 'r') as f:
+                assert json.load(f) == {'foo': 'bar'}
+    finally:
+        shutil.rmtree(os.path.join(private_data_dir))

awx/main/tasks.py

@@ -659,6 +659,7 @@ class BaseTask(LogErrorsTask):
         # Derived class should call add_ansible_venv() or add_awx_venv()
         if self.should_use_proot(instance, **kwargs):
             env['PROOT_TMP_DIR'] = settings.AWX_PROOT_BASE_PATH
+        env['AWX_PRIVATE_DATA_DIR'] = kwargs['private_data_dir']
         return env
 
     def build_safe_env(self, env, **kwargs):
@@ -1307,6 +1308,21 @@ class RunJob(BaseTask):
                 kwargs['private_data_dir'],
                 kwargs['fact_modification_times']
             )
+
+        # persist artifacts set via `set_stat` (if any)
+        custom_stats_path = os.path.join(kwargs['private_data_dir'], 'artifacts', 'custom')
+        if os.path.exists(custom_stats_path):
+            with open(custom_stats_path, 'r') as f:
+                custom_stat_data = None
+                try:
+                    custom_stat_data = json.load(f)
+                except ValueError:
+                    logger.warning('Could not parse custom `set_fact` data for job {}'.format(job.id))
+
+                if custom_stat_data:
+                    job.artifacts = custom_stat_data
+                    job.save(update_fields=['artifacts'])
+
         try:
             inventory = job.inventory
         except Inventory.DoesNotExist: