From 66142e7d75f76b86d5022c26150e1c2df81db97a Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 8 Sep 2018 22:17:39 +0100
Subject: [PATCH] ansible: fork isolated tasks from correct parent.

Closes #355.
---
 ansible_mitogen/connection.py                 |  5 +++-
 tests/ansible/integration/runner/all.yml      |  4 ++-
 ...rking_behaviour.yml => forking_active.yml} | 21 ++-------------
 .../runner/forking_correct_parent.yml         | 26 +++++++++++++++++++
 .../integration/runner/forking_inactive.yml   | 23 ++++++++++++++++
 5 files changed, 58 insertions(+), 21 deletions(-)
 rename tests/ansible/integration/runner/{forking_behaviour.yml => forking_active.yml} (59%)
 create mode 100644 tests/ansible/integration/runner/forking_correct_parent.yml
 create mode 100644 tests/ansible/integration/runner/forking_inactive.yml

diff --git a/ansible_mitogen/connection.py b/ansible_mitogen/connection.py
index fd81b99c..21ecb5f0 100644
--- a/ansible_mitogen/connection.py
+++ b/ansible_mitogen/connection.py
@@ -693,6 +693,8 @@ class Connection(ansible.plugins.connection.ConnectionBase):
 
         if kwargs.pop('use_login_context', None):
             call_context = self.login_context
+        elif kwargs.pop('use_fork_context', None):
+            call_context = self.fork_context
         else:
             call_context = self.context
 
@@ -743,7 +745,8 @@ class Connection(ansible.plugins.connection.ConnectionBase):
         :returns:
             mitogen.core.Context of the new child.
         """
-        return self.call(ansible_mitogen.target.create_fork_child)
+        return self.call(ansible_mitogen.target.create_fork_child,
+                         use_fork_context=True)
 
     def get_default_cwd(self):
         """
diff --git a/tests/ansible/integration/runner/all.yml b/tests/ansible/integration/runner/all.yml
index ffb263fb..69c22edb 100644
--- a/tests/ansible/integration/runner/all.yml
+++ b/tests/ansible/integration/runner/all.yml
@@ -14,5 +14,7 @@
 - import_playbook: custom_script_interpreter.yml
 - import_playbook: environment_isolation.yml
 - import_playbook: etc_environment.yml
-- import_playbook: forking_behaviour.yml
+- import_playbook: forking_active.yml
+- import_playbook: forking_inactive.yml
+- import_playbook: forking_correct_parent.yml
 - import_playbook: missing_module.yml
diff --git a/tests/ansible/integration/runner/forking_behaviour.yml b/tests/ansible/integration/runner/forking_active.yml
similarity index 59%
rename from tests/ansible/integration/runner/forking_behaviour.yml
rename to tests/ansible/integration/runner/forking_active.yml
index 7268fce0..15e6d2bc 100644
--- a/tests/ansible/integration/runner/forking_behaviour.yml
+++ b/tests/ansible/integration/runner/forking_active.yml
@@ -1,26 +1,8 @@
-
-- name: integration/runner/forking_behaviour.yml
+- name: integration/runner/forking_active.yml
   hosts: test-targets
   any_errors_fatal: true
   tasks:
 
-  # Verify non-async jobs run in-process.
-
-  - name: get process ID.
-    custom_python_detect_environment:
-    register: sync_proc1
-    when: is_mitogen
-
-  - name: get process ID again.
-    custom_python_detect_environment:
-    register: sync_proc2
-    when: is_mitogen
-
-  - assert:
-      that:
-      - sync_proc1.pid == sync_proc2.pid
-    when: is_mitogen
-
   # Verify mitogen_task_isolation=fork triggers forking.
 
   - name: get force-forked process ID.
@@ -42,3 +24,4 @@
       - fork_proc1.pid != sync_proc1.pid
       - fork_proc1.pid != fork_proc2.pid
     when: is_mitogen
+
diff --git a/tests/ansible/integration/runner/forking_correct_parent.yml b/tests/ansible/integration/runner/forking_correct_parent.yml
new file mode 100644
index 00000000..e8207676
--- /dev/null
+++ b/tests/ansible/integration/runner/forking_correct_parent.yml
@@ -0,0 +1,26 @@
+
+- name: integration/runner/forking_correct_parent.yml
+  hosts: test-targets
+  any_errors_fatal: true
+  tasks:
+
+  # Verify mitogen_task_isolation=fork forks from "virginal fork parent", not
+  # shared interpreter.
+
+  - name: get regular process ID.
+    custom_python_detect_environment:
+    register: regular_proc
+    when: is_mitogen
+
+  - name: get force-forked process ID again.
+    custom_python_detect_environment:
+    register: fork_proc
+    vars:
+      mitogen_task_isolation: fork
+    when: is_mitogen
+
+  - assert:
+      that:
+      - fork_proc.pid != regular_proc.pid
+      - fork_proc.ppid != regular_proc.pid
+    when: is_mitogen
diff --git a/tests/ansible/integration/runner/forking_inactive.yml b/tests/ansible/integration/runner/forking_inactive.yml
new file mode 100644
index 00000000..b84cec7e
--- /dev/null
+++ b/tests/ansible/integration/runner/forking_inactive.yml
@@ -0,0 +1,23 @@
+# Verify non-async jobs run in-process.
+
+- name: integration/runner/forking_inactive.yml
+  hosts: test-targets
+  any_errors_fatal: true
+  tasks:
+
+  - name: get process ID.
+    custom_python_detect_environment:
+    register: sync_proc1
+    when: is_mitogen
+
+  - name: get process ID again.
+    custom_python_detect_environment:
+    register: sync_proc2
+    when: is_mitogen
+
+  - assert:
+      that:
+      - sync_proc1.pid == sync_proc2.pid
+    when: is_mitogen
+
+