Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Merge remote-tracking branch 'origin/dmw'
- Context serialization fix
- #370: functioning reboot module.
  • Loading branch information
dw committed Nov 1, 2018
2 parents 6bf4d79 + 7fd4549 commit 484aa44
Show file tree
Hide file tree
Showing 6 changed files with 125 additions and 27 deletions.
44 changes: 27 additions & 17 deletions ansible_mitogen/services.py
Expand Up @@ -268,6 +268,23 @@ def _update_lru(self, new_context, spec, via):
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):
"""
Expand All @@ -280,23 +297,19 @@ def shutdown_all(self):
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()

Expand Down Expand Up @@ -362,13 +375,10 @@ def _connect(self, key, spec, via=None):
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(
Expand Down
15 changes: 13 additions & 2 deletions docs/changelog.rst
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
18 changes: 10 additions & 8 deletions mitogen/core.py
Expand Up @@ -553,7 +553,7 @@ def __init__(self, **kwargs):
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)
Expand Down Expand Up @@ -1498,14 +1498,16 @@ def __repr__(self):
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):
Expand Down
1 change: 1 addition & 0 deletions 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
46 changes: 46 additions & 0 deletions 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
28 changes: 28 additions & 0 deletions tests/serialization_test.py
Expand Up @@ -6,11 +6,14 @@
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)
Expand All @@ -33,5 +36,30 @@ def test_empty_bytes(self):
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()

0 comments on commit 484aa44

Please sign in to comment.