237 lines
7.5 KiB
Python
237 lines
7.5 KiB
Python
|
# Copyright 2019, 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.
|
||
|
|
|
||
|
|
# !mitogen: minify_safe
|
||
|
|
|
||
|
|
import datetime
|
||
|
|
import functools
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import mitogen.core
|
||
|
|
import mitogen.master
|
||
|
|
|
||
|
|
|
||
|
|
iteritems = getattr(dict, 'iteritems', dict.items)
|
||
|
|
|
||
|
|
if mitogen.core.PY3:
|
||
|
|
iteritems = dict.items
|
||
|
|
else:
|
||
|
|
iteritems = dict.iteritems
|
||
|
|
|
||
|
|
|
||
|
|
def setup_gil():
|
||
|
|
"""
|
||
|
|
Set extremely long GIL release interval to let threads naturally progress
|
||
|
|
through CPU-heavy sequences without forcing the wake of another thread that
|
||
|
|
may contend trying to run the same CPU-heavy code. For the new-style
|
||
|
|
Ansible work, this drops runtime ~33% and involuntary context switches by
|
||
|
|
>80%, essentially making threads cooperatively scheduled.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
# Python 2.
|
||
|
|
sys.setcheckinterval(100000)
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Python 3.
|
||
|
|
sys.setswitchinterval(10)
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def disable_site_packages():
|
||
|
|
"""
|
||
|
|
Remove all entries mentioning ``site-packages`` or ``Extras`` from
|
||
|
|
:attr:sys.path. Used primarily for testing on OS X within a virtualenv,
|
||
|
|
where OS X bundles some ancient version of the :mod:`six` module.
|
||
|
|
"""
|
||
|
|
for entry in sys.path[:]:
|
||
|
|
if 'site-packages' in entry or 'Extras' in entry:
|
||
|
|
sys.path.remove(entry)
|
||
|
|
|
||
|
|
|
||
|
|
def _formatTime(record, datefmt=None):
|
||
|
|
dt = datetime.datetime.fromtimestamp(record.created)
|
||
|
|
return dt.strftime(datefmt)
|
||
|
|
|
||
|
|
|
||
|
|
def log_get_formatter():
|
||
|
|
datefmt = '%H:%M:%S'
|
||
|
|
if sys.version_info > (2, 6):
|
||
|
|
datefmt += '.%f'
|
||
|
|
fmt = '%(asctime)s %(levelname).1s %(name)s: %(message)s'
|
||
|
|
formatter = logging.Formatter(fmt, datefmt)
|
||
|
|
formatter.formatTime = _formatTime
|
||
|
|
return formatter
|
||
|
|
|
||
|
|
|
||
|
|
def log_to_file(path=None, io=False, level='INFO'):
|
||
|
|
"""
|
||
|
|
Install a new :class:`logging.Handler` writing applications logs to the
|
||
|
|
filesystem. Useful when debugging slave IO problems.
|
||
|
|
|
||
|
|
Parameters to this function may be overridden at runtime using environment
|
||
|
|
variables. See :ref:`logging-env-vars`.
|
||
|
|
|
||
|
|
:param str path:
|
||
|
|
If not :data:`None`, a filesystem path to write logs to. Otherwise,
|
||
|
|
logs are written to :data:`sys.stderr`.
|
||
|
|
|
||
|
|
:param bool io:
|
||
|
|
If :data:`True`, include extremely verbose IO logs in the output.
|
||
|
|
Useful for debugging hangs, less useful for debugging application code.
|
||
|
|
|
||
|
|
:param str level:
|
||
|
|
Name of the :mod:`logging` package constant that is the minimum level
|
||
|
|
to log at. Useful levels are ``DEBUG``, ``INFO``, ``WARNING``, and
|
||
|
|
``ERROR``.
|
||
|
|
"""
|
||
|
|
log = logging.getLogger('')
|
||
|
|
if path:
|
||
|
|
fp = open(path, 'w', 1)
|
||
|
|
mitogen.core.set_cloexec(fp.fileno())
|
||
|
|
else:
|
||
|
|
fp = sys.stderr
|
||
|
|
|
||
|
|
level = os.environ.get('MITOGEN_LOG_LEVEL', level).upper()
|
||
|
|
io = level == 'IO'
|
||
|
|
if io:
|
||
|
|
level = 'DEBUG'
|
||
|
|
logging.getLogger('mitogen.io').setLevel(level)
|
||
|
|
|
||
|
|
level = getattr(logging, level, logging.INFO)
|
||
|
|
log.setLevel(level)
|
||
|
|
|
||
|
|
# Prevent accidental duplicate log_to_file() calls from generating
|
||
|
|
# duplicate output.
|
||
|
|
for handler_ in reversed(log.handlers):
|
||
|
|
if getattr(handler_, 'is_mitogen', None):
|
||
|
|
log.handlers.remove(handler_)
|
||
|
|
|
||
|
|
handler = logging.StreamHandler(fp)
|
||
|
|
handler.is_mitogen = True
|
||
|
|
handler.formatter = log_get_formatter()
|
||
|
|
log.handlers.insert(0, handler)
|
||
|
|
|
||
|
|
|
||
|
|
def run_with_router(func, *args, **kwargs):
|
||
|
|
"""
|
||
|
|
Arrange for `func(router, *args, **kwargs)` to run with a temporary
|
||
|
|
:class:`mitogen.master.Router`, ensuring the Router and Broker are
|
||
|
|
correctly shut down during normal or exceptional return.
|
||
|
|
|
||
|
|
:returns:
|
||
|
|
`func`'s return value.
|
||
|
|
"""
|
||
|
|
broker = mitogen.master.Broker()
|
||
|
|
router = mitogen.master.Router(broker)
|
||
|
|
try:
|
||
|
|
return func(router, *args, **kwargs)
|
||
|
|
finally:
|
||
|
|
broker.shutdown()
|
||
|
|
broker.join()
|
||
|
|
|
||
|
|
|
||
|
|
def with_router(func):
|
||
|
|
"""
|
||
|
|
Decorator version of :func:`run_with_router`. Example:
|
||
|
|
|
||
|
|
.. code-block:: python
|
||
|
|
|
||
|
|
@with_router
|
||
|
|
def do_stuff(router, arg):
|
||
|
|
pass
|
||
|
|
|
||
|
|
do_stuff(blah, 123)
|
||
|
|
"""
|
||
|
|
@functools.wraps(func)
|
||
|
|
def wrapper(*args, **kwargs):
|
||
|
|
return run_with_router(func, *args, **kwargs)
|
||
|
|
return wrapper
|
||
|
|
|
||
|
|
|
||
|
|
PASSTHROUGH = (
|
||
|
|
int, float, bool,
|
||
|
|
type(None),
|
||
|
|
mitogen.core.Context,
|
||
|
|
mitogen.core.CallError,
|
||
|
|
mitogen.core.Blob,
|
||
|
|
mitogen.core.Secret,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def cast(obj):
|
||
|
|
"""
|
||
|
|
Return obj (or a copy) with subtypes of builtins cast to their supertype.
|
||
|
|
Subtypes of those in :data:`PASSTHROUGH` are not modified.
|
||
|
|
|
||
|
|
Many tools love to subclass built-in types in order to implement useful
|
||
|
|
functionality, such as annotating the safety of a Unicode string, or adding
|
||
|
|
additional methods to a dict. However :py:mod:`pickle` serializes these
|
||
|
|
exactly, leading to :exc:`mitogen.CallError` during :meth:`Context.call
|
||
|
|
<mitogen.parent.Context.call>` in the target when it tries to deserialize
|
||
|
|
the data.
|
||
|
|
|
||
|
|
This function walks the object graph `obj`, producing a copy with any
|
||
|
|
custom sub-types removed. The functionality is not default since the
|
||
|
|
resulting walk may be computationally expensive given a large enough graph.
|
||
|
|
|
||
|
|
Raises :py:exc:`TypeError` if an unknown subtype is encountered, or
|
||
|
|
casting does not return the desired supertype.
|
||
|
|
|
||
|
|
See :ref:`serialization-rules` for a list of supported types.
|
||
|
|
|
||
|
|
:param obj:
|
||
|
|
Object to undecorate.
|
||
|
|
:returns:
|
||
|
|
Undecorated object.
|
||
|
|
"""
|
||
|
|
if isinstance(obj, dict):
|
||
|
|
return dict((cast(k), cast(v)) for k, v in iteritems(obj))
|
||
|
|
if isinstance(obj, (list, tuple)):
|
||
|
|
return [cast(v) for v in obj]
|
||
|
|
if isinstance(obj, PASSTHROUGH):
|
||
|
|
return obj
|
||
|
|
if isinstance(obj, mitogen.core.UnicodeType):
|
||
|
|
return _cast(obj, mitogen.core.UnicodeType)
|
||
|
|
if isinstance(obj, mitogen.core.BytesType):
|
||
|
|
return _cast(obj, mitogen.core.BytesType)
|
||
|
|
|
||
|
|
raise TypeError("Cannot serialize: %r: %r" % (type(obj), obj))
|
||
|
|
|
||
|
|
|
||
|
|
def _cast(obj, desired_type):
|
||
|
|
result = desired_type(obj)
|
||
|
|
if type(result) is not desired_type:
|
||
|
|
raise TypeError("Cast of %r to %r failed, got %r"
|
||
|
|
% (type(obj), desired_type, type(result)))
|
||
|
|
return result
|