issue #297: local commands must execute with WorkerProcess environment.

ansible_mitogen/connection.py

@@ -550,7 +550,14 @@ class Connection(ansible.plugins.connection.ConnectionBase):
     def get_default_cwd(self):
         """
         Overridden by connections/mitogen_local.py to emulate behaviour of CWD
-        inherited from WorkerProcess.
+        being fixed to that of ActionBase._loader.get_basedir().
+        """
+        return None
+
+    def get_default_env(self):
+        """
+        Overridden by connections/mitogen_local.py to emulate behaviour of
+        WorkProcess environment inherited from WorkerProcess.
         """
         return None
 

ansible_mitogen/planner.py

@@ -176,6 +176,7 @@ class Planner(object):
         new = dict((mitogen.core.UnicodeType(k), kwargs[k])
                    for k in kwargs)
         new.setdefault('cwd', self._inv.connection.get_default_cwd())
+        new.setdefault('extra_env', self._inv.connection.get_default_env())
         new.setdefault('emulate_tty', True)
         new.setdefault('service_context', self._inv.connection.parent)
         return new

ansible_mitogen/plugins/connection/mitogen_local.py

@@ -37,6 +37,32 @@ except ImportError:
     del base_dir
 
 import ansible_mitogen.connection
+import ansible_mitogen.process
+
+
+if sys.version_info > (2, 7):
+    viewkeys = dict.viewkeys
+else:
+    viewkeys = lambda dct: set(dct)
+
+
+def dict_diff(old, new):
+    """
+    Return a dict representing the differences between the dicts `old` and
+    `new`. Deleted keys appear as a key with the value :data:`None`, added and
+    changed keys appear as a key with the new value.
+    """
+    old_keys = viewkeys(old)
+    new_keys = viewkeys(dict(new))
+    out = {}
+    for key in new_keys - old_keys:
+        out[key] = new[key]
+    for key in old_keys - new_keys:
+        out[key] = None
+    for key in old_keys & new_keys:
+        if old[key] != new[key]:
+            out[key] = new[key]
+    return out
 
 
 class Connection(ansible_mitogen.connection.Connection):
@@ -45,3 +71,13 @@ class Connection(ansible_mitogen.connection.Connection):
     def get_default_cwd(self):
         # https://github.com/ansible/ansible/issues/14489
         return self.loader_basedir
+
+    def get_default_env(self):
+        """
+        Vanilla Ansible local commands execute with an environment inherited
+        from WorkerProcess, we must emulate that.
+        """
+        return dict_diff(
+            old=ansible_mitogen.process.MuxProcess.original_env,
+            new=os.environ,
+        )

ansible_mitogen/process.py

@@ -84,6 +84,12 @@ class MuxProcess(object):
     #: that was spawned.
     worker_pid = None
 
+    #: A copy of :data:`os.environ` at the time the multiplexer process was
+    #: started. It's used by mitogen_local.py to find changes made to the
+    #: top-level environment (e.g. vars plugins -- issue #297) that must be
+    #: applied to locally executed commands and modules.
+    original_env = None
+
     #: In both processes, this is the temporary UNIX socket used for
     #: forked WorkerProcesses to contact the MuxProcess
     unix_listener_path = None
@@ -108,6 +114,7 @@ class MuxProcess(object):
         mitogen.core.set_cloexec(cls.worker_sock.fileno())
         mitogen.core.set_cloexec(cls.child_sock.fileno())
 
+        cls.original_env = dict(os.environ)
         cls.child_pid = os.fork()
         ansible_mitogen.logging.setup()
         if cls.child_pid:

ansible_mitogen/runner.py

@@ -137,13 +137,14 @@ class Runner(object):
         When :data:`True`, indicate the runner should detach the context from
         its parent after setup has completed successfully.
     """
-    def __init__(self, module, service_context, json_args, cwd=None, env=None,
+    def __init__(self, module, service_context, json_args, extra_env=None,
-                 econtext=None, detach=False):
+                 cwd=None, env=None, econtext=None, detach=False):
         self.module = module
         self.service_context = service_context
         self.econtext = econtext
         self.detach = detach
         self.args = json.loads(json_args)
+        self.extra_env = extra_env
         self.env = env
         self.cwd = cwd
 
@@ -154,7 +155,10 @@ class Runner(object):
         execution. The base implementation simply prepares the environment.
         """
         self._cwd = TemporaryCwd(self.cwd)
-        self._env = TemporaryEnvironment(self.env)
+        env = dict(self.extra_env or {})
+        if self.env:
+            env.update(self.env)
+        self._env = TemporaryEnvironment(env)
 
     def revert(self):
         """
@@ -275,12 +279,22 @@ class TemporaryCwd(object):
 
 
 class TemporaryEnvironment(object):
+    """
+    Apply environment changes from `env` until :meth:`revert` is called. Values
+    in the dict may be :data:`None` to indicate the relevant key should be
+    deleted.
+    """
     def __init__(self, env=None):
-        self.original = os.environ.copy()
+        self.original = dict(os.environ)
         self.env = env or {}
-        os.environ.update((k, str(v)) for k, v in iteritems(self.env))
+        for key, value in iteritems(self.env):
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = str(value)
 
     def revert(self):
+        if self.env:
             os.environ.clear()
             os.environ.update(self.original)
 

tests/ansible/ansible.cfg

@@ -5,6 +5,7 @@ strategy_plugins = ../../ansible_mitogen/plugins/strategy
 action_plugins = lib/action
 callback_plugins = lib/callback
 stdout_callback = nice_stdout
+vars_plugins = lib/vars
 library = lib/modules
 # module_utils = lib/module_utils
 retry_files_enabled = False

tests/ansible/integration/local/all.yml

@@ -1,4 +1,4 @@
 
 - import_playbook: cwd_preserved.yml
-#- import_playbook: env_preserved.yml
+- import_playbook: env_preserved.yml
 

tests/ansible/integration/local/env_preserved.yml

@@ -0,0 +1,8 @@
+
+# Execution environment should be that of WorkerProcess --
+# https://github.com/dw/mitogen/issues/297
+
+- hosts: localhost
+  connection: local
+  tasks:
+    - shell: "env | grep EVIL_VARS_PLUGIN"

tests/ansible/lib/vars/custom_modifies_os_environ.py

@@ -0,0 +1,16 @@
+# https://github.com/dw/mitogen/issues/297
+
+from __future__ import (absolute_import, division, print_function)
+__metaclass__ = type
+
+from ansible.plugins.vars import BaseVarsPlugin
+import os
+
+class VarsModule(BaseVarsPlugin):
+    def __init__(self, *args):
+        super(VarsModule, self).__init__(*args)
+        os.environ['EVIL_VARS_PLUGIN'] = 'YIPEEE'
+
+    def get_vars(self, loader, path, entities, cache=True):
+        super(VarsModule, self).get_vars(loader, path, entities)
+        return {}