mitogen/ansible_mitogen/services.py

# Copyright 2017, David Wilson
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Classes in this file define Mitogen 'services' that run (initially) within the
connection multiplexer process that is forked off the top-level controller
process.

Once a worker process connects to a multiplexer process
(Connection._connect()), it communicates with these services to establish new
connections, grant access to files by children, and register for notification
when a child has completed a job.
"""

from __future__ import absolute_import
import logging
import os
import os.path
import pprint
import sys
import threading
import zlib

import mitogen
import mitogen.service
import ansible_mitogen.target


LOG = logging.getLogger(__name__)


class Error(Exception):
    pass


class ContextService(mitogen.service.Service):
    """
    Used by workers to fetch the single Context instance corresponding to a
    connection configuration, creating the matching connection if it does not
    exist.

    For connection methods and their parameters, see:
        https://mitogen.readthedocs.io/en/latest/api.html#context-factories

    This concentrates connections in the top-level process, which may become a
    bottleneck. The bottleneck can be removed using per-CPU connection
    processes and arranging for the worker to select one according to a hash of
    the connection parameters (sharding).
    """
    handle = 500
    max_message_size = 1000
    max_interpreters = int(os.getenv('MITOGEN_MAX_INTERPRETERS', '20'))

    def __init__(self, *args, **kwargs):
        super(ContextService, self).__init__(*args, **kwargs)
        self._lock = threading.Lock()
        #: Records the :meth:`get` result dict for successful calls, returned
        #: for identical subsequent calls. Keyed by :meth:`key_from_kwargs`.
        self._response_by_key = {}
        #: List of :class:`mitogen.core.Message` waiting for the result dict
        #: for a particular connection config. Keyed as sbove.
        self._waiters_by_key = {}
        #: Mapping of :class:`mitogen.core.Context` -> reference count. Each
        #: call to :meth:`get` increases this by one. Calls to :meth:`put`
        #: decrease it by one.
        self._refs_by_context = {}
        #: List of contexts in creation order by via= parameter. When
        #: :attr:`max_interpreters` is reached, the most recently used context
        #: is destroyed to make room for any additional context.
        self._lru_by_via = {}
        #: :meth:`key_from_kwargs` result by Context.
        self._key_by_context = {}

    @mitogen.service.expose(mitogen.service.AllowParents())
    @mitogen.service.arg_spec({
        'context': mitogen.core.Context
    })
    def put(self, context):
        """
        Return a reference, making it eligable for recycling once its reference
        count reaches zero.
        """
        LOG.debug('%r.put(%r)', self, context)
        if self._refs_by_context.get(context, 0) == 0:
            LOG.warning('%r.put(%r): refcount was 0. shutdown_all called?',
                        self, context)
            return
        self._refs_by_context[context] -= 1

    def key_from_kwargs(self, **kwargs):
        """
        Generate a deduplication key from the request. The default
        implementation returns a string based on a stable representation of the
        input dictionary generated by :py:func:`pprint.pformat`.
        """
        return pprint.pformat(kwargs)

    def _produce_response(self, key, response):
        """
        Reply to every waiting request matching a configuration key with a
        response dictionary, deleting the list of waiters when done.

        :param str key:
            Result of :meth:`key_from_kwargs`
        :param dict response:
            Response dictionary
        :returns:
            Number of waiters that were replied to.
        """
        self._lock.acquire()
        try:
            waiters = self._waiters_by_key.pop(key)
            count = len(waiters)
            for msg in waiters:
                msg.reply(response)
        finally:
            self._lock.release()
        return count

    def _shutdown(self, context, lru=None, new_context=None):
        """
        Arrange for `context` to be shut down, and optionally add `new_context`
        to the LRU list while holding the lock.
        """
        LOG.info('%r._shutdown(): shutting down %r', self, context)
        context.shutdown()

        key = self._key_by_context[context]

        self._lock.acquire()
        try:
            del self._response_by_key[key]
            del self._refs_by_context[context]
            del self._key_by_context[context]
            if lru:
                lru.remove(context)
            if new_context:
                lru.append(new_context)
        finally:
            self._lock.release()

    def _update_lru(self, new_context, **kwargs):
        """
        Update the LRU ("MRU"?) list associated with the connection described
        by `kwargs`, destroying the most recently created context if the list
        is full. Finally add `new_context` to the list.
        """
        via = kwargs.get('via')
        if via is None:
            # We don't have a limit on the number of directly connections.
            return

        lru = self._lru_by_via.setdefault(via, [])
        if len(lru) < self.max_interpreters:
            lru.append(new_context)
            return

        for context in reversed(lru):
            if self._refs_by_context[context] == 0:
                break
        else:
            LOG.warning('via=%r reached maximum number of interpreters, '
                        'but they are all marked as in-use.', via)
            return

        self._shutdown(context, lru=lru, new_context=new_context)

    @mitogen.service.expose(mitogen.service.AllowParents())
    def shutdown_all(self):
        """
        For testing use, arrange for all connections to be shut down.
        """
        for context in list(self._key_by_context):
            self._shutdown(context)
        self._lru_by_via = {}

    def _on_stream_disconnect(self, stream):
        """
        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.
        """
        # 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 _waiters_by_key below.
        self._lock.acquire()
        try:
            for context, key in list(self._key_by_context.items()):
                if context.context_id in stream.routes:
                    LOG.info('Dropping %r due to disconnect of %r',
                             context, stream)
                    self._response_by_key.pop(key, None)
                    self._waiters_by_key.pop(key, None)
                    self._refs_by_context.pop(context, None)
                    self._lru_by_via.pop(context, None)
                    self._refs_by_context.pop(context, None)
        finally:
            self._lock.release()

    def _connect(self, key, method_name, **kwargs):
        """
        Actual connect implementation. Arranges for the Mitogen connection to
        be created and enqueues an asynchronous call to start the forked task
        parent in the remote context.

        :param key:
            Deduplication key representing the connection configuration.
        :param method_name:
            :class:`mitogen.parent.Router` method implementing the connection
            type.
        :param kwargs:
            Keyword arguments passed to the router method.
        :returns:
            Dict like::

                {
                    'context': mitogen.core.Context or None,
                    'home_dir': str or None,
                    'msg': str or None
                }

            Where either `msg` is an error message and the remaining fields are
            :data:`None`, or `msg` is :data:`None` and the remaining fields are
            set.
        """
        method = getattr(self.router, method_name, None)
        if method is None:
            raise Error('no such Router method: %s' % (method_name,))

        try:
            context = method(**kwargs)
        except mitogen.core.StreamError as e:
            return {
                'context': None,
                'home_dir': None,
                'msg': str(e),
            }

        if kwargs.get('via'):
            self._update_lru(context, method_name=method_name, **kwargs)
        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))

        home_dir = context.call(os.path.expanduser, '~')

        # We don't need to wait for the result of this. Ideally we'd check its
        # return value somewhere, but logs will catch a failure anyway.
        context.call_async(ansible_mitogen.target.start_fork_parent)

        if os.environ.get('MITOGEN_DUMP_THREAD_STACKS'):
            from mitogen import debug
            context.call(debug.dump_to_logger)

        self._key_by_context[context] = key
        self._refs_by_context[context] = 0
        return {
            'context': context,
            'home_dir': home_dir,
            'msg': None,
        }

    @mitogen.service.expose(mitogen.service.AllowParents())
    @mitogen.service.arg_spec({
        'method_name': str
    })
    def get(self, msg, **kwargs):
        """
        Return a Context referring to an established connection with the given
        configuration, establishing a new connection as necessary.

        :param str method_name:
            The :class:`mitogen.parent.Router` connection method to use.
        :param dict kwargs:
            Keyword arguments passed to `mitogen.master.Router.[method_name]()`.

        :returns tuple:
            Tuple of `(context, home_dir)`, where:
                * `context` is the mitogen.master.Context referring to the
                  target context.
                * `home_dir` is a cached copy of the remote directory.
        """
        key = self.key_from_kwargs(**kwargs)
        self._lock.acquire()
        try:
            response = self._response_by_key.get(key)
            if response is not None:
                self._refs_by_context[response['context']] += 1
                return response

            waiters = self._waiters_by_key.get(key)
            if waiters is not None:
                waiters.append(msg)
                return self.NO_REPLY

            self._waiters_by_key[key] = [msg]
        finally:
            self._lock.release()

        # I'm the first thread to wait, so I will create the connection.
        try:
            response = self._connect(key, **kwargs)
            count = self._produce_response(key, response)
            if response['msg'] is None:
                # Only record the response for non-error results.
                self._response_by_key[key] = response
                # Set the reference count to the number of waiters.
                self._refs_by_context[response['context']] += count
        except mitogen.core.CallError:
            e = sys.exc_info()[1]
            self._produce_response(key, e)
        except Exception:
            e = sys.exc_info()[1]
            self._produce_response(key, mitogen.core.CallError(e))

        return self.NO_REPLY


class FileService(mitogen.service.Service):
    """
    Primitive latency-inducing file server for old-style incantations of the
    module runner. This is to be replaced later with a scheme that forwards
    files known to be missing without the target having to ask for them,
    avoiding a corresponding roundtrip per file.

    Paths must be explicitly added to the service by a trusted context before
    they will be served to an untrusted context.
    """
    handle = 501
    max_message_size = 1000
    unregistered_msg = 'Path is not registered with FileService.'

    def __init__(self, router):
        super(FileService, self).__init__(router)
        self._paths = {}

    @mitogen.service.expose(policy=mitogen.service.AllowParents())
    @mitogen.service.arg_spec({
        'path': basestring
    })
    def register(self, path):
        """
        Authorize a path for access by child contexts. Calling this repeatedly
        with the same path is harmless.

        :param str path:
            File path.
        """
        if path not in self._paths:
            LOG.debug('%r: registering %r', self, path)
            with open(path, 'rb') as fp:
                self._paths[path] = zlib.compress(fp.read())

    @mitogen.service.expose(policy=mitogen.service.AllowAny())
    @mitogen.service.arg_spec({
        'path': basestring
    })
    def fetch(self, path):
        """
        Fetch a file's data.

        :param str path:
            File path.

        :returns:
            The file data.

        :raises mitogen.core.CallError:
            The path was not registered.
        """
        if path not in self._paths:
            raise mitogen.core.CallError(self.unregistered_msg)

        LOG.debug('Serving %r', path)
        return self._paths[path]