Merge remote-tracking branch 'origin/dmw'

- Context serialization fix
- #370: functioning reboot module.

ansible_mitogen/services.py

@@ -268,6 +268,23 @@ class ContextService(mitogen.service.Service):
         finally:
             self._lock.release()
 
+    @mitogen.service.expose(mitogen.service.AllowParents())
+    def dump(self):
+        """
+        For testing, return a list of dicts describing every currently
+        connected context.
+        """
+        return [
+            {
+                'context_name': context.name,
+                'via': getattr(self._via_by_context.get(context),
+                               'name', None),
+                'refs': self._refs_by_context.get(context),
+            }
+            for context, key in sorted(self._key_by_context.items(),
+                                       key=lambda c_k: c_k[0].context_id)
+        ]
+
     @mitogen.service.expose(mitogen.service.AllowParents())
     def shutdown_all(self):
         """
@@ -280,23 +297,19 @@ class ContextService(mitogen.service.Service):
         finally:
             self._lock.release()
 
-    def _on_stream_disconnect(self, stream):
+    def _on_context_disconnect(self, context):
         """
-        Respond to Stream disconnection by deleting any record of contexts
-        reached via that stream. This method runs in the Broker thread and must
-        not to block.
+        Respond to Context disconnect event by deleting any record of the no
+        longer reachable context.  This method runs in the Broker thread and
+        must not to block.
         """
         # TODO: there is a race between creation of a context and disconnection
         # of its related stream. An error reply should be sent to any message
         # in _latches_by_key below.
         self._lock.acquire()
         try:
-            routes = self.router.route_monitor.get_routes(stream)
-            for context in list(self._key_by_context):
-                if context.context_id in routes:
-                    LOG.info('Dropping %r due to disconnect of %r',
-                             context, stream)
-                    self._forget_context_unlocked(context)
+            LOG.info('Forgetting %r due to stream disconnect', context)
+            self._forget_context_unlocked(context)
         finally:
             self._lock.release()
 
@@ -362,13 +375,10 @@ class ContextService(mitogen.service.Service):
         context = method(via=via, unidirectional=True, **spec['kwargs'])
         if via and spec.get('enable_lru'):
             self._update_lru(context, spec, via)
-        else:
-            # For directly connected contexts, listen to the associated
-            # Stream's disconnect event and use it to invalidate dependent
-            # Contexts.
-            stream = self.router.stream_by_id(context.context_id)
-            mitogen.core.listen(stream, 'disconnect',
-                                lambda: self._on_stream_disconnect(stream))
+
+        # Forget the context when its disconnect event fires.
+        mitogen.core.listen(context, 'disconnect',
+            lambda: self._on_context_disconnect(context))
 
         self._send_module_forwards(context)
         init_child_result = context.call(

docs/changelog.rst

@@ -52,6 +52,10 @@ Fixes
   was invoked using sudo without appropriate flags to cause the ``HOME``
   environment variable to be reset to match the target account.
 
+* `#370 <https://github.com/dw/mitogen/issues/370>`_: the Ansible
+  `reboot <https://docs.ansible.com/ansible/latest/modules/reboot_module.html>`_
+  module is supported.
+
 * `#373 <https://github.com/dw/mitogen/issues/373>`_: the LXC and LXD methods
   now print a useful hint when Python fails to start, as no useful error is
   normally logged to the console by these tools.
@@ -110,8 +114,9 @@ bug reports, features and fixes in this release contributed by
 `Brian Candler <https://github.com/candlerb>`_,
 `Guy Knights <https://github.com/knightsg>`_,
 `Jiří Vávra <https://github.com/Houbovo>`_,
-`Jonathan Rosser <https://github.com/jrosser>`_, and
-`Mehdi <https://github.com/mehdisat7>`_.
+`Jonathan Rosser <https://github.com/jrosser>`_,
+`Mehdi <https://github.com/mehdisat7>`_, and
+`Mohammed Naser <https://github.com/mnaser/>`_.
 
 
 v0.2.3 (2018-10-23)
@@ -499,6 +504,12 @@ Mitogen for Ansible
 
 **Known Issues**
 
+* On OS X when a SSH password is specified and the default connection type of
+  ``smart`` is used, Ansible may select the Paramiko plug-in rather than
+  Mitogen. If you specify a password on OS X, ensure ``connection: ssh``
+  appears in your playbook, ``ansible.cfg``, or as ``-c ssh`` on the
+  command-line.
+
 * The ``raw`` action executes as a regular Mitogen connection, which requires
   Python on the target, precluding its use for installing Python. This will be
   addressed in a future 0.2 release. For now, simply mix Mitogen and vanilla

mitogen/core.py

@@ -553,7 +553,7 @@ class Message(object):
         assert isinstance(self.data, BytesType)
 
     def _unpickle_context(self, context_id, name):
-        return _unpickle_context(self.router, context_id, name)
+        return _unpickle_context(context_id, name, router=self.router)
 
     def _unpickle_sender(self, context_id, dst_handle):
         return _unpickle_sender(self.router, context_id, dst_handle)
@@ -1498,14 +1498,16 @@ class Context(object):
         return 'Context(%s, %r)' % (self.context_id, self.name)
 
 
-def _unpickle_context(router, context_id, name):
-    if not (isinstance(router, Router) and
-            isinstance(context_id, (int, long)) and context_id >= 0 and (
-                (name is None) or
-                (isinstance(name, UnicodeType) and len(name) < 100))
-            ):
+def _unpickle_context(context_id, name, router=None):
+    if not (isinstance(context_id, (int, long)) and context_id >= 0 and (
+        (name is None) or
+        (isinstance(name, UnicodeType) and len(name) < 100))
+    ):
         raise TypeError('cannot unpickle Context: bad input')
-    return router.context_by_id(context_id, name=name)
+
+    if isinstance(router, Router):
+        return router.context_by_id(context_id, name=name)
+    return Context(None, context_id, name)  # For plain Jane pickle.
 
 
 class Poller(object):

tests/ansible/integration/context_service/all.yml

@@ -1,2 +1,3 @@
+- import_playbook: disconnect_cleanup.yml
 - import_playbook: lru_one_target.yml
 - import_playbook: reconnection.yml

tests/ansible/integration/context_service/disconnect_cleanup.yml

@@ -0,0 +1,46 @@
+# issue #76, #370: ensure context state is forgotten on disconnect, including
+# state of dependent contexts (e.g. sudo, connection delegation, ..).
+
+- name: integration/context_service/disconnect_cleanup.yml
+  hosts: test-targets
+  any_errors_fatal: true
+  tasks:
+  - meta: end_play
+    when: not is_mitogen
+
+  # Start with a clean slate.
+  - mitogen_shutdown_all:
+
+  # Connect a few users.
+  - shell: "true"
+    become: true
+    become_user: "mitogen__user{{item}}"
+    with_items: [1, 2, 3]
+
+  # Verify current state.
+  - mitogen_action_script:
+      script: |
+        self._connection._connect()
+        result['dump'] = self._connection.parent.call_service(
+          service_name='ansible_mitogen.services.ContextService',
+          method_name='dump'
+        )
+    register: out
+
+  - assert:
+      that: out.dump|length == 4  # ssh account + 3 sudo accounts
+
+  - meta: reset_connection
+
+  # Verify current state.
+  - mitogen_action_script:
+      script: |
+        self._connection._connect()
+        result['dump'] = self._connection.parent.call_service(
+          service_name='ansible_mitogen.services.ContextService',
+          method_name='dump'
+        )
+    register: out
+
+  - assert:
+      that: out.dump|length == 1  # just the ssh account

tests/serialization_test.py

@@ -6,11 +6,14 @@ except ImportError:
     from StringIO import StringIO as StringIO
     from StringIO import StringIO as BytesIO
 
+import pickle
 import unittest2
 
 import mitogen.core
 from mitogen.core import b
 
+import testlib
+
 
 def roundtrip(v):
     msg = mitogen.core.Message.pickled(v)
@@ -33,5 +36,30 @@ class BlobTest(unittest2.TestCase):
         self.assertEquals(b(''), roundtrip(v))
 
 
+class ContextTest(testlib.RouterMixin, unittest2.TestCase):
+    klass = mitogen.core.Context
+
+    # Ensure Context can be round-tripped by regular pickle in addition to
+    # Mitogen's hacked pickle. Users may try to call pickle on a Context in
+    # strange circumstances, and it's often used to glue pieces of an app
+    # together (e.g. Ansible).
+
+    def test_mitogen_roundtrip(self):
+        c = self.router.fork()
+        r = mitogen.core.Receiver(self.router)
+        r.to_sender().send(c)
+        c2 = r.get().unpickle()
+        self.assertEquals(None, c2.router)
+        self.assertEquals(c.context_id, c2.context_id)
+        self.assertEquals(c.name, c2.name)
+
+    def test_vanilla_roundtrip(self):
+        c = self.router.fork()
+        c2 = pickle.loads(pickle.dumps(c))
+        self.assertEquals(None, c2.router)
+        self.assertEquals(c.context_id, c2.context_id)
+        self.assertEquals(c.name, c2.name)
+
+
 if __name__ == '__main__':
     unittest2.main()