Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found

Target

Select target project
  • varac/stackspin
  • xeruf/stackspin
  • stackspin/stackspin
3 results
Show changes
Showing
with 0 additions and 11307 deletions
# 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
"""
On the Mitogen master, this is imported from ``mitogen/__init__.py`` as would
be expected. On the slave, it is built dynamically during startup.
"""
#: Library version as a tuple.
__version__ = (0, 2, 9)
#: This is :data:`False` in slave contexts. Previously it was used to prevent
#: re-execution of :mod:`__main__` in single file programs, however that now
#: happens automatically.
is_master = True
#: This is `0` in a master, otherwise it is the master-assigned ID unique to
#: the slave context used for message routing.
context_id = 0
#: This is :data:`None` in a master, otherwise it is the master-assigned ID
#: unique to the slave's parent context.
parent_id = None
#: This is an empty list in a master, otherwise it is a list of parent context
#: IDs ordered from most direct to least direct.
parent_ids = []
import os
_default_profiling = os.environ.get('MITOGEN_PROFILING') is not None
del os
def main(log_level='INFO', profiling=_default_profiling):
"""
Convenience decorator primarily useful for writing discardable test
scripts.
In the master process, when `func` is defined in the :mod:`__main__`
module, arranges for `func(router)` to be invoked immediately, with
:py:class:`mitogen.master.Router` construction and destruction handled just
as in :py:func:`mitogen.utils.run_with_router`. In slaves, this function
does nothing.
:param str log_level:
Logging package level to configure via
:py:func:`mitogen.utils.log_to_file`.
:param bool profiling:
If :py:data:`True`, equivalent to setting
:py:attr:`mitogen.master.Router.profiling` prior to router
construction. This causes ``/tmp`` files to be created everywhere at
the end of a successful run with :py:mod:`cProfile` output for every
thread.
Example:
::
import mitogen
import requests
def get_url(url):
return requests.get(url).text
@mitogen.main()
def main(router):
z = router.ssh(hostname='k3')
print(z.call(get_url, 'https://example.org/')))))
"""
def wrapper(func):
if func.__module__ != '__main__':
return func
import mitogen.parent
import mitogen.utils
if profiling:
mitogen.core.enable_profiling()
mitogen.master.Router.profiling = profiling
mitogen.utils.log_to_file(level=log_level)
return mitogen.core._profile_hook(
'app.main',
mitogen.utils.run_with_router,
func,
)
return wrapper
# 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 logging
import mitogen.core
import mitogen.parent
LOG = logging.getLogger(__name__)
class Options(mitogen.parent.Options):
container = None
username = None
buildah_path = 'buildah'
def __init__(self, container=None, buildah_path=None, username=None,
**kwargs):
super(Options, self).__init__(**kwargs)
assert container is not None
self.container = container
if buildah_path:
self.buildah_path = buildah_path
if username:
self.username = username
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = False
# TODO: better way of capturing errors such as "No such container."
create_child_args = {
'merge_stdio': True
}
def _get_name(self):
return u'buildah.' + self.options.container
def get_boot_command(self):
args = [self.options.buildah_path, 'run']
if self.options.username:
args += ['--user=' + self.options.username]
args += ['--', self.options.container]
return args + super(Connection, self).get_boot_command()
"""Utilities to support packages."""
# !mitogen: minify_safe
# NOTE: This module must remain compatible with Python 2.3, as it is shared
# by setuptools for distribution with Python 2.3 and up.
import os
import sys
import imp
import os.path
from types import ModuleType
__all__ = [
'get_importer', 'iter_importers', 'get_loader', 'find_loader',
'walk_packages', 'iter_modules', 'get_data',
'ImpImporter', 'ImpLoader', 'read_code', 'extend_path',
]
def read_code(stream):
# This helper is needed in order for the PEP 302 emulation to
# correctly handle compiled files
import marshal
magic = stream.read(4)
if magic != imp.get_magic():
return None
stream.read(4) # Skip timestamp
return marshal.load(stream)
def simplegeneric(func):
"""Make a trivial single-dispatch generic function"""
registry = {}
def wrapper(*args, **kw):
ob = args[0]
try:
cls = ob.__class__
except AttributeError:
cls = type(ob)
try:
mro = cls.__mro__
except AttributeError:
try:
class cls(cls, object):
pass
mro = cls.__mro__[1:]
except TypeError:
mro = object, # must be an ExtensionClass or some such :(
for t in mro:
if t in registry:
return registry[t](*args, **kw)
else:
return func(*args, **kw)
try:
wrapper.__name__ = func.__name__
except (TypeError, AttributeError):
pass # Python 2.3 doesn't allow functions to be renamed
def register(typ, func=None):
if func is None:
return lambda f: register(typ, f)
registry[typ] = func
return func
wrapper.__dict__ = func.__dict__
wrapper.__doc__ = func.__doc__
wrapper.register = register
return wrapper
def walk_packages(path=None, prefix='', onerror=None):
"""Yields (module_loader, name, ispkg) for all modules recursively
on path, or, if path is None, all accessible modules.
'path' should be either None or a list of paths to look for
modules in.
'prefix' is a string to output on the front of every module name
on output.
Note that this function must import all *packages* (NOT all
modules!) on the given path, in order to access the __path__
attribute to find submodules.
'onerror' is a function which gets called with one argument (the
name of the package which was being imported) if any exception
occurs while trying to import a package. If no onerror function is
supplied, ImportErrors are caught and ignored, while all other
exceptions are propagated, terminating the search.
Examples:
# list all modules python can access
walk_packages()
# list all submodules of ctypes
walk_packages(ctypes.__path__, ctypes.__name__+'.')
"""
def seen(p, m={}):
if p in m:
return True
m[p] = True
for importer, name, ispkg in iter_modules(path, prefix):
yield importer, name, ispkg
if ispkg:
try:
__import__(name)
except ImportError:
if onerror is not None:
onerror(name)
except Exception:
if onerror is not None:
onerror(name)
else:
raise
else:
path = getattr(sys.modules[name], '__path__', None) or []
# don't traverse path items we've seen before
path = [p for p in path if not seen(p)]
for item in walk_packages(path, name+'.', onerror):
yield item
def iter_modules(path=None, prefix=''):
"""Yields (module_loader, name, ispkg) for all submodules on path,
or, if path is None, all top-level modules on sys.path.
'path' should be either None or a list of paths to look for
modules in.
'prefix' is a string to output on the front of every module name
on output.
"""
if path is None:
importers = iter_importers()
else:
importers = map(get_importer, path)
yielded = {}
for i in importers:
for name, ispkg in iter_importer_modules(i, prefix):
if name not in yielded:
yielded[name] = 1
yield i, name, ispkg
#@simplegeneric
def iter_importer_modules(importer, prefix=''):
if not hasattr(importer, 'iter_modules'):
return []
return importer.iter_modules(prefix)
iter_importer_modules = simplegeneric(iter_importer_modules)
class ImpImporter:
"""PEP 302 Importer that wraps Python's "classic" import algorithm
ImpImporter(dirname) produces a PEP 302 importer that searches that
directory. ImpImporter(None) produces a PEP 302 importer that searches
the current sys.path, plus any modules that are frozen or built-in.
Note that ImpImporter does not currently support being used by placement
on sys.meta_path.
"""
def __init__(self, path=None):
self.path = path
def find_module(self, fullname, path=None):
# Note: we ignore 'path' argument since it is only used via meta_path
subname = fullname.split(".")[-1]
if subname != fullname and self.path is None:
return None
if self.path is None:
path = None
else:
path = [os.path.realpath(self.path)]
try:
file, filename, etc = imp.find_module(subname, path)
except ImportError:
return None
return ImpLoader(fullname, file, filename, etc)
def iter_modules(self, prefix=''):
if self.path is None or not os.path.isdir(self.path):
return
yielded = {}
import inspect
try:
filenames = os.listdir(self.path)
except OSError:
# ignore unreadable directories like import does
filenames = []
filenames.sort() # handle packages before same-named modules
for fn in filenames:
modname = inspect.getmodulename(fn)
if modname=='__init__' or modname in yielded:
continue
path = os.path.join(self.path, fn)
ispkg = False
if not modname and os.path.isdir(path) and '.' not in fn:
modname = fn
try:
dircontents = os.listdir(path)
except OSError:
# ignore unreadable directories like import does
dircontents = []
for fn in dircontents:
subname = inspect.getmodulename(fn)
if subname=='__init__':
ispkg = True
break
else:
continue # not a package
if modname and '.' not in modname:
yielded[modname] = 1
yield prefix + modname, ispkg
class ImpLoader:
"""PEP 302 Loader that wraps Python's "classic" import algorithm
"""
code = source = None
def __init__(self, fullname, file, filename, etc):
self.file = file
self.filename = filename
self.fullname = fullname
self.etc = etc
def load_module(self, fullname):
self._reopen()
try:
mod = imp.load_module(fullname, self.file, self.filename, self.etc)
finally:
if self.file:
self.file.close()
# Note: we don't set __loader__ because we want the module to look
# normal; i.e. this is just a wrapper for standard import machinery
return mod
def get_data(self, pathname):
return open(pathname, "rb").read()
def _reopen(self):
if self.file and self.file.closed:
mod_type = self.etc[2]
if mod_type==imp.PY_SOURCE:
self.file = open(self.filename, 'rU')
elif mod_type in (imp.PY_COMPILED, imp.C_EXTENSION):
self.file = open(self.filename, 'rb')
def _fix_name(self, fullname):
if fullname is None:
fullname = self.fullname
elif fullname != self.fullname:
raise ImportError("Loader for module %s cannot handle "
"module %s" % (self.fullname, fullname))
return fullname
def is_package(self, fullname):
fullname = self._fix_name(fullname)
return self.etc[2]==imp.PKG_DIRECTORY
def get_code(self, fullname=None):
fullname = self._fix_name(fullname)
if self.code is None:
mod_type = self.etc[2]
if mod_type==imp.PY_SOURCE:
source = self.get_source(fullname)
self.code = compile(source, self.filename, 'exec')
elif mod_type==imp.PY_COMPILED:
self._reopen()
try:
self.code = read_code(self.file)
finally:
self.file.close()
elif mod_type==imp.PKG_DIRECTORY:
self.code = self._get_delegate().get_code()
return self.code
def get_source(self, fullname=None):
fullname = self._fix_name(fullname)
if self.source is None:
mod_type = self.etc[2]
if mod_type==imp.PY_SOURCE:
self._reopen()
try:
self.source = self.file.read()
finally:
self.file.close()
elif mod_type==imp.PY_COMPILED:
if os.path.exists(self.filename[:-1]):
f = open(self.filename[:-1], 'rU')
self.source = f.read()
f.close()
elif mod_type==imp.PKG_DIRECTORY:
self.source = self._get_delegate().get_source()
return self.source
def _get_delegate(self):
return ImpImporter(self.filename).find_module('__init__')
def get_filename(self, fullname=None):
fullname = self._fix_name(fullname)
mod_type = self.etc[2]
if self.etc[2]==imp.PKG_DIRECTORY:
return self._get_delegate().get_filename()
elif self.etc[2] in (imp.PY_SOURCE, imp.PY_COMPILED, imp.C_EXTENSION):
return self.filename
return None
try:
import zipimport
from zipimport import zipimporter
def iter_zipimport_modules(importer, prefix=''):
dirlist = zipimport._zip_directory_cache[importer.archive].keys()
dirlist.sort()
_prefix = importer.prefix
plen = len(_prefix)
yielded = {}
import inspect
for fn in dirlist:
if not fn.startswith(_prefix):
continue
fn = fn[plen:].split(os.sep)
if len(fn)==2 and fn[1].startswith('__init__.py'):
if fn[0] not in yielded:
yielded[fn[0]] = 1
yield fn[0], True
if len(fn)!=1:
continue
modname = inspect.getmodulename(fn[0])
if modname=='__init__':
continue
if modname and '.' not in modname and modname not in yielded:
yielded[modname] = 1
yield prefix + modname, False
iter_importer_modules.register(zipimporter, iter_zipimport_modules)
except ImportError:
pass
def get_importer(path_item):
"""Retrieve a PEP 302 importer for the given path item
The returned importer is cached in sys.path_importer_cache
if it was newly created by a path hook.
If there is no importer, a wrapper around the basic import
machinery is returned. This wrapper is never inserted into
the importer cache (None is inserted instead).
The cache (or part of it) can be cleared manually if a
rescan of sys.path_hooks is necessary.
"""
try:
importer = sys.path_importer_cache[path_item]
except KeyError:
for path_hook in sys.path_hooks:
try:
importer = path_hook(path_item)
break
except ImportError:
pass
else:
importer = None
sys.path_importer_cache.setdefault(path_item, importer)
if importer is None:
try:
importer = ImpImporter(path_item)
except ImportError:
importer = None
return importer
def iter_importers(fullname=""):
"""Yield PEP 302 importers for the given module name
If fullname contains a '.', the importers will be for the package
containing fullname, otherwise they will be importers for sys.meta_path,
sys.path, and Python's "classic" import machinery, in that order. If
the named module is in a package, that package is imported as a side
effect of invoking this function.
Non PEP 302 mechanisms (e.g. the Windows registry) used by the
standard import machinery to find files in alternative locations
are partially supported, but are searched AFTER sys.path. Normally,
these locations are searched BEFORE sys.path, preventing sys.path
entries from shadowing them.
For this to cause a visible difference in behaviour, there must
be a module or package name that is accessible via both sys.path
and one of the non PEP 302 file system mechanisms. In this case,
the emulation will find the former version, while the builtin
import mechanism will find the latter.
Items of the following types can be affected by this discrepancy:
imp.C_EXTENSION, imp.PY_SOURCE, imp.PY_COMPILED, imp.PKG_DIRECTORY
"""
if fullname.startswith('.'):
raise ImportError("Relative module names not supported")
if '.' in fullname:
# Get the containing package's __path__
pkg = '.'.join(fullname.split('.')[:-1])
if pkg not in sys.modules:
__import__(pkg)
path = getattr(sys.modules[pkg], '__path__', None) or []
else:
for importer in sys.meta_path:
yield importer
path = sys.path
for item in path:
yield get_importer(item)
if '.' not in fullname:
yield ImpImporter()
def get_loader(module_or_name):
"""Get a PEP 302 "loader" object for module_or_name
If the module or package is accessible via the normal import
mechanism, a wrapper around the relevant part of that machinery
is returned. Returns None if the module cannot be found or imported.
If the named module is not already imported, its containing package
(if any) is imported, in order to establish the package __path__.
This function uses iter_importers(), and is thus subject to the same
limitations regarding platform-specific special import locations such
as the Windows registry.
"""
if module_or_name in sys.modules:
module_or_name = sys.modules[module_or_name]
if isinstance(module_or_name, ModuleType):
module = module_or_name
loader = getattr(module, '__loader__', None)
if loader is not None:
return loader
fullname = module.__name__
else:
fullname = module_or_name
return find_loader(fullname)
def find_loader(fullname):
"""Find a PEP 302 "loader" object for fullname
If fullname contains dots, path must be the containing package's __path__.
Returns None if the module cannot be found or imported. This function uses
iter_importers(), and is thus subject to the same limitations regarding
platform-specific special import locations such as the Windows registry.
"""
for importer in iter_importers(fullname):
loader = importer.find_module(fullname)
if loader is not None:
return loader
return None
def extend_path(path, name):
"""Extend a package's path.
Intended use is to place the following code in a package's __init__.py:
from pkgutil import extend_path
__path__ = extend_path(__path__, __name__)
This will add to the package's __path__ all subdirectories of
directories on sys.path named after the package. This is useful
if one wants to distribute different parts of a single logical
package as multiple directories.
It also looks for *.pkg files beginning where * matches the name
argument. This feature is similar to *.pth files (see site.py),
except that it doesn't special-case lines starting with 'import'.
A *.pkg file is trusted at face value: apart from checking for
duplicates, all entries found in a *.pkg file are added to the
path, regardless of whether they are exist the filesystem. (This
is a feature.)
If the input path is not a list (as is the case for frozen
packages) it is returned unchanged. The input path is not
modified; an extended copy is returned. Items are only appended
to the copy at the end.
It is assumed that sys.path is a sequence. Items of sys.path that
are not (unicode or 8-bit) strings referring to existing
directories are ignored. Unicode items of sys.path that cause
errors when used as filenames may cause this function to raise an
exception (in line with os.path.isdir() behavior).
"""
if not isinstance(path, list):
# This could happen e.g. when this is called from inside a
# frozen package. Return the path unchanged in that case.
return path
pname = os.path.join(*name.split('.')) # Reconstitute as relative path
# Just in case os.extsep != '.'
sname = os.extsep.join(name.split('.'))
sname_pkg = sname + os.extsep + "pkg"
init_py = "__init__" + os.extsep + "py"
path = path[:] # Start with a copy of the existing path
for dir in sys.path:
if not isinstance(dir, basestring) or not os.path.isdir(dir):
continue
subdir = os.path.join(dir, pname)
# XXX This may still add duplicate entries to path on
# case-insensitive filesystems
initfile = os.path.join(subdir, init_py)
if subdir not in path and os.path.isfile(initfile):
path.append(subdir)
# XXX Is this the right thing for subpackages like zope.app?
# It looks for a file named "zope.app.pkg"
pkgfile = os.path.join(dir, sname_pkg)
if os.path.isfile(pkgfile):
try:
f = open(pkgfile)
except IOError:
msg = sys.exc_info()[1]
sys.stderr.write("Can't open %s: %s\n" %
(pkgfile, msg))
else:
for line in f:
line = line.rstrip('\n')
if not line or line.startswith('#'):
continue
path.append(line) # Don't check for existence!
f.close()
return path
def get_data(package, resource):
"""Get a resource from a package.
This is a wrapper round the PEP 302 loader get_data API. The package
argument should be the name of a package, in standard module format
(foo.bar). The resource argument should be in the form of a relative
filename, using '/' as the path separator. The parent directory name '..'
is not allowed, and nor is a rooted name (starting with a '/').
The function returns a binary string, which is the contents of the
specified resource.
For packages located in the filesystem, which have already been imported,
this is the rough equivalent of
d = os.path.dirname(sys.modules[package].__file__)
data = open(os.path.join(d, resource), 'rb').read()
If the package cannot be located or loaded, or it uses a PEP 302 loader
which does not support get_data(), then None is returned.
"""
loader = get_loader(package)
if loader is None or not hasattr(loader, 'get_data'):
return None
mod = sys.modules.get(package) or loader.load_module(package)
if mod is None or not hasattr(mod, '__file__'):
return None
# Modify the resource name to be compatible with the loader.get_data
# signature - an os.path format "filename" starting with the dirname of
# the package's __file__
parts = resource.split('/')
parts.insert(0, os.path.dirname(mod.__file__))
resource_name = os.path.join(*parts)
return loader.get_data(resource_name)
"""Tokenization help for Python programs.
generate_tokens(readline) is a generator that breaks a stream of
text into Python tokens. It accepts a readline-like method which is called
repeatedly to get the next line of input (or "" for EOF). It generates
5-tuples with these members:
the token type (see token.py)
the token (a string)
the starting (row, column) indices of the token (a 2-tuple of ints)
the ending (row, column) indices of the token (a 2-tuple of ints)
the original line (string)
It is designed to match the working of the Python tokenizer exactly, except
that it produces COMMENT tokens for comments and gives type OP for all
operators
Older entry points
tokenize_loop(readline, tokeneater)
tokenize(readline, tokeneater=printtoken)
are the same, except instead of generating tokens, tokeneater is a callback
function to which the 5 fields described above are passed as 5 arguments,
each time a new token is found."""
# !mitogen: minify_safe
__author__ = 'Ka-Ping Yee <ping@lfw.org>'
__credits__ = ('GvR, ESR, Tim Peters, Thomas Wouters, Fred Drake, '
'Skip Montanaro, Raymond Hettinger')
from itertools import chain
import string, re
from token import *
import token
__all__ = [x for x in dir(token) if not x.startswith("_")]
__all__ += ["COMMENT", "tokenize", "generate_tokens", "NL", "untokenize"]
del token
COMMENT = N_TOKENS
tok_name[COMMENT] = 'COMMENT'
NL = N_TOKENS + 1
tok_name[NL] = 'NL'
N_TOKENS += 2
def group(*choices): return '(' + '|'.join(choices) + ')'
def any(*choices): return group(*choices) + '*'
def maybe(*choices): return group(*choices) + '?'
Whitespace = r'[ \f\t]*'
Comment = r'#[^\r\n]*'
Ignore = Whitespace + any(r'\\\r?\n' + Whitespace) + maybe(Comment)
Name = r'[a-zA-Z_]\w*'
Hexnumber = r'0[xX][\da-fA-F]+[lL]?'
Octnumber = r'(0[oO][0-7]+)|(0[0-7]*)[lL]?'
Binnumber = r'0[bB][01]+[lL]?'
Decnumber = r'[1-9]\d*[lL]?'
Intnumber = group(Hexnumber, Binnumber, Octnumber, Decnumber)
Exponent = r'[eE][-+]?\d+'
Pointfloat = group(r'\d+\.\d*', r'\.\d+') + maybe(Exponent)
Expfloat = r'\d+' + Exponent
Floatnumber = group(Pointfloat, Expfloat)
Imagnumber = group(r'\d+[jJ]', Floatnumber + r'[jJ]')
Number = group(Imagnumber, Floatnumber, Intnumber)
# Tail end of ' string.
Single = r"[^'\\]*(?:\\.[^'\\]*)*'"
# Tail end of " string.
Double = r'[^"\\]*(?:\\.[^"\\]*)*"'
# Tail end of ''' string.
Single3 = r"[^'\\]*(?:(?:\\.|'(?!''))[^'\\]*)*'''"
# Tail end of """ string.
Double3 = r'[^"\\]*(?:(?:\\.|"(?!""))[^"\\]*)*"""'
Triple = group("[uUbB]?[rR]?'''", '[uUbB]?[rR]?"""')
# Single-line ' or " string.
String = group(r"[uUbB]?[rR]?'[^\n'\\]*(?:\\.[^\n'\\]*)*'",
r'[uUbB]?[rR]?"[^\n"\\]*(?:\\.[^\n"\\]*)*"')
# Because of leftmost-then-longest match semantics, be sure to put the
# longest operators first (e.g., if = came before ==, == would get
# recognized as two instances of =).
Operator = group(r"\*\*=?", r">>=?", r"<<=?", r"<>", r"!=",
r"//=?",
r"[+\-*/%&|^=<>]=?",
r"~")
Bracket = '[][(){}]'
Special = group(r'\r?\n', r'[:;.,`@]')
Funny = group(Operator, Bracket, Special)
PlainToken = group(Number, Funny, String, Name)
Token = Ignore + PlainToken
# First (or only) line of ' or " string.
ContStr = group(r"[uUbB]?[rR]?'[^\n'\\]*(?:\\.[^\n'\\]*)*" +
group("'", r'\\\r?\n'),
r'[uUbB]?[rR]?"[^\n"\\]*(?:\\.[^\n"\\]*)*' +
group('"', r'\\\r?\n'))
PseudoExtras = group(r'\\\r?\n|\Z', Comment, Triple)
PseudoToken = Whitespace + group(PseudoExtras, Number, Funny, ContStr, Name)
tokenprog, pseudoprog, single3prog, double3prog = map(
re.compile, (Token, PseudoToken, Single3, Double3))
endprogs = {"'": re.compile(Single), '"': re.compile(Double),
"'''": single3prog, '"""': double3prog,
"r'''": single3prog, 'r"""': double3prog,
"u'''": single3prog, 'u"""': double3prog,
"ur'''": single3prog, 'ur"""': double3prog,
"R'''": single3prog, 'R"""': double3prog,
"U'''": single3prog, 'U"""': double3prog,
"uR'''": single3prog, 'uR"""': double3prog,
"Ur'''": single3prog, 'Ur"""': double3prog,
"UR'''": single3prog, 'UR"""': double3prog,
"b'''": single3prog, 'b"""': double3prog,
"br'''": single3prog, 'br"""': double3prog,
"B'''": single3prog, 'B"""': double3prog,
"bR'''": single3prog, 'bR"""': double3prog,
"Br'''": single3prog, 'Br"""': double3prog,
"BR'''": single3prog, 'BR"""': double3prog,
'r': None, 'R': None, 'u': None, 'U': None,
'b': None, 'B': None}
triple_quoted = {}
for t in ("'''", '"""',
"r'''", 'r"""', "R'''", 'R"""',
"u'''", 'u"""', "U'''", 'U"""',
"ur'''", 'ur"""', "Ur'''", 'Ur"""',
"uR'''", 'uR"""', "UR'''", 'UR"""',
"b'''", 'b"""', "B'''", 'B"""',
"br'''", 'br"""', "Br'''", 'Br"""',
"bR'''", 'bR"""', "BR'''", 'BR"""'):
triple_quoted[t] = t
single_quoted = {}
for t in ("'", '"',
"r'", 'r"', "R'", 'R"',
"u'", 'u"', "U'", 'U"',
"ur'", 'ur"', "Ur'", 'Ur"',
"uR'", 'uR"', "UR'", 'UR"',
"b'", 'b"', "B'", 'B"',
"br'", 'br"', "Br'", 'Br"',
"bR'", 'bR"', "BR'", 'BR"' ):
single_quoted[t] = t
tabsize = 8
class TokenError(Exception): pass
class StopTokenizing(Exception): pass
def printtoken(type, token, srow_scol, erow_ecol, line): # for testing
srow, scol = srow_scol
erow, ecol = erow_ecol
print("%d,%d-%d,%d:\t%s\t%s" % \
(srow, scol, erow, ecol, tok_name[type], repr(token)))
def tokenize(readline, tokeneater=printtoken):
"""
The tokenize() function accepts two parameters: one representing the
input stream, and one providing an output mechanism for tokenize().
The first parameter, readline, must be a callable object which provides
the same interface as the readline() method of built-in file objects.
Each call to the function should return one line of input as a string.
The second parameter, tokeneater, must also be a callable object. It is
called once for each token, with five arguments, corresponding to the
tuples generated by generate_tokens().
"""
try:
tokenize_loop(readline, tokeneater)
except StopTokenizing:
pass
# backwards compatible interface
def tokenize_loop(readline, tokeneater):
for token_info in generate_tokens(readline):
tokeneater(*token_info)
class Untokenizer:
def __init__(self):
self.tokens = []
self.prev_row = 1
self.prev_col = 0
def add_whitespace(self, start):
row, col = start
if row < self.prev_row or row == self.prev_row and col < self.prev_col:
raise ValueError("start ({},{}) precedes previous end ({},{})"
.format(row, col, self.prev_row, self.prev_col))
row_offset = row - self.prev_row
if row_offset:
self.tokens.append("\\\n" * row_offset)
self.prev_col = 0
col_offset = col - self.prev_col
if col_offset:
self.tokens.append(" " * col_offset)
def untokenize(self, iterable):
it = iter(iterable)
indents = []
startline = False
for t in it:
if len(t) == 2:
self.compat(t, it)
break
tok_type, token, start, end, line = t
if tok_type == ENDMARKER:
break
if tok_type == INDENT:
indents.append(token)
continue
elif tok_type == DEDENT:
indents.pop()
self.prev_row, self.prev_col = end
continue
elif tok_type in (NEWLINE, NL):
startline = True
elif startline and indents:
indent = indents[-1]
if start[1] >= len(indent):
self.tokens.append(indent)
self.prev_col = len(indent)
startline = False
self.add_whitespace(start)
self.tokens.append(token)
self.prev_row, self.prev_col = end
if tok_type in (NEWLINE, NL):
self.prev_row += 1
self.prev_col = 0
return "".join(self.tokens)
def compat(self, token, iterable):
indents = []
toks_append = self.tokens.append
startline = token[0] in (NEWLINE, NL)
prevstring = False
for tok in chain([token], iterable):
toknum, tokval = tok[:2]
if toknum in (NAME, NUMBER):
tokval += ' '
# Insert a space between two consecutive strings
if toknum == STRING:
if prevstring:
tokval = ' ' + tokval
prevstring = True
else:
prevstring = False
if toknum == INDENT:
indents.append(tokval)
continue
elif toknum == DEDENT:
indents.pop()
continue
elif toknum in (NEWLINE, NL):
startline = True
elif startline and indents:
toks_append(indents[-1])
startline = False
toks_append(tokval)
def untokenize(iterable):
"""Transform tokens back into Python source code.
Each element returned by the iterable must be a token sequence
with at least two elements, a token number and token value. If
only two tokens are passed, the resulting output is poor.
Round-trip invariant for full input:
Untokenized source will match input source exactly
Round-trip invariant for limited intput:
# Output text will tokenize the back to the input
t1 = [tok[:2] for tok in generate_tokens(f.readline)]
newcode = untokenize(t1)
readline = iter(newcode.splitlines(1)).next
t2 = [tok[:2] for tok in generate_tokens(readline)]
assert t1 == t2
"""
ut = Untokenizer()
return ut.untokenize(iterable)
def generate_tokens(readline):
"""
The generate_tokens() generator requires one argument, readline, which
must be a callable object which provides the same interface as the
readline() method of built-in file objects. Each call to the function
should return one line of input as a string. Alternately, readline
can be a callable function terminating with StopIteration:
readline = open(myfile).next # Example of alternate readline
The generator produces 5-tuples with these members: the token type; the
token string; a 2-tuple (srow, scol) of ints specifying the row and
column where the token begins in the source; a 2-tuple (erow, ecol) of
ints specifying the row and column where the token ends in the source;
and the line on which the token was found. The line passed is the
logical line; continuation lines are included.
"""
lnum = parenlev = continued = 0
namechars, numchars = string.ascii_letters + '_', '0123456789'
contstr, needcont = '', 0
contline = None
indents = [0]
while 1: # loop over lines in stream
try:
line = readline()
except StopIteration:
line = ''
lnum += 1
pos, max = 0, len(line)
if contstr: # continued string
if not line:
raise TokenError("EOF in multi-line string", strstart)
endmatch = endprog.match(line)
if endmatch:
pos = end = endmatch.end(0)
yield (STRING, contstr + line[:end],
strstart, (lnum, end), contline + line)
contstr, needcont = '', 0
contline = None
elif needcont and line[-2:] != '\\\n' and line[-3:] != '\\\r\n':
yield (ERRORTOKEN, contstr + line,
strstart, (lnum, len(line)), contline)
contstr = ''
contline = None
continue
else:
contstr = contstr + line
contline = contline + line
continue
elif parenlev == 0 and not continued: # new statement
if not line: break
column = 0
while pos < max: # measure leading whitespace
if line[pos] == ' ':
column += 1
elif line[pos] == '\t':
column = (column//tabsize + 1)*tabsize
elif line[pos] == '\f':
column = 0
else:
break
pos += 1
if pos == max:
break
if line[pos] in '#\r\n': # skip comments or blank lines
if line[pos] == '#':
comment_token = line[pos:].rstrip('\r\n')
nl_pos = pos + len(comment_token)
yield (COMMENT, comment_token,
(lnum, pos), (lnum, pos + len(comment_token)), line)
yield (NL, line[nl_pos:],
(lnum, nl_pos), (lnum, len(line)), line)
else:
yield ((NL, COMMENT)[line[pos] == '#'], line[pos:],
(lnum, pos), (lnum, len(line)), line)
continue
if column > indents[-1]: # count indents or dedents
indents.append(column)
yield (INDENT, line[:pos], (lnum, 0), (lnum, pos), line)
while column < indents[-1]:
if column not in indents:
raise IndentationError(
"unindent does not match any outer indentation level",
("<tokenize>", lnum, pos, line))
indents = indents[:-1]
yield (DEDENT, '', (lnum, pos), (lnum, pos), line)
else: # continued statement
if not line:
raise TokenError("EOF in multi-line statement", (lnum, 0))
continued = 0
while pos < max:
pseudomatch = pseudoprog.match(line, pos)
if pseudomatch: # scan for tokens
start, end = pseudomatch.span(1)
spos, epos, pos = (lnum, start), (lnum, end), end
if start == end:
continue
token, initial = line[start:end], line[start]
if initial in numchars or \
(initial == '.' and token != '.'): # ordinary number
yield (NUMBER, token, spos, epos, line)
elif initial in '\r\n':
if parenlev > 0:
n = NL
else:
n = NEWLINE
yield (n, token, spos, epos, line)
elif initial == '#':
assert not token.endswith("\n")
yield (COMMENT, token, spos, epos, line)
elif token in triple_quoted:
endprog = endprogs[token]
endmatch = endprog.match(line, pos)
if endmatch: # all on one line
pos = endmatch.end(0)
token = line[start:pos]
yield (STRING, token, spos, (lnum, pos), line)
else:
strstart = (lnum, start) # multiple lines
contstr = line[start:]
contline = line
break
elif initial in single_quoted or \
token[:2] in single_quoted or \
token[:3] in single_quoted:
if token[-1] == '\n': # continued string
strstart = (lnum, start)
endprog = (endprogs[initial] or endprogs[token[1]] or
endprogs[token[2]])
contstr, needcont = line[start:], 1
contline = line
break
else: # ordinary string
yield (STRING, token, spos, epos, line)
elif initial in namechars: # ordinary name
yield (NAME, token, spos, epos, line)
elif initial == '\\': # continued stmt
continued = 1
else:
if initial in '([{':
parenlev += 1
elif initial in ')]}':
parenlev -= 1
yield (OP, token, spos, epos, line)
else:
yield (ERRORTOKEN, line[pos],
(lnum, pos), (lnum, pos+1), line)
pos += 1
for indent in indents[1:]: # pop remaining indent levels
yield (DEDENT, '', (lnum, 0), (lnum, 0), '')
yield (ENDMARKER, '', (lnum, 0), (lnum, 0), '')
if __name__ == '__main__': # testing
import sys
if len(sys.argv) > 1:
tokenize(open(sys.argv[1]).readline)
else:
tokenize(sys.stdin.readline)
# 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
"""
This module implements most package functionality, but remains separate from
non-essential code in order to reduce its size, since it is also serves as the
bootstrap implementation sent to every new slave context.
"""
import binascii
import collections
import encodings.latin_1
import encodings.utf_8
import errno
import fcntl
import itertools
import linecache
import logging
import os
import pickle as py_pickle
import pstats
import signal
import socket
import struct
import sys
import syslog
import threading
import time
import traceback
import warnings
import weakref
import zlib
# Python >3.7 deprecated the imp module.
warnings.filterwarnings('ignore', message='the imp module is deprecated')
import imp
# Absolute imports for <2.5.
select = __import__('select')
try:
import cProfile
except ImportError:
cProfile = None
try:
import thread
except ImportError:
import threading as thread
try:
import cPickle as pickle
except ImportError:
import pickle
try:
from cStringIO import StringIO as BytesIO
except ImportError:
from io import BytesIO
try:
BaseException
except NameError:
BaseException = Exception
try:
ModuleNotFoundError
except NameError:
ModuleNotFoundError = ImportError
# TODO: usage of 'import' after setting __name__, but before fixing up
# sys.modules generates a warning. This happens when profiling = True.
warnings.filterwarnings('ignore',
"Parent module 'mitogen' not found while handling absolute import")
LOG = logging.getLogger('mitogen')
IOLOG = logging.getLogger('mitogen.io')
IOLOG.setLevel(logging.INFO)
# str.encode() may take import lock. Deadlock possible if broker calls
# .encode() on behalf of thread currently waiting for module.
LATIN1_CODEC = encodings.latin_1.Codec()
_v = False
_vv = False
GET_MODULE = 100
CALL_FUNCTION = 101
FORWARD_LOG = 102
ADD_ROUTE = 103
DEL_ROUTE = 104
ALLOCATE_ID = 105
SHUTDOWN = 106
LOAD_MODULE = 107
FORWARD_MODULE = 108
DETACHING = 109
CALL_SERVICE = 110
STUB_CALL_SERVICE = 111
#: Special value used to signal disconnection or the inability to route a
#: message, when it appears in the `reply_to` field. Usually causes
#: :class:`mitogen.core.ChannelError` to be raised when it is received.
#:
#: It indicates the sender did not know how to process the message, or wishes
#: no further messages to be delivered to it. It is used when:
#:
#: * a remote receiver is disconnected or explicitly closed.
#: * a related message could not be delivered due to no route existing for it.
#: * a router is being torn down, as a sentinel value to notify
#: :meth:`mitogen.core.Router.add_handler` callbacks to clean up.
IS_DEAD = 999
try:
BaseException
except NameError:
BaseException = Exception
PY24 = sys.version_info < (2, 5)
PY3 = sys.version_info > (3,)
if PY3:
b = str.encode
BytesType = bytes
UnicodeType = str
FsPathTypes = (str,)
BufferType = lambda buf, start: memoryview(buf)[start:]
long = int
else:
b = str
BytesType = str
FsPathTypes = (str, unicode)
BufferType = buffer
UnicodeType = unicode
AnyTextType = (BytesType, UnicodeType)
try:
next
except NameError:
next = lambda it: it.next()
# #550: prehistoric WSL did not advertise itself in uname output.
try:
fp = open('/proc/sys/kernel/osrelease')
IS_WSL = 'Microsoft' in fp.read()
fp.close()
except IOError:
IS_WSL = False
#: Default size for calls to :meth:`Side.read` or :meth:`Side.write`, and the
#: size of buffers configured by :func:`mitogen.parent.create_socketpair`. This
#: value has many performance implications, 128KiB seems to be a sweet spot.
#:
#: * When set low, large messages cause many :class:`Broker` IO loop
#: iterations, burning CPU and reducing throughput.
#: * When set high, excessive RAM is reserved by the OS for socket buffers (2x
#: per child), and an identically sized temporary userspace buffer is
#: allocated on each read that requires zeroing, and over a particular size
#: may require two system calls to allocate/deallocate.
#:
#: Care must be taken to ensure the underlying kernel object and receiving
#: program support the desired size. For example,
#:
#: * Most UNIXes have TTYs with fixed 2KiB-4KiB buffers, making them unsuitable
#: for efficient IO.
#: * Different UNIXes have varying presets for pipes, which may not be
#: configurable. On recent Linux the default pipe buffer size is 64KiB, but
#: under memory pressure may be as low as 4KiB for unprivileged processes.
#: * When communication is via an intermediary process, its internal buffers
#: effect the speed OS buffers will drain. For example OpenSSH uses 64KiB
#: reads.
#:
#: An ideal :class:`Message` has a size that is a multiple of
#: :data:`CHUNK_SIZE` inclusive of headers, to avoid wasting IO loop iterations
#: writing small trailer chunks.
CHUNK_SIZE = 131072
_tls = threading.local()
if __name__ == 'mitogen.core':
# When loaded using import mechanism, ExternalContext.main() will not have
# a chance to set the synthetic mitogen global, so just import it here.
import mitogen
else:
# When loaded as __main__, ensure classes and functions gain a __module__
# attribute consistent with the host process, so that pickling succeeds.
__name__ = 'mitogen.core'
class Error(Exception):
"""
Base for all exceptions raised by Mitogen.
:param str fmt:
Exception text, or format string if `args` is non-empty.
:param tuple args:
Format string arguments.
"""
def __init__(self, fmt=None, *args):
if args:
fmt %= args
if fmt and not isinstance(fmt, UnicodeType):
fmt = fmt.decode('utf-8')
Exception.__init__(self, fmt)
class LatchError(Error):
"""
Raised when an attempt is made to use a :class:`mitogen.core.Latch` that
has been marked closed.
"""
pass
class Blob(BytesType):
"""
A serializable bytes subclass whose content is summarized in repr() output,
making it suitable for logging binary data.
"""
def __repr__(self):
return '[blob: %d bytes]' % len(self)
def __reduce__(self):
return (Blob, (BytesType(self),))
class Secret(UnicodeType):
"""
A serializable unicode subclass whose content is masked in repr() output,
making it suitable for logging passwords.
"""
def __repr__(self):
return '[secret]'
if not PY3:
# TODO: what is this needed for in 2.x?
def __str__(self):
return UnicodeType(self)
def __reduce__(self):
return (Secret, (UnicodeType(self),))
class Kwargs(dict):
"""
A serializable dict subclass that indicates its keys should be coerced to
Unicode on Python 3 and bytes on Python<2.6.
Python 2 produces keyword argument dicts whose keys are bytes, requiring a
helper to ensure compatibility with Python 3 where Unicode is required,
whereas Python 3 produces keyword argument dicts whose keys are Unicode,
requiring a helper for Python 2.4/2.5, where bytes are required.
"""
if PY3:
def __init__(self, dct):
for k, v in dct.items():
if type(k) is bytes:
self[k.decode()] = v
else:
self[k] = v
elif sys.version_info < (2, 6, 5):
def __init__(self, dct):
for k, v in dct.iteritems():
if type(k) is unicode:
k, _ = encodings.utf_8.encode(k)
self[k] = v
def __repr__(self):
return 'Kwargs(%s)' % (dict.__repr__(self),)
def __reduce__(self):
return (Kwargs, (dict(self),))
class CallError(Error):
"""
Serializable :class:`Error` subclass raised when :meth:`Context.call()
<mitogen.parent.Context.call>` fails. A copy of the traceback from the
external context is appended to the exception message.
"""
def __init__(self, fmt=None, *args):
if not isinstance(fmt, BaseException):
Error.__init__(self, fmt, *args)
else:
e = fmt
cls = e.__class__
fmt = '%s.%s: %s' % (cls.__module__, cls.__name__, e)
tb = sys.exc_info()[2]
if tb:
fmt += '\n'
fmt += ''.join(traceback.format_tb(tb))
Error.__init__(self, fmt)
def __reduce__(self):
return (_unpickle_call_error, (self.args[0],))
def _unpickle_call_error(s):
if not (type(s) is UnicodeType and len(s) < 10000):
raise TypeError('cannot unpickle CallError: bad input')
return CallError(s)
class ChannelError(Error):
"""
Raised when a channel dies or has been closed.
"""
remote_msg = 'Channel closed by remote end.'
local_msg = 'Channel closed by local end.'
class StreamError(Error):
"""
Raised when a stream cannot be established.
"""
pass
class TimeoutError(Error):
"""
Raised when a timeout occurs on a stream.
"""
pass
def to_text(o):
"""
Coerce `o` to Unicode by decoding it from UTF-8 if it is an instance of
:class:`bytes`, otherwise pass it to the :class:`str` constructor. The
returned object is always a plain :class:`str`, any subclass is removed.
"""
if isinstance(o, BytesType):
return o.decode('utf-8')
return UnicodeType(o)
# Documented in api.rst to work around Sphinx limitation.
now = getattr(time, 'monotonic', time.time)
# Python 2.4
try:
any
except NameError:
def any(it):
for elem in it:
if elem:
return True
def _partition(s, sep, find):
"""
(str|unicode).(partition|rpartition) for Python 2.4/2.5.
"""
idx = find(sep)
if idx != -1:
left = s[0:idx]
return left, sep, s[len(left)+len(sep):]
if hasattr(UnicodeType, 'rpartition'):
str_partition = UnicodeType.partition
str_rpartition = UnicodeType.rpartition
bytes_partition = BytesType.partition
else:
def str_partition(s, sep):
return _partition(s, sep, s.find) or (s, u'', u'')
def str_rpartition(s, sep):
return _partition(s, sep, s.rfind) or (u'', u'', s)
def bytes_partition(s, sep):
return _partition(s, sep, s.find) or (s, '', '')
def _has_parent_authority(context_id):
return (
(context_id == mitogen.context_id) or
(context_id in mitogen.parent_ids)
)
def has_parent_authority(msg, _stream=None):
"""
Policy function for use with :class:`Receiver` and
:meth:`Router.add_handler` that requires incoming messages to originate
from a parent context, or on a :class:`Stream` whose :attr:`auth_id
<Stream.auth_id>` has been set to that of a parent context or the current
context.
"""
return _has_parent_authority(msg.auth_id)
def _signals(obj, signal):
return (
obj.__dict__
.setdefault('_signals', {})
.setdefault(signal, [])
)
def listen(obj, name, func):
"""
Arrange for `func()` to be invoked when signal `name` is fired on `obj`.
"""
_signals(obj, name).append(func)
def unlisten(obj, name, func):
"""
Remove `func()` from the list of functions invoked when signal `name` is
fired by `obj`.
:raises ValueError:
`func()` was not on the list.
"""
_signals(obj, name).remove(func)
def fire(obj, name, *args, **kwargs):
"""
Arrange for `func(*args, **kwargs)` to be invoked for every function
registered for signal `name` on `obj`.
"""
for func in _signals(obj, name):
func(*args, **kwargs)
def takes_econtext(func):
"""
Decorator that marks a function or class method to automatically receive a
kwarg named `econtext`, referencing the
:class:`mitogen.core.ExternalContext` active in the context in which the
function is being invoked in. The decorator is only meaningful when the
function is invoked via :data:`CALL_FUNCTION <mitogen.core.CALL_FUNCTION>`.
When the function is invoked directly, `econtext` must still be passed to
it explicitly.
"""
func.mitogen_takes_econtext = True
return func
def takes_router(func):
"""
Decorator that marks a function or class method to automatically receive a
kwarg named `router`, referencing the :class:`mitogen.core.Router` active
in the context in which the function is being invoked in. The decorator is
only meaningful when the function is invoked via :data:`CALL_FUNCTION
<mitogen.core.CALL_FUNCTION>`.
When the function is invoked directly, `router` must still be passed to it
explicitly.
"""
func.mitogen_takes_router = True
return func
def is_blacklisted_import(importer, fullname):
"""
Return :data:`True` if `fullname` is part of a blacklisted package, or if
any packages have been whitelisted and `fullname` is not part of one.
NB:
- If a package is on both lists, then it is treated as blacklisted.
- If any package is whitelisted, then all non-whitelisted packages are
treated as blacklisted.
"""
return ((not any(fullname.startswith(s) for s in importer.whitelist)) or
(any(fullname.startswith(s) for s in importer.blacklist)))
def set_cloexec(fd):
"""
Set the file descriptor `fd` to automatically close on :func:`os.execve`.
This has no effect on file descriptors inherited across :func:`os.fork`,
they must be explicitly closed through some other means, such as
:func:`mitogen.fork.on_fork`.
"""
flags = fcntl.fcntl(fd, fcntl.F_GETFD)
assert fd > 2, 'fd %r <= 2' % (fd,)
fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
def set_nonblock(fd):
"""
Set the file descriptor `fd` to non-blocking mode. For most underlying file
types, this causes :func:`os.read` or :func:`os.write` to raise
:class:`OSError` with :data:`errno.EAGAIN` rather than block the thread
when the underlying kernel buffer is exhausted.
"""
flags = fcntl.fcntl(fd, fcntl.F_GETFL)
fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
def set_block(fd):
"""
Inverse of :func:`set_nonblock`, i.e. cause `fd` to block the thread when
the underlying kernel buffer is exhausted.
"""
flags = fcntl.fcntl(fd, fcntl.F_GETFL)
fcntl.fcntl(fd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK)
def io_op(func, *args):
"""
Wrap `func(*args)` that may raise :class:`select.error`, :class:`IOError`,
or :class:`OSError`, trapping UNIX error codes relating to disconnection
and retry events in various subsystems:
* When a signal is delivered to the process on Python 2, system call retry
is signalled through :data:`errno.EINTR`. The invocation is automatically
restarted.
* When performing IO against a TTY, disconnection of the remote end is
signalled by :data:`errno.EIO`.
* When performing IO against a socket, disconnection of the remote end is
signalled by :data:`errno.ECONNRESET`.
* When performing IO against a pipe, disconnection of the remote end is
signalled by :data:`errno.EPIPE`.
:returns:
Tuple of `(return_value, disconnect_reason)`, where `return_value` is
the return value of `func(*args)`, and `disconnected` is an exception
instance when disconnection was detected, otherwise :data:`None`.
"""
while True:
try:
return func(*args), None
except (select.error, OSError, IOError):
e = sys.exc_info()[1]
_vv and IOLOG.debug('io_op(%r) -> OSError: %s', func, e)
if e.args[0] == errno.EINTR:
continue
if e.args[0] in (errno.EIO, errno.ECONNRESET, errno.EPIPE):
return None, e
raise
class PidfulStreamHandler(logging.StreamHandler):
"""
A :class:`logging.StreamHandler` subclass used when
:meth:`Router.enable_debug() <mitogen.master.Router.enable_debug>` has been
called, or the `debug` parameter was specified during context construction.
Verifies the process ID has not changed on each call to :meth:`emit`,
reopening the associated log file when a change is detected.
This ensures logging to the per-process output files happens correctly even
when uncooperative third party components call :func:`os.fork`.
"""
#: PID that last opened the log file.
open_pid = None
#: Output path template.
template = '/tmp/mitogen.%s.%s.log'
def _reopen(self):
self.acquire()
try:
if self.open_pid == os.getpid():
return
ts = time.strftime('%Y%m%d_%H%M%S')
path = self.template % (os.getpid(), ts)
self.stream = open(path, 'w', 1)
set_cloexec(self.stream.fileno())
self.stream.write('Parent PID: %s\n' % (os.getppid(),))
self.stream.write('Created by:\n\n%s\n' % (
''.join(traceback.format_stack()),
))
self.open_pid = os.getpid()
finally:
self.release()
def emit(self, record):
if self.open_pid != os.getpid():
self._reopen()
logging.StreamHandler.emit(self, record)
def enable_debug_logging():
global _v, _vv
_v = True
_vv = True
root = logging.getLogger()
root.setLevel(logging.DEBUG)
IOLOG.setLevel(logging.DEBUG)
handler = PidfulStreamHandler()
handler.formatter = logging.Formatter(
'%(asctime)s %(levelname).1s %(name)s: %(message)s',
'%H:%M:%S'
)
root.handlers.insert(0, handler)
_profile_hook = lambda name, func, *args: func(*args)
_profile_fmt = os.environ.get(
'MITOGEN_PROFILE_FMT',
'/tmp/mitogen.stats.%(pid)s.%(identity)s.%(now)s.%(ext)s',
)
def _profile_hook(name, func, *args):
"""
Call `func(*args)` and return its result. This function is replaced by
:func:`_real_profile_hook` when :func:`enable_profiling` is called. This
interface is obsolete and will be replaced by a signals-based integration
later on.
"""
return func(*args)
def _real_profile_hook(name, func, *args):
profiler = cProfile.Profile()
profiler.enable()
try:
return func(*args)
finally:
path = _profile_fmt % {
'now': int(1e6 * now()),
'identity': name,
'pid': os.getpid(),
'ext': '%s'
}
profiler.dump_stats(path % ('pstats',))
profiler.create_stats()
fp = open(path % ('log',), 'w')
try:
stats = pstats.Stats(profiler, stream=fp)
stats.sort_stats('cumulative')
stats.print_stats()
finally:
fp.close()
def enable_profiling(econtext=None):
global _profile_hook
_profile_hook = _real_profile_hook
def import_module(modname):
"""
Import `module` and return the attribute named `attr`.
"""
return __import__(modname, None, None, [''])
def pipe():
"""
Create a UNIX pipe pair using :func:`os.pipe`, wrapping the returned
descriptors in Python file objects in order to manage their lifetime and
ensure they are closed when their last reference is discarded and they have
not been closed explicitly.
"""
rfd, wfd = os.pipe()
return (
os.fdopen(rfd, 'rb', 0),
os.fdopen(wfd, 'wb', 0)
)
def iter_split(buf, delim, func):
"""
Invoke `func(s)` for each `delim`-delimited chunk in the potentially large
`buf`, avoiding intermediate lists and quadratic string operations. Return
the trailing undelimited portion of `buf`, or any unprocessed portion of
`buf` after `func(s)` returned :data:`False`.
:returns:
`(trailer, cont)`, where `cont` is :data:`False` if the last call to
`func(s)` returned :data:`False`.
"""
dlen = len(delim)
start = 0
cont = True
while cont:
nl = buf.find(delim, start)
if nl == -1:
break
cont = not func(buf[start:nl]) is False
start = nl + dlen
return buf[start:], cont
class Py24Pickler(py_pickle.Pickler):
"""
Exceptions were classic classes until Python 2.5. Sadly for 2.4, cPickle
offers little control over how a classic instance is pickled. Therefore 2.4
uses a pure-Python pickler, so CallError can be made to look as it does on
newer Pythons.
This mess will go away once proper serialization exists.
"""
@classmethod
def dumps(cls, obj, protocol):
bio = BytesIO()
self = cls(bio, protocol=protocol)
self.dump(obj)
return bio.getvalue()
def save_exc_inst(self, obj):
if isinstance(obj, CallError):
func, args = obj.__reduce__()
self.save(func)
self.save(args)
self.write(py_pickle.REDUCE)
else:
py_pickle.Pickler.save_inst(self, obj)
if PY24:
dispatch = py_pickle.Pickler.dispatch.copy()
dispatch[py_pickle.InstanceType] = save_exc_inst
if PY3:
# In 3.x Unpickler is a class exposing find_class as an overridable, but it
# cannot be overridden without subclassing.
class _Unpickler(pickle.Unpickler):
def find_class(self, module, func):
return self.find_global(module, func)
pickle__dumps = pickle.dumps
elif PY24:
# On Python 2.4, we must use a pure-Python pickler.
pickle__dumps = Py24Pickler.dumps
_Unpickler = pickle.Unpickler
else:
pickle__dumps = pickle.dumps
# In 2.x Unpickler is a function exposing a writeable find_global
# attribute.
_Unpickler = pickle.Unpickler
class Message(object):
"""
Messages are the fundamental unit of communication, comprising fields from
the :ref:`stream-protocol` header, an optional reference to the receiving
:class:`mitogen.core.Router` for ingress messages, and helper methods for
deserialization and generating replies.
"""
#: Integer target context ID. :class:`Router` delivers messages locally
#: when their :attr:`dst_id` matches :data:`mitogen.context_id`, otherwise
#: they are routed up or downstream.
dst_id = None
#: Integer source context ID. Used as the target of replies if any are
#: generated.
src_id = None
#: Context ID under whose authority the message is acting. See
#: :ref:`source-verification`.
auth_id = None
#: Integer target handle in the destination context. This is one of the
#: :ref:`standard-handles`, or a dynamically generated handle used to
#: receive a one-time reply, such as the return value of a function call.
handle = None
#: Integer target handle to direct any reply to this message. Used to
#: receive a one-time reply, such as the return value of a function call.
#: :data:`IS_DEAD` has a special meaning when it appears in this field.
reply_to = None
#: Raw message data bytes.
data = b('')
_unpickled = object()
#: The :class:`Router` responsible for routing the message. This is
#: :data:`None` for locally originated messages.
router = None
#: The :class:`Receiver` over which the message was last received. Part of
#: the :class:`mitogen.select.Select` interface. Defaults to :data:`None`.
receiver = None
HEADER_FMT = '>hLLLLLL'
HEADER_LEN = struct.calcsize(HEADER_FMT)
HEADER_MAGIC = 0x4d49 # 'MI'
def __init__(self, **kwargs):
"""
Construct a message from from the supplied `kwargs`. :attr:`src_id` and
:attr:`auth_id` are always set to :data:`mitogen.context_id`.
"""
self.src_id = mitogen.context_id
self.auth_id = mitogen.context_id
vars(self).update(kwargs)
assert isinstance(self.data, BytesType), 'Message data is not Bytes'
def pack(self):
return (
struct.pack(self.HEADER_FMT, self.HEADER_MAGIC, self.dst_id,
self.src_id, self.auth_id, self.handle,
self.reply_to or 0, len(self.data))
+ self.data
)
def _unpickle_context(self, 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)
def _unpickle_bytes(self, s, encoding):
s, n = LATIN1_CODEC.encode(s)
return s
def _find_global(self, module, func):
"""
Return the class implementing `module_name.class_name` or raise
`StreamError` if the module is not whitelisted.
"""
if module == __name__:
if func == '_unpickle_call_error' or func == 'CallError':
return _unpickle_call_error
elif func == '_unpickle_sender':
return self._unpickle_sender
elif func == '_unpickle_context':
return self._unpickle_context
elif func == 'Blob':
return Blob
elif func == 'Secret':
return Secret
elif func == 'Kwargs':
return Kwargs
elif module == '_codecs' and func == 'encode':
return self._unpickle_bytes
elif module == '__builtin__' and func == 'bytes':
return BytesType
raise StreamError('cannot unpickle %r/%r', module, func)
@property
def is_dead(self):
"""
:data:`True` if :attr:`reply_to` is set to the magic value
:data:`IS_DEAD`, indicating the sender considers the channel dead. Dead
messages can be raised in a variety of circumstances, see
:data:`IS_DEAD` for more information.
"""
return self.reply_to == IS_DEAD
@classmethod
def dead(cls, reason=None, **kwargs):
"""
Syntax helper to construct a dead message.
"""
kwargs['data'], _ = encodings.utf_8.encode(reason or u'')
return cls(reply_to=IS_DEAD, **kwargs)
@classmethod
def pickled(cls, obj, **kwargs):
"""
Construct a pickled message, setting :attr:`data` to the serialization
of `obj`, and setting remaining fields using `kwargs`.
:returns:
The new message.
"""
self = cls(**kwargs)
try:
self.data = pickle__dumps(obj, protocol=2)
except pickle.PicklingError:
e = sys.exc_info()[1]
self.data = pickle__dumps(CallError(e), protocol=2)
return self
def reply(self, msg, router=None, **kwargs):
"""
Compose a reply to this message and send it using :attr:`router`, or
`router` is :attr:`router` is :data:`None`.
:param obj:
Either a :class:`Message`, or an object to be serialized in order
to construct a new message.
:param router:
Optional router to use if :attr:`router` is :data:`None`.
:param kwargs:
Optional keyword parameters overriding message fields in the reply.
"""
if not isinstance(msg, Message):
msg = Message.pickled(msg)
msg.dst_id = self.src_id
msg.handle = self.reply_to
vars(msg).update(kwargs)
if msg.handle:
(self.router or router).route(msg)
else:
LOG.debug('dropping reply to message with no return address: %r',
msg)
if PY3:
UNPICKLER_KWARGS = {'encoding': 'bytes'}
else:
UNPICKLER_KWARGS = {}
def _throw_dead(self):
if len(self.data):
raise ChannelError(self.data.decode('utf-8', 'replace'))
elif self.src_id == mitogen.context_id:
raise ChannelError(ChannelError.local_msg)
else:
raise ChannelError(ChannelError.remote_msg)
def unpickle(self, throw=True, throw_dead=True):
"""
Unpickle :attr:`data`, optionally raising any exceptions present.
:param bool throw_dead:
If :data:`True`, raise exceptions, otherwise it is the caller's
responsibility.
:raises CallError:
The serialized data contained CallError exception.
:raises ChannelError:
The `is_dead` field was set.
"""
_vv and IOLOG.debug('%r.unpickle()', self)
if throw_dead and self.is_dead:
self._throw_dead()
obj = self._unpickled
if obj is Message._unpickled:
fp = BytesIO(self.data)
unpickler = _Unpickler(fp, **self.UNPICKLER_KWARGS)
unpickler.find_global = self._find_global
try:
# Must occur off the broker thread.
try:
obj = unpickler.load()
except:
LOG.error('raw pickle was: %r', self.data)
raise
self._unpickled = obj
except (TypeError, ValueError):
e = sys.exc_info()[1]
raise StreamError('invalid message: %s', e)
if throw:
if isinstance(obj, CallError):
raise obj
return obj
def __repr__(self):
return 'Message(%r, %r, %r, %r, %r, %r..%d)' % (
self.dst_id, self.src_id, self.auth_id, self.handle,
self.reply_to, (self.data or '')[:50], len(self.data)
)
class Sender(object):
"""
Senders are used to send pickled messages to a handle in another context,
it is the inverse of :class:`mitogen.core.Receiver`.
Senders may be serialized, making them convenient to wire up data flows.
See :meth:`mitogen.core.Receiver.to_sender` for more information.
:param mitogen.core.Context context:
Context to send messages to.
:param int dst_handle:
Destination handle to send messages to.
"""
def __init__(self, context, dst_handle):
self.context = context
self.dst_handle = dst_handle
def send(self, data):
"""
Send `data` to the remote end.
"""
_vv and IOLOG.debug('%r.send(%r..)', self, repr(data)[:100])
self.context.send(Message.pickled(data, handle=self.dst_handle))
explicit_close_msg = 'Sender was explicitly closed'
def close(self):
"""
Send a dead message to the remote, causing :meth:`ChannelError` to be
raised in any waiting thread.
"""
_vv and IOLOG.debug('%r.close()', self)
self.context.send(
Message.dead(
reason=self.explicit_close_msg,
handle=self.dst_handle
)
)
def __repr__(self):
return 'Sender(%r, %r)' % (self.context, self.dst_handle)
def __reduce__(self):
return _unpickle_sender, (self.context.context_id, self.dst_handle)
def _unpickle_sender(router, context_id, dst_handle):
if not (isinstance(router, Router) and
isinstance(context_id, (int, long)) and context_id >= 0 and
isinstance(dst_handle, (int, long)) and dst_handle > 0):
raise TypeError('cannot unpickle Sender: bad input or missing router')
return Sender(Context(router, context_id), dst_handle)
class Receiver(object):
"""
Receivers maintain a thread-safe queue of messages sent to a handle of this
context from another context.
:param mitogen.core.Router router:
Router to register the handler on.
:param int handle:
If not :data:`None`, an explicit handle to register, otherwise an
unused handle is chosen.
:param bool persist:
If :data:`False`, unregister the handler after one message is received.
Single-message receivers are intended for RPC-like transactions, such
as in the case of :meth:`mitogen.parent.Context.call_async`.
:param mitogen.core.Context respondent:
Context this receiver is receiving from. If not :data:`None`, arranges
for the receiver to receive a dead message if messages can no longer be
routed to the context due to disconnection, and ignores messages that
did not originate from the respondent context.
"""
#: If not :data:`None`, a function invoked as `notify(receiver)` after a
#: message has been received. The function is invoked on :class:`Broker`
#: thread, therefore it must not block. Used by
#: :class:`mitogen.select.Select` to efficiently implement waiting on
#: multiple event sources.
notify = None
raise_channelerror = True
def __init__(self, router, handle=None, persist=True,
respondent=None, policy=None, overwrite=False):
self.router = router
#: The handle.
self.handle = handle # Avoid __repr__ crash in add_handler()
self._latch = Latch() # Must exist prior to .add_handler()
self.handle = router.add_handler(
fn=self._on_receive,
handle=handle,
policy=policy,
persist=persist,
respondent=respondent,
overwrite=overwrite,
)
def __repr__(self):
return 'Receiver(%r, %r)' % (self.router, self.handle)
def __enter__(self):
return self
def __exit__(self, _1, _2, _3):
self.close()
def to_sender(self):
"""
Return a :class:`Sender` configured to deliver messages to this
receiver. As senders are serializable, this makes it convenient to pass
`(context_id, handle)` pairs around::
def deliver_monthly_report(sender):
for line in open('monthly_report.txt'):
sender.send(line)
sender.close()
@mitogen.main()
def main(router):
remote = router.ssh(hostname='mainframe')
recv = mitogen.core.Receiver(router)
remote.call(deliver_monthly_report, recv.to_sender())
for msg in recv:
print(msg)
"""
return Sender(self.router.myself(), self.handle)
def _on_receive(self, msg):
"""
Callback registered for the handle with :class:`Router`; appends data
to the internal queue.
"""
_vv and IOLOG.debug('%r._on_receive(%r)', self, msg)
self._latch.put(msg)
if self.notify:
self.notify(self)
closed_msg = 'the Receiver has been closed'
def close(self):
"""
Unregister the receiver's handle from its associated router, and cause
:class:`ChannelError` to be raised in any thread waiting in :meth:`get`
on this receiver.
"""
if self.handle:
self.router.del_handler(self.handle)
self.handle = None
self._latch.close()
def size(self):
"""
Return the number of items currently buffered.
As with :class:`Queue.Queue`, `0` may be returned even though a
subsequent call to :meth:`get` will succeed, since a message may be
posted at any moment between :meth:`size` and :meth:`get`.
As with :class:`Queue.Queue`, `>0` may be returned even though a
subsequent call to :meth:`get` will block, since another waiting thread
may be woken at any moment between :meth:`size` and :meth:`get`.
:raises LatchError:
The underlying latch has already been marked closed.
"""
return self._latch.size()
def empty(self):
"""
Return `size() == 0`.
.. deprecated:: 0.2.8
Use :meth:`size` instead.
:raises LatchError:
The latch has already been marked closed.
"""
return self._latch.empty()
def get(self, timeout=None, block=True, throw_dead=True):
"""
Sleep waiting for a message to arrive on this receiver.
:param float timeout:
If not :data:`None`, specifies a timeout in seconds.
:raises mitogen.core.ChannelError:
The remote end indicated the channel should be closed,
communication with it was lost, or :meth:`close` was called in the
local process.
:raises mitogen.core.TimeoutError:
Timeout was reached.
:returns:
:class:`Message` that was received.
"""
_vv and IOLOG.debug('%r.get(timeout=%r, block=%r)', self, timeout, block)
try:
msg = self._latch.get(timeout=timeout, block=block)
except LatchError:
raise ChannelError(self.closed_msg)
if msg.is_dead and throw_dead:
msg._throw_dead()
return msg
def __iter__(self):
"""
Yield consecutive :class:`Message` instances delivered to this receiver
until :class:`ChannelError` is raised.
"""
while True:
try:
msg = self.get()
except ChannelError:
return
yield msg
class Channel(Sender, Receiver):
"""
A channel inherits from :class:`mitogen.core.Sender` and
`mitogen.core.Receiver` to provide bidirectional functionality.
.. deprecated:: 0.2.0
This class is incomplete and obsolete, it will be removed in Mitogen
0.3.
Channels were an early attempt at syntax sugar. It is always easier to pass
around unidirectional pairs of senders/receivers, even though the syntax is
baroque:
.. literalinclude:: ../examples/ping_pong.py
Since all handles aren't known until after both ends are constructed, for
both ends to communicate through a channel, it is necessary for one end to
retrieve the handle allocated to the other and reconfigure its own channel
to match. Currently this is a manual task.
"""
def __init__(self, router, context, dst_handle, handle=None):
Sender.__init__(self, context, dst_handle)
Receiver.__init__(self, router, handle)
def close(self):
Receiver.close(self)
Sender.close(self)
def __repr__(self):
return 'Channel(%s, %s)' % (
Sender.__repr__(self),
Receiver.__repr__(self)
)
class Importer(object):
"""
Import protocol implementation that fetches modules from the parent
process.
:param context: Context to communicate via.
"""
# The Mitogen package is handled specially, since the child context must
# construct it manually during startup.
MITOGEN_PKG_CONTENT = [
'buildah',
'compat',
'debug',
'doas',
'docker',
'kubectl',
'fakessh',
'fork',
'jail',
'lxc',
'lxd',
'master',
'minify',
'os_fork',
'parent',
'select',
'service',
'setns',
'ssh',
'su',
'sudo',
'utils',
]
ALWAYS_BLACKLIST = [
# 2.x generates needless imports for 'builtins', while 3.x does the
# same for '__builtin__'. The correct one is built-in, the other always
# a negative round-trip.
'builtins',
'__builtin__',
'thread',
# org.python.core imported by copy, pickle, xml.sax; breaks Jython, but
# very unlikely to trigger a bug report.
'org',
]
if PY3:
ALWAYS_BLACKLIST += ['cStringIO']
def __init__(self, router, context, core_src, whitelist=(), blacklist=()):
self._log = logging.getLogger('mitogen.importer')
self._context = context
self._present = {'mitogen': self.MITOGEN_PKG_CONTENT}
self._lock = threading.Lock()
self.whitelist = list(whitelist) or ['']
self.blacklist = list(blacklist) + self.ALWAYS_BLACKLIST
# Preserve copies of the original server-supplied whitelist/blacklist
# for later use by children.
self.master_whitelist = self.whitelist[:]
self.master_blacklist = self.blacklist[:]
# Presence of an entry in this map indicates in-flight GET_MODULE.
self._callbacks = {}
self._cache = {}
if core_src:
self._update_linecache('x/mitogen/core.py', core_src)
self._cache['mitogen.core'] = (
'mitogen.core',
None,
'x/mitogen/core.py',
zlib.compress(core_src, 9),
[],
)
self._install_handler(router)
def _update_linecache(self, path, data):
"""
The Python 2.4 linecache module, used to fetch source code for
tracebacks and :func:`inspect.getsource`, does not support PEP-302,
meaning it needs extra help to for Mitogen-loaded modules. Directly
populate its cache if a loaded module belongs to the Mitogen package.
"""
if PY24 and 'mitogen' in path:
linecache.cache[path] = (
len(data),
0.0,
[line+'\n' for line in data.splitlines()],
path,
)
def _install_handler(self, router):
router.add_handler(
fn=self._on_load_module,
handle=LOAD_MODULE,
policy=has_parent_authority,
)
def __repr__(self):
return 'Importer'
def builtin_find_module(self, fullname):
# imp.find_module() will always succeed for __main__, because it is a
# built-in module. That means it exists on a special linked list deep
# within the bowels of the interpreter. We must special case it.
if fullname == '__main__':
raise ModuleNotFoundError()
parent, _, modname = str_rpartition(fullname, '.')
if parent:
path = sys.modules[parent].__path__
else:
path = None
fp, pathname, description = imp.find_module(modname, path)
if fp:
fp.close()
def find_module(self, fullname, path=None):
if hasattr(_tls, 'running'):
return None
_tls.running = True
try:
#_v and self._log.debug('Python requested %r', fullname)
fullname = to_text(fullname)
pkgname, dot, _ = str_rpartition(fullname, '.')
pkg = sys.modules.get(pkgname)
if pkgname and getattr(pkg, '__loader__', None) is not self:
self._log.debug('%s is submodule of a locally loaded package',
fullname)
return None
suffix = fullname[len(pkgname+dot):]
if pkgname and suffix not in self._present.get(pkgname, ()):
self._log.debug('%s has no submodule %s', pkgname, suffix)
return None
# #114: explicitly whitelisted prefixes override any
# system-installed package.
if self.whitelist != ['']:
if any(fullname.startswith(s) for s in self.whitelist):
return self
try:
self.builtin_find_module(fullname)
_vv and self._log.debug('%r is available locally', fullname)
except ImportError:
_vv and self._log.debug('we will try to load %r', fullname)
return self
finally:
del _tls.running
blacklisted_msg = (
'%r is present in the Mitogen importer blacklist, therefore this '
'context will not attempt to request it from the master, as the '
'request will always be refused.'
)
pkg_resources_msg = (
'pkg_resources is prohibited from importing __main__, as it causes '
'problems in applications whose main module is not designed to be '
're-imported by children.'
)
absent_msg = (
'The Mitogen master process was unable to serve %r. It may be a '
'native Python extension, or it may be missing entirely. Check the '
'importer debug logs on the master for more information.'
)
def _refuse_imports(self, fullname):
if is_blacklisted_import(self, fullname):
raise ModuleNotFoundError(self.blacklisted_msg % (fullname,))
f = sys._getframe(2)
requestee = f.f_globals['__name__']
if fullname == '__main__' and requestee == 'pkg_resources':
# Anything that imports pkg_resources will eventually cause
# pkg_resources to try and scan __main__ for its __requires__
# attribute (pkg_resources/__init__.py::_build_master()). This
# breaks any app that is not expecting its __main__ to suddenly be
# sucked over a network and injected into a remote process, like
# py.test.
raise ModuleNotFoundError(self.pkg_resources_msg)
if fullname == 'pbr':
# It claims to use pkg_resources to read version information, which
# would result in PEP-302 being used, but it actually does direct
# filesystem access. So instead smodge the environment to override
# any version that was defined. This will probably break something
# later.
os.environ['PBR_VERSION'] = '0.0.0'
def _on_load_module(self, msg):
if msg.is_dead:
return
tup = msg.unpickle()
fullname = tup[0]
_v and self._log.debug('received %s', fullname)
self._lock.acquire()
try:
self._cache[fullname] = tup
if tup[2] is not None and PY24:
self._update_linecache(
path='master:' + tup[2],
data=zlib.decompress(tup[3])
)
callbacks = self._callbacks.pop(fullname, [])
finally:
self._lock.release()
for callback in callbacks:
callback()
def _request_module(self, fullname, callback):
self._lock.acquire()
try:
present = fullname in self._cache
if not present:
funcs = self._callbacks.get(fullname)
if funcs is not None:
_v and self._log.debug('existing request for %s in flight',
fullname)
funcs.append(callback)
else:
_v and self._log.debug('sending new %s request to parent',
fullname)
self._callbacks[fullname] = [callback]
self._context.send(
Message(data=b(fullname), handle=GET_MODULE)
)
finally:
self._lock.release()
if present:
callback()
def load_module(self, fullname):
fullname = to_text(fullname)
_v and self._log.debug('requesting %s', fullname)
self._refuse_imports(fullname)
event = threading.Event()
self._request_module(fullname, event.set)
event.wait()
ret = self._cache[fullname]
if ret[2] is None:
raise ModuleNotFoundError(self.absent_msg % (fullname,))
pkg_present = ret[1]
mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
mod.__file__ = self.get_filename(fullname)
mod.__loader__ = self
if pkg_present is not None: # it's a package.
mod.__path__ = []
mod.__package__ = fullname
self._present[fullname] = pkg_present
else:
mod.__package__ = str_rpartition(fullname, '.')[0] or None
if mod.__package__ and not PY3:
# 2.x requires __package__ to be exactly a string.
mod.__package__, _ = encodings.utf_8.encode(mod.__package__)
source = self.get_source(fullname)
try:
code = compile(source, mod.__file__, 'exec', 0, 1)
except SyntaxError:
LOG.exception('while importing %r', fullname)
raise
if PY3:
exec(code, vars(mod))
else:
exec('exec code in vars(mod)')
# #590: if a module replaces itself in sys.modules during import, below
# is necessary. This matches PyImport_ExecCodeModuleEx()
return sys.modules.get(fullname, mod)
def get_filename(self, fullname):
if fullname in self._cache:
path = self._cache[fullname][2]
if path is None:
# If find_loader() returns self but a subsequent master RPC
# reveals the module can't be loaded, and so load_module()
# throws ImportError, on Python 3.x it is still possible for
# the loader to be called to fetch metadata.
raise ModuleNotFoundError(self.absent_msg % (fullname,))
return u'master:' + self._cache[fullname][2]
def get_source(self, fullname):
if fullname in self._cache:
compressed = self._cache[fullname][3]
if compressed is None:
raise ModuleNotFoundError(self.absent_msg % (fullname,))
source = zlib.decompress(self._cache[fullname][3])
if PY3:
return to_text(source)
return source
class LogHandler(logging.Handler):
"""
A :class:`logging.Handler` subclass that arranges for :data:`FORWARD_LOG`
messages to be sent to a parent context in response to logging messages
generated by the current context. This is installed by default in child
contexts during bootstrap, so that :mod:`logging` events can be viewed and
managed centrally in the master process.
The handler is initially *corked* after construction, such that it buffers
messages until :meth:`uncork` is called. This allows logging to be
installed prior to communication with the target being available, and
avoids any possible race where early log messages might be dropped.
:param mitogen.core.Context context:
The context to send log messages towards. At present this is always
the master process.
"""
def __init__(self, context):
logging.Handler.__init__(self)
self.context = context
self.local = threading.local()
self._buffer = []
# Private synchronization is needed while corked, to ensure no
# concurrent call to _send() exists during uncork().
self._buffer_lock = threading.Lock()
def uncork(self):
"""
#305: during startup :class:`LogHandler` may be installed before it is
possible to route messages, therefore messages are buffered until
:meth:`uncork` is called by :class:`ExternalContext`.
"""
self._buffer_lock.acquire()
try:
self._send = self.context.send
for msg in self._buffer:
self._send(msg)
self._buffer = None
finally:
self._buffer_lock.release()
def _send(self, msg):
self._buffer_lock.acquire()
try:
if self._buffer is None:
# uncork() may run concurrent to _send()
self._send(msg)
else:
self._buffer.append(msg)
finally:
self._buffer_lock.release()
def emit(self, rec):
"""
Send a :data:`FORWARD_LOG` message towards the target context.
"""
if rec.name == 'mitogen.io' or \
getattr(self.local, 'in_emit', False):
return
self.local.in_emit = True
try:
msg = self.format(rec)
encoded = '%s\x00%s\x00%s' % (rec.name, rec.levelno, msg)
if isinstance(encoded, UnicodeType):
# Logging package emits both :(
encoded = encoded.encode('utf-8')
self._send(Message(data=encoded, handle=FORWARD_LOG))
finally:
self.local.in_emit = False
class Stream(object):
"""
A :class:`Stream` is one readable and optionally one writeable file
descriptor (represented by :class:`Side`) aggregated alongside an
associated :class:`Protocol` that knows how to respond to IO readiness
events for those descriptors.
Streams are registered with :class:`Broker`, and callbacks are invoked on
the broker thread in response to IO activity. When registered using
:meth:`Broker.start_receive` or :meth:`Broker._start_transmit`, the broker
may call any of :meth:`on_receive`, :meth:`on_transmit`,
:meth:`on_shutdown` or :meth:`on_disconnect`.
It is expected that the :class:`Protocol` associated with a stream will
change over its life. For example during connection setup, the initial
protocol may be :class:`mitogen.parent.BootstrapProtocol` that knows how to
enter SSH and sudo passwords and transmit the :mod:`mitogen.core` source to
the target, before handing off to :class:`MitogenProtocol` when the target
process is initialized.
Streams connecting to children are in turn aggregated by
:class:`mitogen.parent.Connection`, which contains additional logic for
managing any child process, and a reference to any separate ``stderr``
:class:`Stream` connected to that process.
"""
#: A :class:`Side` representing the stream's receive file descriptor.
receive_side = None
#: A :class:`Side` representing the stream's transmit file descriptor.
transmit_side = None
#: A :class:`Protocol` representing the protocol active on the stream.
protocol = None
#: In parents, the :class:`mitogen.parent.Connection` instance.
conn = None
#: The stream name. This is used in the :meth:`__repr__` output in any log
#: messages, it may be any descriptive string.
name = u'default'
def set_protocol(self, protocol):
"""
Bind a :class:`Protocol` to this stream, by updating
:attr:`Protocol.stream` to refer to this stream, and updating this
stream's :attr:`Stream.protocol` to the refer to the protocol. Any
prior protocol's :attr:`Protocol.stream` is set to :data:`None`.
"""
if self.protocol:
self.protocol.stream = None
self.protocol = protocol
self.protocol.stream = self
def accept(self, rfp, wfp):
"""
Attach a pair of file objects to :attr:`receive_side` and
:attr:`transmit_side`, after wrapping them in :class:`Side` instances.
:class:`Side` will call :func:`set_nonblock` and :func:`set_cloexec`
on the underlying file descriptors during construction.
The same file object may be used for both sides. The default
:meth:`on_disconnect` is handles the possibility that only one
descriptor may need to be closed.
:param file rfp:
The file object to receive from.
:param file wfp:
The file object to transmit to.
"""
self.receive_side = Side(self, rfp)
self.transmit_side = Side(self, wfp)
def __repr__(self):
return "<Stream %s #%04x>" % (self.name, id(self) & 0xffff,)
def on_receive(self, broker):
"""
Invoked by :class:`Broker` when the stream's :attr:`receive_side` has
been marked readable using :meth:`Broker.start_receive` and the broker
has detected the associated file descriptor is ready for reading.
Subclasses must implement this if they are registered using
:meth:`Broker.start_receive`, and the method must invoke
:meth:`on_disconnect` if reading produces an empty string.
The default implementation reads :attr:`Protocol.read_size` bytes and
passes the resulting bytestring to :meth:`Protocol.on_receive`. If the
bytestring is 0 bytes, invokes :meth:`on_disconnect` instead.
"""
buf = self.receive_side.read(self.protocol.read_size)
if not buf:
LOG.debug('%r: empty read, disconnecting', self.receive_side)
return self.on_disconnect(broker)
self.protocol.on_receive(broker, buf)
def on_transmit(self, broker):
"""
Invoked by :class:`Broker` when the stream's :attr:`transmit_side` has
been marked writeable using :meth:`Broker._start_transmit` and the
broker has detected the associated file descriptor is ready for
writing.
Subclasses must implement they are ever registerd with
:meth:`Broker._start_transmit`.
The default implementation invokes :meth:`Protocol.on_transmit`.
"""
self.protocol.on_transmit(broker)
def on_shutdown(self, broker):
"""
Invoked by :meth:`Broker.shutdown` to allow the stream time to
gracefully shutdown.
The default implementation emits a ``shutdown`` signal before
invoking :meth:`on_disconnect`.
"""
fire(self, 'shutdown')
self.protocol.on_shutdown(broker)
def on_disconnect(self, broker):
"""
Invoked by :class:`Broker` to force disconnect the stream during
shutdown, invoked by the default :meth:`on_shutdown` implementation,
and usually invoked by any subclass :meth:`on_receive` implementation
in response to a 0-byte read.
The base implementation fires a ``disconnect`` event, then closes
:attr:`receive_side` and :attr:`transmit_side` after unregistering the
stream from the broker.
"""
fire(self, 'disconnect')
self.protocol.on_disconnect(broker)
class Protocol(object):
"""
Implement the program behaviour associated with activity on a
:class:`Stream`. The protocol in use may vary over a stream's life, for
example to allow :class:`mitogen.parent.BootstrapProtocol` to initialize
the connected child before handing it off to :class:`MitogenProtocol`. A
stream's active protocol is tracked in the :attr:`Stream.protocol`
attribute, and modified via :meth:`Stream.set_protocol`.
Protocols do not handle IO, they are entirely reliant on the interface
provided by :class:`Stream` and :class:`Side`, allowing the underlying IO
implementation to be replaced without modifying behavioural logic.
"""
stream_class = Stream
#: The :class:`Stream` this protocol is currently bound to, or
#: :data:`None`.
stream = None
#: The size of the read buffer used by :class:`Stream` when this is the
#: active protocol for the stream.
read_size = CHUNK_SIZE
@classmethod
def build_stream(cls, *args, **kwargs):
stream = cls.stream_class()
stream.set_protocol(cls(*args, **kwargs))
return stream
def __repr__(self):
return '%s(%s)' % (
self.__class__.__name__,
self.stream and self.stream.name,
)
def on_shutdown(self, broker):
_v and LOG.debug('%r: shutting down', self)
self.stream.on_disconnect(broker)
def on_disconnect(self, broker):
# Normally both sides an FD, so it is important that tranmit_side is
# deregistered from Poller before closing the receive side, as pollers
# like epoll and kqueue unregister all events on FD close, causing
# subsequent attempt to unregister the transmit side to fail.
LOG.debug('%r: disconnecting', self)
broker.stop_receive(self.stream)
if self.stream.transmit_side:
broker._stop_transmit(self.stream)
self.stream.receive_side.close()
if self.stream.transmit_side:
self.stream.transmit_side.close()
class DelimitedProtocol(Protocol):
"""
Provide a :meth:`Protocol.on_receive` implementation for protocols that are
delimited by a fixed string, like text based protocols. Each message is
passed to :meth:`on_line_received` as it arrives, with incomplete messages
passed to :meth:`on_partial_line_received`.
When emulating user input it is often necessary to respond to incomplete
lines, such as when a "Password: " prompt is sent.
:meth:`on_partial_line_received` may be called repeatedly with an
increasingly complete message. When a complete message is finally received,
:meth:`on_line_received` will be called once for it before the buffer is
discarded.
If :func:`on_line_received` returns :data:`False`, remaining data is passed
unprocessed to the stream's current protocol's :meth:`on_receive`. This
allows switching from line-oriented to binary while the input buffer
contains both kinds of data.
"""
#: The delimiter. Defaults to newline.
delimiter = b('\n')
_trailer = b('')
def on_receive(self, broker, buf):
_vv and IOLOG.debug('%r.on_receive()', self)
stream = self.stream
self._trailer, cont = mitogen.core.iter_split(
buf=self._trailer + buf,
delim=self.delimiter,
func=self.on_line_received,
)
if self._trailer:
if cont:
self.on_partial_line_received(self._trailer)
else:
assert stream.protocol is not self, \
'stream protocol is no longer %r' % (self,)
stream.protocol.on_receive(broker, self._trailer)
def on_line_received(self, line):
"""
Receive a line from the stream.
:param bytes line:
The encoded line, excluding the delimiter.
:returns:
:data:`False` to indicate this invocation modified the stream's
active protocol, and any remaining buffered data should be passed
to the new protocol's :meth:`on_receive` method.
Any other return value is ignored.
"""
pass
def on_partial_line_received(self, line):
"""
Receive a trailing unterminated partial line from the stream.
:param bytes line:
The encoded partial line.
"""
pass
class BufferedWriter(object):
"""
Implement buffered output while avoiding quadratic string operations. This
is currently constructed by each protocol, in future it may become fixed
for each stream instead.
"""
def __init__(self, broker, protocol):
self._broker = broker
self._protocol = protocol
self._buf = collections.deque()
self._len = 0
def write(self, s):
"""
Transmit `s` immediately, falling back to enqueuing it and marking the
stream writeable if no OS buffer space is available.
"""
if not self._len:
# Modifying epoll/Kqueue state is expensive, as are needless broker
# loops. Rather than wait for writeability, just write immediately,
# and fall back to the broker loop on error or full buffer.
try:
n = self._protocol.stream.transmit_side.write(s)
if n:
if n == len(s):
return
s = s[n:]
except OSError:
pass
self._broker._start_transmit(self._protocol.stream)
self._buf.append(s)
self._len += len(s)
def on_transmit(self, broker):
"""
Respond to stream writeability by retrying previously buffered
:meth:`write` calls.
"""
if self._buf:
buf = self._buf.popleft()
written = self._protocol.stream.transmit_side.write(buf)
if not written:
_v and LOG.debug('disconnected during write to %r', self)
self._protocol.stream.on_disconnect(broker)
return
elif written != len(buf):
self._buf.appendleft(BufferType(buf, written))
_vv and IOLOG.debug('transmitted %d bytes to %r', written, self)
self._len -= written
if not self._buf:
broker._stop_transmit(self._protocol.stream)
class Side(object):
"""
Represent one side of a :class:`Stream`. This allows unidirectional (e.g.
pipe) and bidirectional (e.g. socket) streams to operate identically.
Sides are also responsible for tracking the open/closed state of the
underlying FD, preventing erroneous duplicate calls to :func:`os.close` due
to duplicate :meth:`Stream.on_disconnect` calls, which would otherwise risk
silently succeeding by closing an unrelated descriptor. For this reason, it
is crucial only one file object exists per unique descriptor.
:param mitogen.core.Stream stream:
The stream this side is associated with.
:param object fp:
The file or socket object managing the underlying file descriptor. Any
object may be used that supports `fileno()` and `close()` methods.
:param bool cloexec:
If :data:`True`, the descriptor has its :data:`fcntl.FD_CLOEXEC` flag
enabled using :func:`fcntl.fcntl`.
:param bool keep_alive:
If :data:`True`, the continued existence of this side will extend the
shutdown grace period until it has been unregistered from the broker.
:param bool blocking:
If :data:`False`, the descriptor has its :data:`os.O_NONBLOCK` flag
enabled using :func:`fcntl.fcntl`.
"""
_fork_refs = weakref.WeakValueDictionary()
closed = False
def __init__(self, stream, fp, cloexec=True, keep_alive=True, blocking=False):
#: The :class:`Stream` for which this is a read or write side.
self.stream = stream
# File or socket object responsible for the lifetime of its underlying
# file descriptor.
self.fp = fp
#: Integer file descriptor to perform IO on, or :data:`None` if
#: :meth:`close` has been called. This is saved separately from the
#: file object, since :meth:`file.fileno` cannot be called on it after
#: it has been closed.
self.fd = fp.fileno()
#: If :data:`True`, causes presence of this side in
#: :class:`Broker`'s active reader set to defer shutdown until the
#: side is disconnected.
self.keep_alive = keep_alive
self._fork_refs[id(self)] = self
if cloexec:
set_cloexec(self.fd)
if not blocking:
set_nonblock(self.fd)
def __repr__(self):
return '<Side of %s fd %s>' % (
self.stream.name or repr(self.stream),
self.fd
)
@classmethod
def _on_fork(cls):
while cls._fork_refs:
_, side = cls._fork_refs.popitem()
_vv and IOLOG.debug('Side._on_fork() closing %r', side)
side.close()
def close(self):
"""
Call :meth:`file.close` on :attr:`fp` if it is not :data:`None`,
then set it to :data:`None`.
"""
_vv and IOLOG.debug('%r.close()', self)
if not self.closed:
self.closed = True
self.fp.close()
def read(self, n=CHUNK_SIZE):
"""
Read up to `n` bytes from the file descriptor, wrapping the underlying
:func:`os.read` call with :func:`io_op` to trap common disconnection
conditions.
:meth:`read` always behaves as if it is reading from a regular UNIX
file; socket, pipe, and TTY disconnection errors are masked and result
in a 0-sized read like a regular file.
:returns:
Bytes read, or the empty string to indicate disconnection was
detected.
"""
if self.closed:
# Refuse to touch the handle after closed, it may have been reused
# by another thread. TODO: synchronize read()/write()/close().
return b('')
s, disconnected = io_op(os.read, self.fd, n)
if disconnected:
LOG.debug('%r: disconnected during read: %s', self, disconnected)
return b('')
return s
def write(self, s):
"""
Write as much of the bytes from `s` as possible to the file descriptor,
wrapping the underlying :func:`os.write` call with :func:`io_op` to
trap common disconnection conditions.
:returns:
Number of bytes written, or :data:`None` if disconnection was
detected.
"""
if self.closed:
# Don't touch the handle after close, it may be reused elsewhere.
return None
written, disconnected = io_op(os.write, self.fd, s)
if disconnected:
LOG.debug('%r: disconnected during write: %s', self, disconnected)
return None
return written
class MitogenProtocol(Protocol):
"""
:class:`Protocol` implementing mitogen's :ref:`stream protocol
<stream-protocol>`.
"""
#: If not :data:`False`, indicates the stream has :attr:`auth_id` set and
#: its value is the same as :data:`mitogen.context_id` or appears in
#: :data:`mitogen.parent_ids`.
is_privileged = False
#: Invoked as `on_message(stream, msg)` each message received from the
#: peer.
on_message = None
def __init__(self, router, remote_id, auth_id=None,
local_id=None, parent_ids=None):
self._router = router
self.remote_id = remote_id
#: If not :data:`None`, :class:`Router` stamps this into
#: :attr:`Message.auth_id` of every message received on this stream.
self.auth_id = auth_id
if parent_ids is None:
parent_ids = mitogen.parent_ids
if local_id is None:
local_id = mitogen.context_id
self.is_privileged = (
(remote_id in parent_ids) or
auth_id in ([local_id] + parent_ids)
)
self.sent_modules = set(['mitogen', 'mitogen.core'])
self._input_buf = collections.deque()
self._input_buf_len = 0
self._writer = BufferedWriter(router.broker, self)
#: Routing records the dst_id of every message arriving from this
#: stream. Any arriving DEL_ROUTE is rebroadcast for any such ID.
self.egress_ids = set()
def on_receive(self, broker, buf):
"""
Handle the next complete message on the stream. Raise
:class:`StreamError` on failure.
"""
_vv and IOLOG.debug('%r.on_receive()', self)
if self._input_buf and self._input_buf_len < 128:
self._input_buf[0] += buf
else:
self._input_buf.append(buf)
self._input_buf_len += len(buf)
while self._receive_one(broker):
pass
corrupt_msg = (
'%s: Corruption detected: frame signature incorrect. This likely means'
' some external process is interfering with the connection. Received:'
'\n\n'
'%r'
)
def _receive_one(self, broker):
if self._input_buf_len < Message.HEADER_LEN:
return False
msg = Message()
msg.router = self._router
(magic, msg.dst_id, msg.src_id, msg.auth_id,
msg.handle, msg.reply_to, msg_len) = struct.unpack(
Message.HEADER_FMT,
self._input_buf[0][:Message.HEADER_LEN],
)
if magic != Message.HEADER_MAGIC:
LOG.error(self.corrupt_msg, self.stream.name, self._input_buf[0][:2048])
self.stream.on_disconnect(broker)
return False
if msg_len > self._router.max_message_size:
LOG.error('%r: Maximum message size exceeded (got %d, max %d)',
self, msg_len, self._router.max_message_size)
self.stream.on_disconnect(broker)
return False
total_len = msg_len + Message.HEADER_LEN
if self._input_buf_len < total_len:
_vv and IOLOG.debug(
'%r: Input too short (want %d, got %d)',
self, msg_len, self._input_buf_len - Message.HEADER_LEN
)
return False
start = Message.HEADER_LEN
prev_start = start
remain = total_len
bits = []
while remain:
buf = self._input_buf.popleft()
bit = buf[start:remain]
bits.append(bit)
remain -= len(bit) + start
prev_start = start
start = 0
msg.data = b('').join(bits)
self._input_buf.appendleft(buf[prev_start+len(bit):])
self._input_buf_len -= total_len
self._router._async_route(msg, self.stream)
return True
def pending_bytes(self):
"""
Return the number of bytes queued for transmission on this stream. This
can be used to limit the amount of data buffered in RAM by an otherwise
unlimited consumer.
For an accurate result, this method should be called from the Broker
thread, for example by using :meth:`Broker.defer_sync`.
"""
return self._writer._len
def on_transmit(self, broker):
"""
Transmit buffered messages.
"""
_vv and IOLOG.debug('%r.on_transmit()', self)
self._writer.on_transmit(broker)
def _send(self, msg):
_vv and IOLOG.debug('%r._send(%r)', self, msg)
self._writer.write(msg.pack())
def send(self, msg):
"""
Send `data` to `handle`, and tell the broker we have output. May be
called from any thread.
"""
self._router.broker.defer(self._send, msg)
def on_shutdown(self, broker):
"""
Disable :class:`Protocol` immediate disconnect behaviour.
"""
_v and LOG.debug('%r: shutting down', self)
class Context(object):
"""
Represent a remote context regardless of the underlying connection method.
Context objects are simple facades that emit messages through an
associated router, and have :ref:`signals` raised against them in response
to various events relating to the context.
**Note:** This is the somewhat limited core version, used by child
contexts. The master subclass is documented below this one.
Contexts maintain no internal state and are thread-safe.
Prefer :meth:`Router.context_by_id` over constructing context objects
explicitly, as that method is deduplicating, and returns the only context
instance :ref:`signals` will be raised on.
:param mitogen.core.Router router:
Router to emit messages through.
:param int context_id:
Context ID.
:param str name:
Context name.
"""
name = None
remote_name = None
def __init__(self, router, context_id, name=None):
self.router = router
self.context_id = context_id
if name:
self.name = to_text(name)
def __reduce__(self):
return _unpickle_context, (self.context_id, self.name)
def on_disconnect(self):
_v and LOG.debug('%r: disconnecting', self)
fire(self, 'disconnect')
def send_async(self, msg, persist=False):
"""
Arrange for `msg` to be delivered to this context, with replies
directed to a newly constructed receiver. :attr:`dst_id
<Message.dst_id>` is set to the target context ID, and :attr:`reply_to
<Message.reply_to>` is set to the newly constructed receiver's handle.
:param bool persist:
If :data:`False`, the handler will be unregistered after a single
message has been received.
:param mitogen.core.Message msg:
The message.
:returns:
:class:`Receiver` configured to receive any replies sent to the
message's `reply_to` handle.
"""
receiver = Receiver(self.router, persist=persist, respondent=self)
msg.dst_id = self.context_id
msg.reply_to = receiver.handle
_v and LOG.debug('sending message to %r: %r', self, msg)
self.send(msg)
return receiver
def call_service_async(self, service_name, method_name, **kwargs):
if isinstance(service_name, BytesType):
service_name = service_name.encode('utf-8')
elif not isinstance(service_name, UnicodeType):
service_name = service_name.name() # Service.name()
_v and LOG.debug('calling service %s.%s of %r, args: %r',
service_name, method_name, self, kwargs)
tup = (service_name, to_text(method_name), Kwargs(kwargs))
msg = Message.pickled(tup, handle=CALL_SERVICE)
return self.send_async(msg)
def send(self, msg):
"""
Arrange for `msg` to be delivered to this context. :attr:`dst_id
<Message.dst_id>` is set to the target context ID.
:param Message msg:
Message.
"""
msg.dst_id = self.context_id
self.router.route(msg)
def call_service(self, service_name, method_name, **kwargs):
recv = self.call_service_async(service_name, method_name, **kwargs)
return recv.get().unpickle()
def send_await(self, msg, deadline=None):
"""
Like :meth:`send_async`, but expect a single reply (`persist=False`)
delivered within `deadline` seconds.
:param mitogen.core.Message msg:
The message.
:param float deadline:
If not :data:`None`, seconds before timing out waiting for a reply.
:returns:
Deserialized reply.
:raises TimeoutError:
No message was received and `deadline` passed.
"""
receiver = self.send_async(msg)
response = receiver.get(deadline)
data = response.unpickle()
_vv and IOLOG.debug('%r._send_await() -> %r', self, data)
return data
def __repr__(self):
return 'Context(%s, %r)' % (self.context_id, self.name)
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')
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):
"""
A poller manages OS file descriptors the user is waiting to become
available for IO. The :meth:`poll` method blocks the calling thread
until one or more become ready. The default implementation is based on
:func:`select.poll`.
Each descriptor has an associated `data` element, which is unique for each
readiness type, and defaults to being the same as the file descriptor. The
:meth:`poll` method yields the data associated with a descriptor, rather
than the descriptor itself, allowing concise loops like::
p = Poller()
p.start_receive(conn.fd, data=conn.on_read)
p.start_transmit(conn.fd, data=conn.on_write)
for callback in p.poll():
callback() # invoke appropriate bound instance method
Pollers may be modified while :meth:`poll` is yielding results. Removals
are processed immediately, causing pending events for the descriptor to be
discarded.
The :meth:`close` method must be called when a poller is discarded to avoid
a resource leak.
Pollers may only be used by one thread at a time.
"""
SUPPORTED = True
# This changed from select() to poll() in Mitogen 0.2.4. Since poll() has
# no upper FD limit, it is suitable for use with Latch, which must handle
# FDs larger than select's limit during many-host runs. We want this
# because poll() requires no setup and teardown: just a single system call,
# which is important because Latch.get() creates a Poller on each
# invocation. In a microbenchmark, poll() vs. epoll_ctl() is 30% faster in
# this scenario. If select() must return in future, it is important
# Latch.poller_class is set from parent.py to point to the industrial
# strength poller for the OS, otherwise Latch will fail randomly.
#: Increments on every poll(). Used to version _rfds and _wfds.
_generation = 1
def __init__(self):
self._rfds = {}
self._wfds = {}
def __repr__(self):
return '%s' % (type(self).__name__,)
def _update(self, fd):
"""
Required by PollPoller subclass.
"""
pass
@property
def readers(self):
"""
Return a list of `(fd, data)` tuples for every FD registered for
receive readiness.
"""
return list((fd, data) for fd, (data, gen) in self._rfds.items())
@property
def writers(self):
"""
Return a list of `(fd, data)` tuples for every FD registered for
transmit readiness.
"""
return list((fd, data) for fd, (data, gen) in self._wfds.items())
def close(self):
"""
Close any underlying OS resource used by the poller.
"""
pass
def start_receive(self, fd, data=None):
"""
Cause :meth:`poll` to yield `data` when `fd` is readable.
"""
self._rfds[fd] = (data or fd, self._generation)
self._update(fd)
def stop_receive(self, fd):
"""
Stop yielding readability events for `fd`.
Redundant calls to :meth:`stop_receive` are silently ignored, this may
change in future.
"""
self._rfds.pop(fd, None)
self._update(fd)
def start_transmit(self, fd, data=None):
"""
Cause :meth:`poll` to yield `data` when `fd` is writeable.
"""
self._wfds[fd] = (data or fd, self._generation)
self._update(fd)
def stop_transmit(self, fd):
"""
Stop yielding writeability events for `fd`.
Redundant calls to :meth:`stop_transmit` are silently ignored, this may
change in future.
"""
self._wfds.pop(fd, None)
self._update(fd)
def _poll(self, timeout):
(rfds, wfds, _), _ = io_op(select.select,
self._rfds,
self._wfds,
(), timeout
)
for fd in rfds:
_vv and IOLOG.debug('%r: POLLIN for %r', self, fd)
data, gen = self._rfds.get(fd, (None, None))
if gen and gen < self._generation:
yield data
for fd in wfds:
_vv and IOLOG.debug('%r: POLLOUT for %r', self, fd)
data, gen = self._wfds.get(fd, (None, None))
if gen and gen < self._generation:
yield data
def poll(self, timeout=None):
"""
Block the calling thread until one or more FDs are ready for IO.
:param float timeout:
If not :data:`None`, seconds to wait without an event before
returning an empty iterable.
:returns:
Iterable of `data` elements associated with ready FDs.
"""
_vv and IOLOG.debug('%r.poll(%r)', self, timeout)
self._generation += 1
return self._poll(timeout)
class Latch(object):
"""
A latch is a :class:`Queue.Queue`-like object that supports mutation and
waiting from multiple threads, however unlike :class:`Queue.Queue`,
waiting threads always remain interruptible, so CTRL+C always succeeds, and
waits where a timeout is set experience no wake up latency. These
properties are not possible in combination using the built-in threading
primitives available in Python 2.x.
Latches implement queues using the UNIX self-pipe trick, and a per-thread
:func:`socket.socketpair` that is lazily created the first time any
latch attempts to sleep on a thread, and dynamically associated with the
waiting Latch only for duration of the wait.
See :ref:`waking-sleeping-threads` for further discussion.
"""
#: The :class:`Poller` implementation to use for waiting. Since the poller
#: will be very short-lived, we prefer :class:`mitogen.parent.PollPoller`
#: if it is available, or :class:`mitogen.core.Poller` otherwise, since
#: these implementations require no system calls to create, configure or
#: destroy.
poller_class = Poller
#: If not :data:`None`, a function invoked as `notify(latch)` after a
#: successful call to :meth:`put`. The function is invoked on the
#: :meth:`put` caller's thread, which may be the :class:`Broker` thread,
#: therefore it must not block. Used by :class:`mitogen.select.Select` to
#: efficiently implement waiting on multiple event sources.
notify = None
# The _cls_ prefixes here are to make it crystal clear in the code which
# state mutation isn't covered by :attr:`_lock`.
#: List of reusable :func:`socket.socketpair` tuples. The list is mutated
#: from multiple threads, the only safe operations are `append()` and
#: `pop()`.
_cls_idle_socketpairs = []
#: List of every socket object that must be closed by :meth:`_on_fork`.
#: Inherited descriptors cannot be reused, as the duplicated handles
#: reference the same underlying kernel object in use by the parent.
_cls_all_sockets = []
def __init__(self):
self.closed = False
self._lock = threading.Lock()
#: List of unconsumed enqueued items.
self._queue = []
#: List of `(wsock, cookie)` awaiting an element, where `wsock` is the
#: socketpair's write side, and `cookie` is the string to write.
self._sleeping = []
#: Number of elements of :attr:`_sleeping` that have already been
#: woken, and have a corresponding element index from :attr:`_queue`
#: assigned to them.
self._waking = 0
@classmethod
def _on_fork(cls):
"""
Clean up any files belonging to the parent process after a fork.
"""
cls._cls_idle_socketpairs = []
while cls._cls_all_sockets:
cls._cls_all_sockets.pop().close()
def close(self):
"""
Mark the latch as closed, and cause every sleeping thread to be woken,
with :class:`mitogen.core.LatchError` raised in each thread.
"""
self._lock.acquire()
try:
self.closed = True
while self._waking < len(self._sleeping):
wsock, cookie = self._sleeping[self._waking]
self._wake(wsock, cookie)
self._waking += 1
finally:
self._lock.release()
def size(self):
"""
Return the number of items currently buffered.
As with :class:`Queue.Queue`, `0` may be returned even though a
subsequent call to :meth:`get` will succeed, since a message may be
posted at any moment between :meth:`size` and :meth:`get`.
As with :class:`Queue.Queue`, `>0` may be returned even though a
subsequent call to :meth:`get` will block, since another waiting thread
may be woken at any moment between :meth:`size` and :meth:`get`.
:raises LatchError:
The latch has already been marked closed.
"""
self._lock.acquire()
try:
if self.closed:
raise LatchError()
return len(self._queue)
finally:
self._lock.release()
def empty(self):
"""
Return `size() == 0`.
.. deprecated:: 0.2.8
Use :meth:`size` instead.
:raises LatchError:
The latch has already been marked closed.
"""
return self.size() == 0
def _get_socketpair(self):
"""
Return an unused socketpair, creating one if none exist.
"""
try:
return self._cls_idle_socketpairs.pop() # pop() must be atomic
except IndexError:
rsock, wsock = socket.socketpair()
rsock.setblocking(False)
set_cloexec(rsock.fileno())
set_cloexec(wsock.fileno())
self._cls_all_sockets.extend((rsock, wsock))
return rsock, wsock
COOKIE_MAGIC, = struct.unpack('L', b('LTCH') * (struct.calcsize('L')//4))
COOKIE_FMT = '>Qqqq' # #545: id() and get_ident() may exceed long on armhfp.
COOKIE_SIZE = struct.calcsize(COOKIE_FMT)
def _make_cookie(self):
"""
Return a string encoding the ID of the process, instance and thread.
This disambiguates legitimate wake-ups, accidental writes to the FD,
and buggy internal FD sharing.
"""
return struct.pack(self.COOKIE_FMT, self.COOKIE_MAGIC,
os.getpid(), id(self), thread.get_ident())
def get(self, timeout=None, block=True):
"""
Return the next enqueued object, or sleep waiting for one.
:param float timeout:
If not :data:`None`, specifies a timeout in seconds.
:param bool block:
If :data:`False`, immediately raise
:class:`mitogen.core.TimeoutError` if the latch is empty.
:raises mitogen.core.LatchError:
:meth:`close` has been called, and the object is no longer valid.
:raises mitogen.core.TimeoutError:
Timeout was reached.
:returns:
The de-queued object.
"""
_vv and IOLOG.debug('%r.get(timeout=%r, block=%r)',
self, timeout, block)
self._lock.acquire()
try:
if self.closed:
raise LatchError()
i = len(self._sleeping)
if len(self._queue) > i:
_vv and IOLOG.debug('%r.get() -> %r', self, self._queue[i])
return self._queue.pop(i)
if not block:
raise TimeoutError()
rsock, wsock = self._get_socketpair()
cookie = self._make_cookie()
self._sleeping.append((wsock, cookie))
finally:
self._lock.release()
poller = self.poller_class()
poller.start_receive(rsock.fileno())
try:
return self._get_sleep(poller, timeout, block, rsock, wsock, cookie)
finally:
poller.close()
def _get_sleep(self, poller, timeout, block, rsock, wsock, cookie):
"""
When a result is not immediately available, sleep waiting for
:meth:`put` to write a byte to our socket pair.
"""
_vv and IOLOG.debug(
'%r._get_sleep(timeout=%r, block=%r, fd=%d/%d)',
self, timeout, block, rsock.fileno(), wsock.fileno()
)
e = None
try:
list(poller.poll(timeout))
except Exception:
e = sys.exc_info()[1]
self._lock.acquire()
try:
i = self._sleeping.index((wsock, cookie))
del self._sleeping[i]
try:
got_cookie = rsock.recv(self.COOKIE_SIZE)
except socket.error:
e2 = sys.exc_info()[1]
if e2.args[0] == errno.EAGAIN:
e = TimeoutError()
else:
e = e2
self._cls_idle_socketpairs.append((rsock, wsock))
if e:
raise e
assert cookie == got_cookie, (
"Cookie incorrect; got %r, expected %r" \
% (binascii.hexlify(got_cookie),
binascii.hexlify(cookie))
)
assert i < self._waking, (
"Cookie correct, but no queue element assigned."
)
self._waking -= 1
if self.closed:
raise LatchError()
_vv and IOLOG.debug('%r.get() wake -> %r', self, self._queue[i])
return self._queue.pop(i)
finally:
self._lock.release()
def put(self, obj=None):
"""
Enqueue an object, waking the first thread waiting for a result, if one
exists.
:param obj:
Object to enqueue. Defaults to :data:`None` as a convenience when
using :class:`Latch` only for synchronization.
:raises mitogen.core.LatchError:
:meth:`close` has been called, and the object is no longer valid.
"""
_vv and IOLOG.debug('%r.put(%r)', self, obj)
self._lock.acquire()
try:
if self.closed:
raise LatchError()
self._queue.append(obj)
wsock = None
if self._waking < len(self._sleeping):
wsock, cookie = self._sleeping[self._waking]
self._waking += 1
_vv and IOLOG.debug('%r.put() -> waking wfd=%r',
self, wsock.fileno())
elif self.notify:
self.notify(self)
finally:
self._lock.release()
if wsock:
self._wake(wsock, cookie)
def _wake(self, wsock, cookie):
written, disconnected = io_op(os.write, wsock.fileno(), cookie)
assert written == len(cookie) and not disconnected
def __repr__(self):
return 'Latch(%#x, size=%d, t=%r)' % (
id(self),
len(self._queue),
threading.currentThread().getName(),
)
class Waker(Protocol):
"""
:class:`Protocol` implementing the `UNIX self-pipe trick`_. Used to wake
:class:`Broker` when another thread needs to modify its state, by enqueing
a function call to run on the :class:`Broker` thread.
.. _UNIX self-pipe trick: https://cr.yp.to/docs/selfpipe.html
"""
read_size = 1
broker_ident = None
@classmethod
def build_stream(cls, broker):
stream = super(Waker, cls).build_stream(broker)
stream.accept(*pipe())
return stream
def __init__(self, broker):
self._broker = broker
self._deferred = collections.deque()
def __repr__(self):
return 'Waker(fd=%r/%r)' % (
self.stream.receive_side and self.stream.receive_side.fd,
self.stream.transmit_side and self.stream.transmit_side.fd,
)
@property
def keep_alive(self):
"""
Prevent immediate Broker shutdown while deferred functions remain.
"""
return len(self._deferred)
def on_receive(self, broker, buf):
"""
Drain the pipe and fire callbacks. Since :attr:`_deferred` is
synchronized, :meth:`defer` and :meth:`on_receive` can conspire to
ensure only one byte needs to be pending regardless of queue length.
"""
_vv and IOLOG.debug('%r.on_receive()', self)
while True:
try:
func, args, kwargs = self._deferred.popleft()
except IndexError:
return
try:
func(*args, **kwargs)
except Exception:
LOG.exception('defer() crashed: %r(*%r, **%r)',
func, args, kwargs)
broker.shutdown()
def _wake(self):
"""
Wake the multiplexer by writing a byte. If Broker is midway through
teardown, the FD may already be closed, so ignore EBADF.
"""
try:
self.stream.transmit_side.write(b(' '))
except OSError:
e = sys.exc_info()[1]
if e.args[0] in (errno.EBADF, errno.EWOULDBLOCK):
raise
broker_shutdown_msg = (
"An attempt was made to enqueue a message with a Broker that has "
"already exitted. It is likely your program called Broker.shutdown() "
"too early."
)
def defer(self, func, *args, **kwargs):
"""
Arrange for `func()` to execute on the broker thread. This function
returns immediately without waiting the result of `func()`. Use
:meth:`defer_sync` to block until a result is available.
:raises mitogen.core.Error:
:meth:`defer` was called after :class:`Broker` has begun shutdown.
"""
if thread.get_ident() == self.broker_ident:
_vv and IOLOG.debug('%r.defer() [immediate]', self)
return func(*args, **kwargs)
if self._broker._exitted:
raise Error(self.broker_shutdown_msg)
_vv and IOLOG.debug('%r.defer() [fd=%r]', self,
self.stream.transmit_side.fd)
self._deferred.append((func, args, kwargs))
self._wake()
class IoLoggerProtocol(DelimitedProtocol):
"""
Attached to one end of a socket pair whose other end overwrites one of the
standard ``stdout`` or ``stderr`` file descriptors in a child context.
Received data is split up into lines, decoded as UTF-8 and logged to the
:mod:`logging` package as either the ``stdout`` or ``stderr`` logger.
Logging in child contexts is in turn forwarded to the master process using
:class:`LogHandler`.
"""
@classmethod
def build_stream(cls, name, dest_fd):
"""
Even though the file descriptor `dest_fd` will hold the opposite end of
the socket open, we must keep a separate dup() of it (i.e. wsock) in
case some code decides to overwrite `dest_fd` later, which would
prevent break :meth:`on_shutdown` from calling :meth:`shutdown()
<socket.socket.shutdown>` on it.
"""
rsock, wsock = socket.socketpair()
os.dup2(wsock.fileno(), dest_fd)
stream = super(IoLoggerProtocol, cls).build_stream(name)
stream.name = name
stream.accept(rsock, wsock)
return stream
def __init__(self, name):
self._log = logging.getLogger(name)
# #453: prevent accidental log initialization in a child creating a
# feedback loop.
self._log.propagate = False
self._log.handlers = logging.getLogger().handlers[:]
def on_shutdown(self, broker):
"""
Shut down the write end of the socket, preventing any further writes to
it by this process, or subprocess that inherited it. This allows any
remaining kernel-buffered data to be drained during graceful shutdown
without the buffer continuously refilling due to some out of control
child process.
"""
_v and LOG.debug('%r: shutting down', self)
if not IS_WSL:
# #333: WSL generates invalid readiness indication on shutdown().
# This modifies the *kernel object* inherited by children, causing
# EPIPE on subsequent writes to any dupped FD in any process. The
# read side can then drain completely of prior buffered data.
self.stream.transmit_side.fp.shutdown(socket.SHUT_WR)
self.stream.transmit_side.close()
def on_line_received(self, line):
"""
Decode the received line as UTF-8 and pass it to the logging framework.
"""
self._log.info('%s', line.decode('utf-8', 'replace'))
class Router(object):
"""
Route messages between contexts, and invoke local handlers for messages
addressed to this context. :meth:`Router.route() <route>` straddles the
:class:`Broker` thread and user threads, it is safe to call anywhere.
**Note:** This is the somewhat limited core version of the Router class
used by child contexts. The master subclass is documented below this one.
"""
#: The :class:`mitogen.core.Context` subclass to use when constructing new
#: :class:`Context` objects in :meth:`myself` and :meth:`context_by_id`.
#: Permits :class:`Router` subclasses to extend the :class:`Context`
#: interface, as done in :class:`mitogen.parent.Router`.
context_class = Context
max_message_size = 128 * 1048576
#: When :data:`True`, permit children to only communicate with the current
#: context or a parent of the current context. Routing between siblings or
#: children of parents is prohibited, ensuring no communication is possible
#: between intentionally partitioned networks, such as when a program
#: simultaneously manipulates hosts spread across a corporate and a
#: production network, or production networks that are otherwise
#: air-gapped.
#:
#: Sending a prohibited message causes an error to be logged and a dead
#: message to be sent in reply to the errant message, if that message has
#: ``reply_to`` set.
#:
#: The value of :data:`unidirectional` becomes the default for the
#: :meth:`local() <mitogen.master.Router.local>` `unidirectional`
#: parameter.
unidirectional = False
duplicate_handle_msg = 'cannot register a handle that already exists'
refused_msg = 'refused by policy'
invalid_handle_msg = 'invalid handle'
too_large_msg = 'message too large (max %d bytes)'
respondent_disconnect_msg = 'the respondent Context has disconnected'
broker_exit_msg = 'Broker has exitted'
no_route_msg = 'no route to %r, my ID is %r'
unidirectional_msg = (
'routing mode prevents forward of message from context %d to '
'context %d via context %d'
)
def __init__(self, broker):
self.broker = broker
listen(broker, 'exit', self._on_broker_exit)
self._setup_logging()
self._write_lock = threading.Lock()
#: context ID -> Stream; must hold _write_lock to edit or iterate
self._stream_by_id = {}
#: List of contexts to notify of shutdown; must hold _write_lock
self._context_by_id = {}
self._last_handle = itertools.count(1000)
#: handle -> (persistent?, func(msg))
self._handle_map = {}
#: Context -> set { handle, .. }
self._handles_by_respondent = {}
self.add_handler(self._on_del_route, DEL_ROUTE)
def __repr__(self):
return 'Router(%r)' % (self.broker,)
def _setup_logging(self):
"""
This is done in the :class:`Router` constructor for historical reasons.
It must be called before ExternalContext logs its first messages, but
after logging has been setup. It must also be called when any router is
constructed for a consumer app.
"""
# Here seems as good a place as any.
global _v, _vv
_v = logging.getLogger().level <= logging.DEBUG
_vv = IOLOG.level <= logging.DEBUG
def _on_del_route(self, msg):
"""
Stub :data:`DEL_ROUTE` handler; fires 'disconnect' events on the
corresponding :attr:`_context_by_id` member. This is replaced by
:class:`mitogen.parent.RouteMonitor` in an upgraded context.
"""
if msg.is_dead:
return
target_id_s, _, name = bytes_partition(msg.data, b(':'))
target_id = int(target_id_s, 10)
LOG.error('%r: deleting route to %s (%d)',
self, to_text(name), target_id)
context = self._context_by_id.get(target_id)
if context:
fire(context, 'disconnect')
else:
LOG.debug('DEL_ROUTE for unknown ID %r: %r', target_id, msg)
def _on_stream_disconnect(self, stream):
notify = []
self._write_lock.acquire()
try:
for context in list(self._context_by_id.values()):
stream_ = self._stream_by_id.get(context.context_id)
if stream_ is stream:
del self._stream_by_id[context.context_id]
notify.append(context)
finally:
self._write_lock.release()
# Happens outside lock as e.g. RouteMonitor wants the same lock.
for context in notify:
context.on_disconnect()
def _on_broker_exit(self):
"""
Called prior to broker exit, informs callbacks registered with
:meth:`add_handler` the connection is dead.
"""
_v and LOG.debug('%r: broker has exitted', self)
while self._handle_map:
_, (_, func, _, _) = self._handle_map.popitem()
func(Message.dead(self.broker_exit_msg))
def myself(self):
"""
Return a :class:`Context` referring to the current process. Since
:class:`Context` is serializable, this is convenient to use in remote
function call parameter lists.
"""
return self.context_class(
router=self,
context_id=mitogen.context_id,
name='self',
)
def context_by_id(self, context_id, via_id=None, create=True, name=None):
"""
Return or construct a :class:`Context` given its ID. An internal
mapping of ID to the canonical :class:`Context` representing that ID,
so that :ref:`signals` can be raised.
This may be called from any thread, lookup and construction are atomic.
:param int context_id:
The context ID to look up.
:param int via_id:
If the :class:`Context` does not already exist, set its
:attr:`Context.via` to the :class:`Context` matching this ID.
:param bool create:
If the :class:`Context` does not already exist, create it.
:param str name:
If the :class:`Context` does not already exist, set its name.
:returns:
:class:`Context`, or return :data:`None` if `create` is
:data:`False` and no :class:`Context` previously existed.
"""
context = self._context_by_id.get(context_id)
if context:
return context
if create and via_id is not None:
via = self.context_by_id(via_id)
else:
via = None
self._write_lock.acquire()
try:
context = self._context_by_id.get(context_id)
if create and not context:
context = self.context_class(self, context_id, name=name)
context.via = via
self._context_by_id[context_id] = context
finally:
self._write_lock.release()
return context
def register(self, context, stream):
"""
Register a newly constructed context and its associated stream, and add
the stream's receive side to the I/O multiplexer. This method remains
public while the design has not yet settled.
"""
_v and LOG.debug('%s: registering %r to stream %r',
self, context, stream)
self._write_lock.acquire()
try:
self._stream_by_id[context.context_id] = stream
self._context_by_id[context.context_id] = context
finally:
self._write_lock.release()
self.broker.start_receive(stream)
listen(stream, 'disconnect', lambda: self._on_stream_disconnect(stream))
def stream_by_id(self, dst_id):
"""
Return the :class:`Stream` that should be used to communicate with
`dst_id`. If a specific route for `dst_id` is not known, a reference to
the parent context's stream is returned. If the parent is disconnected,
or when running in the master context, return :data:`None` instead.
This can be used from any thread, but its output is only meaningful
from the context of the :class:`Broker` thread, as disconnection or
replacement could happen in parallel on the broker thread at any
moment.
"""
return (
self._stream_by_id.get(dst_id) or
self._stream_by_id.get(mitogen.parent_id)
)
def del_handler(self, handle):
"""
Remove the handle registered for `handle`
:raises KeyError:
The handle wasn't registered.
"""
_, _, _, respondent = self._handle_map.pop(handle)
if respondent:
self._handles_by_respondent[respondent].discard(handle)
def add_handler(self, fn, handle=None, persist=True,
policy=None, respondent=None,
overwrite=False):
"""
Invoke `fn(msg)` on the :class:`Broker` thread for each Message sent to
`handle` from this context. Unregister after one invocation if
`persist` is :data:`False`. If `handle` is :data:`None`, a new handle
is allocated and returned.
:param int handle:
If not :data:`None`, an explicit handle to register, usually one of
the ``mitogen.core.*`` constants. If unspecified, a new unused
handle will be allocated.
:param bool persist:
If :data:`False`, the handler will be unregistered after a single
message has been received.
:param mitogen.core.Context respondent:
Context that messages to this handle are expected to be sent from.
If specified, arranges for a dead message to be delivered to `fn`
when disconnection of the context is detected.
In future `respondent` will likely also be used to prevent other
contexts from sending messages to the handle.
:param function policy:
Function invoked as `policy(msg, stream)` where `msg` is a
:class:`mitogen.core.Message` about to be delivered, and `stream`
is the :class:`mitogen.core.Stream` on which it was received. The
function must return :data:`True`, otherwise an error is logged and
delivery is refused.
Two built-in policy functions exist:
* :func:`has_parent_authority`: requires the message arrived from a
parent context, or a context acting with a parent context's
authority (``auth_id``).
* :func:`mitogen.parent.is_immediate_child`: requires the
message arrived from an immediately connected child, for use in
messaging patterns where either something becomes buggy or
insecure by permitting indirect upstream communication.
In case of refusal, and the message's ``reply_to`` field is
nonzero, a :class:`mitogen.core.CallError` is delivered to the
sender indicating refusal occurred.
:param bool overwrite:
If :data:`True`, allow existing handles to be silently overwritten.
:return:
`handle`, or if `handle` was :data:`None`, the newly allocated
handle.
:raises Error:
Attemp to register handle that was already registered.
"""
handle = handle or next(self._last_handle)
_vv and IOLOG.debug('%r.add_handler(%r, %r, %r)', self, fn, handle, persist)
if handle in self._handle_map and not overwrite:
raise Error(self.duplicate_handle_msg)
self._handle_map[handle] = persist, fn, policy, respondent
if respondent:
if respondent not in self._handles_by_respondent:
self._handles_by_respondent[respondent] = set()
listen(respondent, 'disconnect',
lambda: self._on_respondent_disconnect(respondent))
self._handles_by_respondent[respondent].add(handle)
return handle
def _on_respondent_disconnect(self, context):
for handle in self._handles_by_respondent.pop(context, ()):
_, fn, _, _ = self._handle_map[handle]
fn(Message.dead(self.respondent_disconnect_msg))
del self._handle_map[handle]
def _maybe_send_dead(self, unreachable, msg, reason, *args):
"""
Send a dead message to either the original sender or the intended
recipient of `msg`, if the original sender was expecting a reply
(because its `reply_to` was set), otherwise assume the message is a
reply of some sort, and send the dead message to the original
destination.
:param bool unreachable:
If :data:`True`, the recipient is known to be dead or routing
failed due to a security precaution, so don't attempt to fallback
to sending the dead message to the recipient if the original sender
did not include a reply address.
:param mitogen.core.Message msg:
Message that triggered the dead message.
:param str reason:
Human-readable error reason.
:param tuple args:
Elements to interpolate with `reason`.
"""
if args:
reason %= args
LOG.debug('%r: %r is dead: %r', self, msg, reason)
if msg.reply_to and not msg.is_dead:
msg.reply(Message.dead(reason=reason), router=self)
elif not unreachable:
self._async_route(
Message.dead(
dst_id=msg.dst_id,
handle=msg.handle,
reason=reason,
)
)
def _invoke(self, msg, stream):
# IOLOG.debug('%r._invoke(%r)', self, msg)
try:
persist, fn, policy, respondent = self._handle_map[msg.handle]
except KeyError:
self._maybe_send_dead(True, msg, reason=self.invalid_handle_msg)
return
if respondent and not (msg.is_dead or
msg.src_id == respondent.context_id):
self._maybe_send_dead(True, msg, 'reply from unexpected context')
return
if policy and not policy(msg, stream):
self._maybe_send_dead(True, msg, self.refused_msg)
return
if not persist:
self.del_handler(msg.handle)
try:
fn(msg)
except Exception:
LOG.exception('%r._invoke(%r): %r crashed', self, msg, fn)
def _async_route(self, msg, in_stream=None):
"""
Arrange for `msg` to be forwarded towards its destination. If its
destination is the local context, then arrange for it to be dispatched
using the local handlers.
This is a lower overhead version of :meth:`route` that may only be
called from the :class:`Broker` thread.
:param Stream in_stream:
If not :data:`None`, the stream the message arrived on. Used for
performing source route verification, to ensure sensitive messages
such as ``CALL_FUNCTION`` arrive only from trusted contexts.
"""
_vv and IOLOG.debug('%r._async_route(%r, %r)', self, msg, in_stream)
if len(msg.data) > self.max_message_size:
self._maybe_send_dead(False, msg, self.too_large_msg % (
self.max_message_size,
))
return
parent_stream = self._stream_by_id.get(mitogen.parent_id)
src_stream = self._stream_by_id.get(msg.src_id, parent_stream)
# When the ingress stream is known, verify the message was received on
# the same as the stream we would expect to receive messages from the
# src_id and auth_id. This is like Reverse Path Filtering in IP, and
# ensures messages from a privileged context cannot be spoofed by a
# child.
if in_stream:
auth_stream = self._stream_by_id.get(msg.auth_id, parent_stream)
if in_stream != auth_stream:
LOG.error('%r: bad auth_id: got %r via %r, not %r: %r',
self, msg.auth_id, in_stream, auth_stream, msg)
return
if msg.src_id != msg.auth_id and in_stream != src_stream:
LOG.error('%r: bad src_id: got %r via %r, not %r: %r',
self, msg.src_id, in_stream, src_stream, msg)
return
# If the stream's MitogenProtocol has auth_id set, copy it to the
# message. This allows subtrees to become privileged by stamping a
# parent's context ID. It is used by mitogen.unix to mark client
# streams (like Ansible WorkerProcess) as having the same rights as
# the parent.
if in_stream.protocol.auth_id is not None:
msg.auth_id = in_stream.protocol.auth_id
if in_stream.protocol.on_message is not None:
in_stream.protocol.on_message(in_stream, msg)
# Record the IDs the source ever communicated with.
in_stream.protocol.egress_ids.add(msg.dst_id)
if msg.dst_id == mitogen.context_id:
return self._invoke(msg, in_stream)
out_stream = self._stream_by_id.get(msg.dst_id)
if (not out_stream) and (parent_stream != src_stream or not in_stream):
# No downstream route exists. The message could be from a child or
# ourselves for a parent, in which case we must forward it
# upstream, or it could be from a parent for a dead child, in which
# case its src_id/auth_id would fail verification if returned to
# the parent, so in that case reply with a dead message instead.
out_stream = parent_stream
if out_stream is None:
self._maybe_send_dead(True, msg, self.no_route_msg,
msg.dst_id, mitogen.context_id)
return
if in_stream and self.unidirectional and not \
(in_stream.protocol.is_privileged or
out_stream.protocol.is_privileged):
self._maybe_send_dead(True, msg, self.unidirectional_msg,
in_stream.protocol.remote_id,
out_stream.protocol.remote_id,
mitogen.context_id)
return
out_stream.protocol._send(msg)
def route(self, msg):
"""
Arrange for the :class:`Message` `msg` to be delivered to its
destination using any relevant downstream context, or if none is found,
by forwarding the message upstream towards the master context. If `msg`
is destined for the local context, it is dispatched using the handles
registered with :meth:`add_handler`.
This may be called from any thread.
"""
self.broker.defer(self._async_route, msg)
class NullTimerList(object):
def get_timeout(self):
return None
class Broker(object):
"""
Responsible for handling I/O multiplexing in a private thread.
**Note:** This somewhat limited core version is used by children. The
master subclass is documented below.
"""
poller_class = Poller
_waker = None
_thread = None
# :func:`mitogen.parent._upgrade_broker` replaces this with
# :class:`mitogen.parent.TimerList` during upgrade.
timers = NullTimerList()
#: Seconds grace to allow :class:`streams <Stream>` to shutdown gracefully
#: before force-disconnecting them during :meth:`shutdown`.
shutdown_timeout = 3.0
def __init__(self, poller_class=None, activate_compat=True):
self._alive = True
self._exitted = False
self._waker = Waker.build_stream(self)
#: Arrange for `func(\*args, \**kwargs)` to be executed on the broker
#: thread, or immediately if the current thread is the broker thread.
#: Safe to call from any thread.
self.defer = self._waker.protocol.defer
self.poller = self.poller_class()
self.poller.start_receive(
self._waker.receive_side.fd,
(self._waker.receive_side, self._waker.on_receive)
)
self._thread = threading.Thread(
target=self._broker_main,
name='mitogen.broker'
)
self._thread.start()
if activate_compat:
self._py24_25_compat()
def _py24_25_compat(self):
"""
Python 2.4/2.5 have grave difficulties with threads/fork. We
mandatorily quiesce all running threads during fork using a
monkey-patch there.
"""
if sys.version_info < (2, 6):
# import_module() is used to avoid dep scanner.
os_fork = import_module('mitogen.os_fork')
os_fork._notice_broker_or_pool(self)
def start_receive(self, stream):
"""
Mark the :attr:`receive_side <Stream.receive_side>` on `stream` as
ready for reading. Safe to call from any thread. When the associated
file descriptor becomes ready for reading,
:meth:`BasicStream.on_receive` will be called.
"""
_vv and IOLOG.debug('%r.start_receive(%r)', self, stream)
side = stream.receive_side
assert side and not side.closed
self.defer(self.poller.start_receive,
side.fd, (side, stream.on_receive))
def stop_receive(self, stream):
"""
Mark the :attr:`receive_side <Stream.receive_side>` on `stream` as not
ready for reading. Safe to call from any thread.
"""
_vv and IOLOG.debug('%r.stop_receive(%r)', self, stream)
self.defer(self.poller.stop_receive, stream.receive_side.fd)
def _start_transmit(self, stream):
"""
Mark the :attr:`transmit_side <Stream.transmit_side>` on `stream` as
ready for writing. Must only be called from the Broker thread. When the
associated file descriptor becomes ready for writing,
:meth:`BasicStream.on_transmit` will be called.
"""
_vv and IOLOG.debug('%r._start_transmit(%r)', self, stream)
side = stream.transmit_side
assert side and not side.closed
self.poller.start_transmit(side.fd, (side, stream.on_transmit))
def _stop_transmit(self, stream):
"""
Mark the :attr:`transmit_side <Stream.receive_side>` on `stream` as not
ready for writing.
"""
_vv and IOLOG.debug('%r._stop_transmit(%r)', self, stream)
self.poller.stop_transmit(stream.transmit_side.fd)
def keep_alive(self):
"""
Return :data:`True` if any reader's :attr:`Side.keep_alive` attribute
is :data:`True`, or any :class:`Context` is still registered that is
not the master. Used to delay shutdown while some important work is in
progress (e.g. log draining).
"""
it = (side.keep_alive for (_, (side, _)) in self.poller.readers)
return sum(it, 0) > 0 or self.timers.get_timeout() is not None
def defer_sync(self, func):
"""
Arrange for `func()` to execute on :class:`Broker` thread, blocking the
current thread until a result or exception is available.
:returns:
Return value of `func()`.
"""
latch = Latch()
def wrapper():
try:
latch.put(func())
except Exception:
latch.put(sys.exc_info()[1])
self.defer(wrapper)
res = latch.get()
if isinstance(res, Exception):
raise res
return res
def _call(self, stream, func):
"""
Call `func(self)`, catching any exception that might occur, logging it,
and force-disconnecting the related `stream`.
"""
try:
func(self)
except Exception:
LOG.exception('%r crashed', stream)
stream.on_disconnect(self)
def _loop_once(self, timeout=None):
"""
Execute a single :class:`Poller` wait, dispatching any IO events that
caused the wait to complete.
:param float timeout:
If not :data:`None`, maximum time in seconds to wait for events.
"""
_vv and IOLOG.debug('%r._loop_once(%r, %r)',
self, timeout, self.poller)
timer_to = self.timers.get_timeout()
if timeout is None:
timeout = timer_to
elif timer_to is not None and timer_to < timeout:
timeout = timer_to
#IOLOG.debug('readers =\n%s', pformat(self.poller.readers))
#IOLOG.debug('writers =\n%s', pformat(self.poller.writers))
for side, func in self.poller.poll(timeout):
self._call(side.stream, func)
if timer_to is not None:
self.timers.expire()
def _broker_exit(self):
"""
Forcefully call :meth:`Stream.on_disconnect` on any streams that failed
to shut down gracefully, then discard the :class:`Poller`.
"""
for _, (side, _) in self.poller.readers + self.poller.writers:
LOG.debug('%r: force disconnecting %r', self, side)
side.stream.on_disconnect(self)
self.poller.close()
def _broker_shutdown(self):
"""
Invoke :meth:`Stream.on_shutdown` for every active stream, then allow
up to :attr:`shutdown_timeout` seconds for the streams to unregister
themselves, logging an error if any did not unregister during the grace
period.
"""
for _, (side, _) in self.poller.readers + self.poller.writers:
self._call(side.stream, side.stream.on_shutdown)
deadline = now() + self.shutdown_timeout
while self.keep_alive() and now() < deadline:
self._loop_once(max(0, deadline - now()))
if self.keep_alive():
LOG.error('%r: pending work still existed %d seconds after '
'shutdown began. This may be due to a timer that is yet '
'to expire, or a child connection that did not fully '
'shut down.', self, self.shutdown_timeout)
def _do_broker_main(self):
"""
Broker thread main function. Dispatches IO events until
:meth:`shutdown` is called.
"""
# For Python 2.4, no way to retrieve ident except on thread.
self._waker.protocol.broker_ident = thread.get_ident()
try:
while self._alive:
self._loop_once()
fire(self, 'before_shutdown')
fire(self, 'shutdown')
self._broker_shutdown()
except Exception:
e = sys.exc_info()[1]
LOG.exception('broker crashed')
syslog.syslog(syslog.LOG_ERR, 'broker crashed: %s' % (e,))
syslog.closelog() # prevent test 'fd leak'.
self._alive = False # Ensure _alive is consistent on crash.
self._exitted = True
self._broker_exit()
def _broker_main(self):
try:
_profile_hook('mitogen.broker', self._do_broker_main)
finally:
# 'finally' to ensure _on_broker_exit() can always SIGTERM.
fire(self, 'exit')
def shutdown(self):
"""
Request broker gracefully disconnect streams and stop. Safe to call
from any thread.
"""
_v and LOG.debug('%r: shutting down', self)
def _shutdown():
self._alive = False
if self._alive and not self._exitted:
self.defer(_shutdown)
def join(self):
"""
Wait for the broker to stop, expected to be called after
:meth:`shutdown`.
"""
self._thread.join()
def __repr__(self):
return 'Broker(%04x)' % (id(self) & 0xffff,)
class Dispatcher(object):
"""
Implementation of the :data:`CALL_FUNCTION` handle for a child context.
Listens on the child's main thread for messages sent by
:class:`mitogen.parent.CallChain` and dispatches the function calls they
describe.
If a :class:`mitogen.parent.CallChain` sending a message is in pipelined
mode, any exception that occurs is recorded, and causes all subsequent
calls with the same `chain_id` to fail with the same exception.
"""
_service_recv = None
def __repr__(self):
return 'Dispatcher'
def __init__(self, econtext):
self.econtext = econtext
#: Chain ID -> CallError if prior call failed.
self._error_by_chain_id = {}
self.recv = Receiver(
router=econtext.router,
handle=CALL_FUNCTION,
policy=has_parent_authority,
)
#: The :data:`CALL_SERVICE` :class:`Receiver` that will eventually be
#: reused by :class:`mitogen.service.Pool`, should it ever be loaded.
#: This is necessary for race-free reception of all service requests
#: delivered regardless of whether the stub or real service pool are
#: loaded. See #547 for related sorrows.
Dispatcher._service_recv = Receiver(
router=econtext.router,
handle=CALL_SERVICE,
policy=has_parent_authority,
)
self._service_recv.notify = self._on_call_service
listen(econtext.broker, 'shutdown', self._on_broker_shutdown)
def _on_broker_shutdown(self):
if self._service_recv.notify == self._on_call_service:
self._service_recv.notify = None
self.recv.close()
@classmethod
@takes_econtext
def forget_chain(cls, chain_id, econtext):
econtext.dispatcher._error_by_chain_id.pop(chain_id, None)
def _parse_request(self, msg):
data = msg.unpickle(throw=False)
_v and LOG.debug('%r: dispatching %r', self, data)
chain_id, modname, klass, func, args, kwargs = data
obj = import_module(modname)
if klass:
obj = getattr(obj, klass)
fn = getattr(obj, func)
if getattr(fn, 'mitogen_takes_econtext', None):
kwargs.setdefault('econtext', self.econtext)
if getattr(fn, 'mitogen_takes_router', None):
kwargs.setdefault('router', self.econtext.router)
return chain_id, fn, args, kwargs
def _dispatch_one(self, msg):
try:
chain_id, fn, args, kwargs = self._parse_request(msg)
except Exception:
return None, CallError(sys.exc_info()[1])
if chain_id in self._error_by_chain_id:
return chain_id, self._error_by_chain_id[chain_id]
try:
return chain_id, fn(*args, **kwargs)
except Exception:
e = CallError(sys.exc_info()[1])
if chain_id is not None:
self._error_by_chain_id[chain_id] = e
return chain_id, e
def _on_call_service(self, recv):
"""
Notifier for the :data:`CALL_SERVICE` receiver. This is called on the
:class:`Broker` thread for any service messages arriving at this
context, for as long as no real service pool implementation is loaded.
In order to safely bootstrap the service pool implementation a sentinel
message is enqueued on the :data:`CALL_FUNCTION` receiver in order to
wake the main thread, where the importer can run without any
possibility of suffering deadlock due to concurrent uses of the
importer.
Should the main thread be blocked indefinitely, preventing the import
from ever running, if it is blocked waiting on a service call, then it
means :mod:`mitogen.service` has already been imported and
:func:`mitogen.service.get_or_create_pool` has already run, meaning the
service pool is already active and the duplicate initialization was not
needed anyway.
#547: This trickery is needed to avoid the alternate option of spinning
a temporary thread to import the service pool, which could deadlock if
a custom import hook executing on the main thread (under the importer
lock) would block waiting for some data that was in turn received by a
service. Main thread import lock can't be released until service is
running, service cannot satisfy request until import lock is released.
"""
self.recv._on_receive(Message(handle=STUB_CALL_SERVICE))
def _init_service_pool(self):
import mitogen.service
mitogen.service.get_or_create_pool(router=self.econtext.router)
def _dispatch_calls(self):
for msg in self.recv:
if msg.handle == STUB_CALL_SERVICE:
if msg.src_id == mitogen.context_id:
self._init_service_pool()
continue
chain_id, ret = self._dispatch_one(msg)
_v and LOG.debug('%r: %r -> %r', self, msg, ret)
if msg.reply_to:
msg.reply(ret)
elif isinstance(ret, CallError) and chain_id is None:
LOG.error('No-reply function call failed: %s', ret)
def run(self):
if self.econtext.config.get('on_start'):
self.econtext.config['on_start'](self.econtext)
_profile_hook('mitogen.child_main', self._dispatch_calls)
class ExternalContext(object):
"""
External context implementation.
This class contains the main program implementation for new children. It is
responsible for setting up everything about the process environment, import
hooks, standard IO redirection, logging, configuring a :class:`Router` and
:class:`Broker`, and finally arranging for :class:`Dispatcher` to take over
the main thread after initialization is complete.
.. attribute:: broker
The :class:`mitogen.core.Broker` instance.
.. attribute:: context
The :class:`mitogen.core.Context` instance.
.. attribute:: channel
The :class:`mitogen.core.Channel` over which :data:`CALL_FUNCTION`
requests are received.
.. attribute:: importer
The :class:`mitogen.core.Importer` instance.
.. attribute:: stdout_log
The :class:`IoLogger` connected to :data:`sys.stdout`.
.. attribute:: stderr_log
The :class:`IoLogger` connected to :data:`sys.stderr`.
"""
detached = False
def __init__(self, config):
self.config = config
def _on_broker_exit(self):
if not self.config['profiling']:
os.kill(os.getpid(), signal.SIGTERM)
def _on_shutdown_msg(self, msg):
if not msg.is_dead:
_v and LOG.debug('shutdown request from context %d', msg.src_id)
self.broker.shutdown()
def _on_parent_disconnect(self):
if self.detached:
mitogen.parent_ids = []
mitogen.parent_id = None
LOG.info('Detachment complete')
else:
_v and LOG.debug('parent stream is gone, dying.')
self.broker.shutdown()
def detach(self):
self.detached = True
stream = self.router.stream_by_id(mitogen.parent_id)
if stream: # not double-detach()'d
os.setsid()
self.parent.send_await(Message(handle=DETACHING))
LOG.info('Detaching from %r; parent is %s', stream, self.parent)
for x in range(20):
pending = self.broker.defer_sync(stream.protocol.pending_bytes)
if not pending:
break
time.sleep(0.05)
if pending:
LOG.error('Stream had %d bytes after 2000ms', pending)
self.broker.defer(stream.on_disconnect, self.broker)
def _setup_master(self):
Router.max_message_size = self.config['max_message_size']
if self.config['profiling']:
enable_profiling()
self.broker = Broker(activate_compat=False)
self.router = Router(self.broker)
self.router.debug = self.config.get('debug', False)
self.router.unidirectional = self.config['unidirectional']
self.router.add_handler(
fn=self._on_shutdown_msg,
handle=SHUTDOWN,
policy=has_parent_authority,
)
self.master = Context(self.router, 0, 'master')
parent_id = self.config['parent_ids'][0]
if parent_id == 0:
self.parent = self.master
else:
self.parent = Context(self.router, parent_id, 'parent')
in_fd = self.config.get('in_fd', 100)
in_fp = os.fdopen(os.dup(in_fd), 'rb', 0)
os.close(in_fd)
out_fp = os.fdopen(os.dup(self.config.get('out_fd', 1)), 'wb', 0)
self.stream = MitogenProtocol.build_stream(
self.router,
parent_id,
local_id=self.config['context_id'],
parent_ids=self.config['parent_ids']
)
self.stream.accept(in_fp, out_fp)
self.stream.name = 'parent'
self.stream.receive_side.keep_alive = False
listen(self.stream, 'disconnect', self._on_parent_disconnect)
listen(self.broker, 'exit', self._on_broker_exit)
def _reap_first_stage(self):
try:
os.wait() # Reap first stage.
except OSError:
pass # No first stage exists (e.g. fakessh)
def _setup_logging(self):
self.log_handler = LogHandler(self.master)
root = logging.getLogger()
root.setLevel(self.config['log_level'])
root.handlers = [self.log_handler]
if self.config['debug']:
enable_debug_logging()
def _setup_importer(self):
importer = self.config.get('importer')
if importer:
importer._install_handler(self.router)
importer._context = self.parent
else:
core_src_fd = self.config.get('core_src_fd', 101)
if core_src_fd:
fp = os.fdopen(core_src_fd, 'rb', 1)
try:
core_src = fp.read()
# Strip "ExternalContext.main()" call from last line.
core_src = b('\n').join(core_src.splitlines()[:-1])
finally:
fp.close()
else:
core_src = None
importer = Importer(
self.router,
self.parent,
core_src,
self.config.get('whitelist', ()),
self.config.get('blacklist', ()),
)
self.importer = importer
self.router.importer = importer
sys.meta_path.insert(0, self.importer)
def _setup_package(self):
global mitogen
mitogen = imp.new_module('mitogen')
mitogen.__package__ = 'mitogen'
mitogen.__path__ = []
mitogen.__loader__ = self.importer
mitogen.main = lambda *args, **kwargs: (lambda func: None)
mitogen.core = sys.modules['__main__']
mitogen.core.__file__ = 'x/mitogen/core.py' # For inspect.getsource()
mitogen.core.__loader__ = self.importer
sys.modules['mitogen'] = mitogen
sys.modules['mitogen.core'] = mitogen.core
del sys.modules['__main__']
def _setup_globals(self):
mitogen.is_master = False
mitogen.__version__ = self.config['version']
mitogen.context_id = self.config['context_id']
mitogen.parent_ids = self.config['parent_ids'][:]
mitogen.parent_id = mitogen.parent_ids[0]
def _nullify_stdio(self):
"""
Open /dev/null to replace stdio temporarily. In case of odd startup,
assume we may be allocated a standard handle.
"""
for stdfd, mode in ((0, os.O_RDONLY), (1, os.O_RDWR), (2, os.O_RDWR)):
fd = os.open('/dev/null', mode)
if fd != stdfd:
os.dup2(fd, stdfd)
os.close(fd)
def _preserve_tty_fp(self):
"""
#481: when stderr is a TTY due to being started via tty_create_child()
or hybrid_tty_create_child(), and some privilege escalation tool like
prehistoric versions of sudo exec this process over the top of itself,
there is nothing left to keep the slave PTY open after we replace our
stdio. Therefore if stderr is a TTY, keep around a permanent dup() to
avoid receiving SIGHUP.
"""
try:
if os.isatty(2):
self.reserve_tty_fp = os.fdopen(os.dup(2), 'r+b', 0)
set_cloexec(self.reserve_tty_fp.fileno())
except OSError:
pass
def _setup_stdio(self):
self._preserve_tty_fp()
# When sys.stdout was opened by the runtime, overwriting it will not
# close FD 1. However when forking from a child that previously used
# fdopen(), overwriting it /will/ close FD 1. So we must swallow the
# close before IoLogger overwrites FD 1, otherwise its new FD 1 will be
# clobbered. Additionally, stdout must be replaced with /dev/null prior
# to stdout.close(), since if block buffering was active in the parent,
# any pre-fork buffered data will be flushed on close(), corrupting the
# connection to the parent.
self._nullify_stdio()
sys.stdout.close()
self._nullify_stdio()
self.loggers = []
for name, fd in (('stdout', 1), ('stderr', 2)):
log = IoLoggerProtocol.build_stream(name, fd)
self.broker.start_receive(log)
self.loggers.append(log)
# Reopen with line buffering.
sys.stdout = os.fdopen(1, 'w', 1)
def main(self):
self._setup_master()
try:
try:
self._setup_logging()
self._setup_importer()
self._reap_first_stage()
if self.config.get('setup_package', True):
self._setup_package()
self._setup_globals()
if self.config.get('setup_stdio', True):
self._setup_stdio()
self.dispatcher = Dispatcher(self)
self.router.register(self.parent, self.stream)
self.router._setup_logging()
_v and LOG.debug('Python version is %s', sys.version)
_v and LOG.debug('Parent is context %r (%s); my ID is %r',
self.parent.context_id, self.parent.name,
mitogen.context_id)
_v and LOG.debug('pid:%r ppid:%r uid:%r/%r, gid:%r/%r host:%r',
os.getpid(), os.getppid(), os.geteuid(),
os.getuid(), os.getegid(), os.getgid(),
socket.gethostname())
sys.executable = os.environ.pop('ARGV0', sys.executable)
_v and LOG.debug('Recovered sys.executable: %r', sys.executable)
if self.config.get('send_ec2', True):
self.stream.transmit_side.write(b('MITO002\n'))
self.broker._py24_25_compat()
self.log_handler.uncork()
self.dispatcher.run()
_v and LOG.debug('ExternalContext.main() normal exit')
except KeyboardInterrupt:
LOG.debug('KeyboardInterrupt received, exiting gracefully.')
except BaseException:
LOG.exception('ExternalContext.main() crashed')
raise
finally:
self.broker.shutdown()
# 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
"""
Basic signal handler for dumping thread stacks.
"""
import difflib
import logging
import os
import gc
import signal
import sys
import threading
import time
import traceback
import mitogen.core
import mitogen.parent
LOG = logging.getLogger(__name__)
_last = None
def enable_evil_interrupts():
signal.signal(signal.SIGALRM, (lambda a, b: None))
signal.setitimer(signal.ITIMER_REAL, 0.01, 0.01)
def disable_evil_interrupts():
signal.setitimer(signal.ITIMER_REAL, 0, 0)
def _hex(n):
return '%08x' % n
def get_subclasses(klass):
"""
Rather than statically import every interesting subclass, forcing it all to
be transferred and potentially disrupting the debugged environment,
enumerate only those loaded in memory. Also returns the original class.
"""
stack = [klass]
seen = set()
while stack:
klass = stack.pop()
seen.add(klass)
stack.extend(klass.__subclasses__())
return seen
def get_routers():
return dict(
(_hex(id(router)), router)
for klass in get_subclasses(mitogen.core.Router)
for router in gc.get_referrers(klass)
if isinstance(router, mitogen.core.Router)
)
def get_router_info():
return {
'routers': dict(
(id_, {
'id': id_,
'streams': len(set(router._stream_by_id.values())),
'contexts': len(set(router._context_by_id.values())),
'handles': len(router._handle_map),
})
for id_, router in get_routers().items()
)
}
def get_stream_info(router_id):
router = get_routers().get(router_id)
return {
'streams': dict(
(_hex(id(stream)), ({
'name': stream.name,
'remote_id': stream.remote_id,
'sent_module_count': len(getattr(stream, 'sent_modules', [])),
'routes': sorted(getattr(stream, 'routes', [])),
'type': type(stream).__module__,
}))
for via_id, stream in router._stream_by_id.items()
)
}
def format_stacks():
name_by_id = dict(
(t.ident, t.name)
for t in threading.enumerate()
)
l = ['', '']
for threadId, stack in sys._current_frames().items():
l += ["# PID %d ThreadID: (%s) %s; %r" % (
os.getpid(),
name_by_id.get(threadId, '<no name>'),
threadId,
stack,
)]
#stack = stack.f_back.f_back
for filename, lineno, name, line in traceback.extract_stack(stack):
l += [
'File: "%s", line %d, in %s' % (
filename,
lineno,
name
)
]
if line:
l += [' ' + line.strip()]
l += ['']
l += ['', '']
return '\n'.join(l)
def get_snapshot():
global _last
s = format_stacks()
snap = s
if _last:
snap += '\n'
diff = list(difflib.unified_diff(
a=_last.splitlines(),
b=s.splitlines(),
fromfile='then',
tofile='now'
))
if diff:
snap += '\n'.join(diff) + '\n'
else:
snap += '(no change since last time)\n'
_last = s
return snap
def _handler(*_):
fp = open('/dev/tty', 'w', 1)
fp.write(get_snapshot())
fp.close()
def install_handler():
signal.signal(signal.SIGUSR2, _handler)
def _logging_main(secs):
while True:
time.sleep(secs)
LOG.info('PERIODIC THREAD DUMP\n\n%s', get_snapshot())
def dump_to_logger(secs=5):
th = threading.Thread(
target=_logging_main,
kwargs={'secs': secs},
name='mitogen.debug.dump_to_logger',
)
th.setDaemon(True)
th.start()
class ContextDebugger(object):
@classmethod
@mitogen.core.takes_econtext
def _configure_context(cls, econtext):
mitogen.parent.upgrade_router(econtext)
econtext.debugger = cls(econtext.router)
def __init__(self, router):
self.router = router
self.router.add_handler(
func=self._on_debug_msg,
handle=mitogen.core.DEBUG,
persist=True,
policy=mitogen.core.has_parent_authority,
)
mitogen.core.listen(router, 'register', self._on_stream_register)
LOG.debug('Context debugging configured.')
def _on_stream_register(self, context, stream):
LOG.debug('_on_stream_register: sending configure() to %r', stream)
context.call_async(ContextDebugger._configure_context)
def _on_debug_msg(self, msg):
if msg != mitogen.core._DEAD:
threading.Thread(
target=self._handle_debug_msg,
name='ContextDebuggerHandler',
args=(msg,)
).start()
def _handle_debug_msg(self, msg):
try:
method, args, kwargs = msg.unpickle()
msg.reply(getattr(self, method)(*args, **kwargs))
except Exception:
e = sys.exc_info()[1]
msg.reply(mitogen.core.CallError(e))
# 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 logging
import re
import mitogen.core
import mitogen.parent
LOG = logging.getLogger(__name__)
password_incorrect_msg = 'doas password is incorrect'
password_required_msg = 'doas password is required'
class PasswordError(mitogen.core.StreamError):
pass
class Options(mitogen.parent.Options):
username = u'root'
password = None
doas_path = 'doas'
password_prompt = u'Password:'
incorrect_prompts = (
u'doas: authentication failed', # slicer69/doas
u'doas: Authorization failed', # openbsd/src
)
def __init__(self, username=None, password=None, doas_path=None,
password_prompt=None, incorrect_prompts=None, **kwargs):
super(Options, self).__init__(**kwargs)
if username is not None:
self.username = mitogen.core.to_text(username)
if password is not None:
self.password = mitogen.core.to_text(password)
if doas_path is not None:
self.doas_path = doas_path
if password_prompt is not None:
self.password_prompt = mitogen.core.to_text(password_prompt)
if incorrect_prompts is not None:
self.incorrect_prompts = [
mitogen.core.to_text(p)
for p in incorrect_prompts
]
class BootstrapProtocol(mitogen.parent.RegexProtocol):
password_sent = False
def setup_patterns(self, conn):
prompt_pattern = re.compile(
re.escape(conn.options.password_prompt).encode('utf-8'),
re.I
)
incorrect_prompt_pattern = re.compile(
u'|'.join(
re.escape(s)
for s in conn.options.incorrect_prompts
).encode('utf-8'),
re.I
)
self.PATTERNS = [
(incorrect_prompt_pattern, type(self)._on_incorrect_password),
]
self.PARTIAL_PATTERNS = [
(prompt_pattern, type(self)._on_password_prompt),
]
def _on_incorrect_password(self, line, match):
if self.password_sent:
self.stream.conn._fail_connection(
PasswordError(password_incorrect_msg)
)
def _on_password_prompt(self, line, match):
if self.stream.conn.options.password is None:
self.stream.conn._fail_connection(
PasswordError(password_required_msg)
)
return
if self.password_sent:
self.stream.conn._fail_connection(
PasswordError(password_incorrect_msg)
)
return
LOG.debug('sending password')
self.stream.transmit_side.write(
(self.stream.conn.options.password + '\n').encode('utf-8')
)
self.password_sent = True
class Connection(mitogen.parent.Connection):
options_class = Options
diag_protocol_class = BootstrapProtocol
create_child = staticmethod(mitogen.parent.hybrid_tty_create_child)
child_is_immediate_subprocess = False
def _get_name(self):
return u'doas.' + self.options.username
def stderr_stream_factory(self):
stream = super(Connection, self).stderr_stream_factory()
stream.protocol.setup_patterns(self)
return stream
def get_boot_command(self):
bits = [self.options.doas_path, '-u', self.options.username, '--']
return bits + super(Connection, self).get_boot_command()
# 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 logging
import mitogen.core
import mitogen.parent
LOG = logging.getLogger(__name__)
class Options(mitogen.parent.Options):
container = None
image = None
username = None
docker_path = u'docker'
def __init__(self, container=None, image=None, docker_path=None,
username=None, **kwargs):
super(Options, self).__init__(**kwargs)
assert container or image
if container:
self.container = mitogen.core.to_text(container)
if image:
self.image = mitogen.core.to_text(image)
if docker_path:
self.docker_path = mitogen.core.to_text(docker_path)
if username:
self.username = mitogen.core.to_text(username)
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = False
# TODO: better way of capturing errors such as "No such container."
create_child_args = {
'merge_stdio': True
}
def _get_name(self):
return u'docker.' + (self.options.container or self.options.image)
def get_boot_command(self):
args = ['--interactive']
if self.options.username:
args += ['--user=' + self.options.username]
bits = [self.options.docker_path]
if self.options.container:
bits += ['exec'] + args + [self.options.container]
elif self.options.image:
bits += ['run'] + args + ['--rm', self.options.image]
return bits + super(Connection, self).get_boot_command()
# 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
"""
:mod:`mitogen.fakessh` is a stream implementation that starts a subprocess with
its environment modified such that ``PATH`` searches for `ssh` return a Mitogen
implementation of SSH. When invoked, this implementation arranges for the
command line supplied by the caller to be executed in a remote context, reusing
the parent context's (possibly proxied) connection to that remote context.
This allows tools like `rsync` and `scp` to transparently reuse the connections
and tunnels already established by the host program to connect to a target
machine, without wasteful redundant SSH connection setup, 3-way handshakes, or
firewall hopping configurations, and enables these tools to be used in
impossible scenarios, such as over `sudo` with ``requiretty`` enabled.
The fake `ssh` command source is written to a temporary file on disk, and
consists of a copy of the :py:mod:`mitogen.core` source code (just like any
other child context), with a line appended to cause it to connect back to the
host process over an FD it inherits. As there is no reliance on an existing
filesystem file, it is possible for child contexts to use fakessh.
As a consequence of connecting back through an inherited FD, only one SSH
invocation is possible, which is fine for tools like `rsync`, however in future
this restriction will be lifted.
Sequence:
1. ``fakessh`` Context and Stream created by parent context. The stream's
buffer has a :py:func:`_fakessh_main` :py:data:`CALL_FUNCTION
<mitogen.core.CALL_FUNCTION>` enqueued.
2. Target program (`rsync/scp/sftp`) invoked, which internally executes
`ssh` from ``PATH``.
3. :py:mod:`mitogen.core` bootstrap begins, recovers the stream FD
inherited via the target program, established itself as the fakessh
context.
4. :py:func:`_fakessh_main` :py:data:`CALL_FUNCTION
<mitogen.core.CALL_FUNCTION>` is read by fakessh context,
a. sets up :py:class:`IoPump` for stdio, registers
stdin_handle for local context.
b. Enqueues :py:data:`CALL_FUNCTION <mitogen.core.CALL_FUNCTION>` for
:py:func:`_start_slave` invoked in target context,
i. the program from the `ssh` command line is started
ii. sets up :py:class:`IoPump` for `ssh` command line process's
stdio pipes
iii. returns `(control_handle, stdin_handle)` to
:py:func:`_fakessh_main`
5. :py:func:`_fakessh_main` receives control/stdin handles from from
:py:func:`_start_slave`,
a. registers remote's stdin_handle with local :py:class:`IoPump`.
b. sends `("start", local_stdin_handle)` to remote's control_handle
c. registers local :py:class:`IoPump` with
:py:class:`mitogen.core.Broker`.
d. loops waiting for `local stdout closed && remote stdout closed`
6. :py:func:`_start_slave` control channel receives `("start", stdin_handle)`,
a. registers remote's stdin_handle with local :py:class:`IoPump`
b. registers local :py:class:`IoPump` with
:py:class:`mitogen.core.Broker`.
c. loops waiting for `local stdout closed && remote stdout closed`
"""
import getopt
import inspect
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import threading
import mitogen.core
import mitogen.master
import mitogen.parent
from mitogen.core import LOG, IOLOG
SSH_GETOPTS = (
"1246ab:c:e:fgi:kl:m:no:p:qstvx"
"ACD:E:F:I:KL:MNO:PQ:R:S:TVw:W:XYy"
)
_mitogen = None
class IoPump(mitogen.core.Protocol):
_output_buf = ''
_closed = False
def __init__(self, broker):
self._broker = broker
def write(self, s):
self._output_buf += s
self._broker._start_transmit(self)
def close(self):
self._closed = True
# If local process hasn't exitted yet, ensure its write buffer is
# drained before lazily triggering disconnect in on_transmit.
if self.transmit_side.fp.fileno() is not None:
self._broker._start_transmit(self)
def on_shutdown(self, stream, broker):
self.close()
def on_transmit(self, stream, broker):
written = self.transmit_side.write(self._output_buf)
IOLOG.debug('%r.on_transmit() -> len %r', self, written)
if written is None:
self.on_disconnect(broker)
else:
self._output_buf = self._output_buf[written:]
if not self._output_buf:
broker._stop_transmit(self)
if self._closed:
self.on_disconnect(broker)
def on_receive(self, stream, broker):
s = stream.receive_side.read()
IOLOG.debug('%r.on_receive() -> len %r', self, len(s))
if s:
mitogen.core.fire(self, 'receive', s)
else:
self.on_disconnect(broker)
def __repr__(self):
return 'IoPump(%r, %r)' % (
self.receive_side.fp.fileno(),
self.transmit_side.fp.fileno(),
)
class Process(object):
"""
Manages the lifetime and pipe connections of the SSH command running in the
slave.
"""
def __init__(self, router, stdin, stdout, proc=None):
self.router = router
self.stdin = stdin
self.stdout = stdout
self.proc = proc
self.control_handle = router.add_handler(self._on_control)
self.stdin_handle = router.add_handler(self._on_stdin)
self.pump = IoPump.build_stream(router.broker)
self.pump.accept(stdin, stdout)
self.stdin = None
self.control = None
self.wake_event = threading.Event()
mitogen.core.listen(self.pump, 'disconnect', self._on_pump_disconnect)
mitogen.core.listen(self.pump, 'receive', self._on_pump_receive)
if proc:
pmon = mitogen.parent.ProcessMonitor.instance()
pmon.add(proc.pid, self._on_proc_exit)
def __repr__(self):
return 'Process(%r, %r)' % (self.stdin, self.stdout)
def _on_proc_exit(self, status):
LOG.debug('%r._on_proc_exit(%r)', self, status)
self.control.put(('exit', status))
def _on_stdin(self, msg):
if msg.is_dead:
IOLOG.debug('%r._on_stdin() -> %r', self, data)
self.pump.protocol.close()
return
data = msg.unpickle()
IOLOG.debug('%r._on_stdin() -> len %d', self, len(data))
self.pump.protocol.write(data)
def _on_control(self, msg):
if not msg.is_dead:
command, arg = msg.unpickle(throw=False)
LOG.debug('%r._on_control(%r, %s)', self, command, arg)
func = getattr(self, '_on_%s' % (command,), None)
if func:
return func(msg, arg)
LOG.warning('%r: unknown command %r', self, command)
def _on_start(self, msg, arg):
dest = mitogen.core.Context(self.router, msg.src_id)
self.control = mitogen.core.Sender(dest, arg[0])
self.stdin = mitogen.core.Sender(dest, arg[1])
self.router.broker.start_receive(self.pump)
def _on_exit(self, msg, arg):
LOG.debug('on_exit: proc = %r', self.proc)
if self.proc:
self.proc.terminate()
else:
self.router.broker.shutdown()
def _on_pump_receive(self, s):
IOLOG.info('%r._on_pump_receive(len %d)', self, len(s))
self.stdin.put(s)
def _on_pump_disconnect(self):
LOG.debug('%r._on_pump_disconnect()', self)
mitogen.core.fire(self, 'disconnect')
self.stdin.close()
self.wake_event.set()
def start_master(self, stdin, control):
self.stdin = stdin
self.control = control
control.put(('start', (self.control_handle, self.stdin_handle)))
self.router.broker.start_receive(self.pump)
def wait(self):
while not self.wake_event.isSet():
# Timeout is used so that sleep is interruptible, as blocking
# variants of libc thread operations cannot be interrupted e.g. via
# KeyboardInterrupt. isSet() test and wait() are separate since in
# <2.7 wait() always returns None.
self.wake_event.wait(0.1)
@mitogen.core.takes_router
def _start_slave(src_id, cmdline, router):
"""
This runs in the target context, it is invoked by _fakessh_main running in
the fakessh context immediately after startup. It starts the slave process
(the the point where it has a stdin_handle to target but not stdout_chan to
write to), and waits for main to.
"""
LOG.debug('_start_slave(%r, %r)', router, cmdline)
proc = subprocess.Popen(
cmdline,
# SSH server always uses user's shell.
shell=True,
# SSH server always executes new commands in the user's HOME.
cwd=os.path.expanduser('~'),
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
)
process = Process(router, proc.stdin, proc.stdout, proc)
return process.control_handle, process.stdin_handle
#
# SSH client interface.
#
def exit():
_mitogen.broker.shutdown()
def die(msg, *args):
if args:
msg %= args
sys.stderr.write('%s\n' % (msg,))
exit()
def parse_args():
hostname = None
remain = sys.argv[1:]
allopts = []
restarted = 0
while remain and restarted < 2:
opts, args = getopt.getopt(remain, SSH_GETOPTS)
remain = remain[:] # getopt bug!
allopts += opts
if not args:
break
if not hostname:
hostname = args.pop(0)
remain = remain[remain.index(hostname) + 1:]
restarted += 1
return hostname, allopts, args
@mitogen.core.takes_econtext
def _fakessh_main(dest_context_id, econtext):
hostname, opts, args = parse_args()
if not hostname:
die('Missing hostname')
subsystem = False
for opt, optarg in opts:
if opt == '-s':
subsystem = True
else:
LOG.debug('Warning option %s %s is ignored.', opt, optarg)
LOG.debug('hostname: %r', hostname)
LOG.debug('opts: %r', opts)
LOG.debug('args: %r', args)
if subsystem:
die('-s <subsystem> is not yet supported')
if not args:
die('fakessh: login mode not supported and no command specified')
dest = mitogen.parent.Context(econtext.router, dest_context_id)
# Even though SSH receives an argument vector, it still cats the vector
# together before sending to the server, the server just uses /bin/sh -c to
# run the command. We must remain puke-for-puke compatible.
control_handle, stdin_handle = dest.call(_start_slave,
mitogen.context_id, ' '.join(args))
LOG.debug('_fakessh_main: received control_handle=%r, stdin_handle=%r',
control_handle, stdin_handle)
process = Process(econtext.router,
stdin=os.fdopen(1, 'w+b', 0),
stdout=os.fdopen(0, 'r+b', 0))
process.start_master(
stdin=mitogen.core.Sender(dest, stdin_handle),
control=mitogen.core.Sender(dest, control_handle),
)
process.wait()
process.control.put(('exit', None))
def _get_econtext_config(context, sock2):
parent_ids = mitogen.parent_ids[:]
parent_ids.insert(0, mitogen.context_id)
return {
'context_id': context.context_id,
'core_src_fd': None,
'debug': getattr(context.router, 'debug', False),
'in_fd': sock2.fileno(),
'log_level': mitogen.parent.get_log_level(),
'max_message_size': context.router.max_message_size,
'out_fd': sock2.fileno(),
'parent_ids': parent_ids,
'profiling': getattr(context.router, 'profiling', False),
'unidirectional': getattr(context.router, 'unidirectional', False),
'setup_stdio': False,
'version': mitogen.__version__,
}
#
# Public API.
#
@mitogen.core.takes_econtext
@mitogen.core.takes_router
def run(dest, router, args, deadline=None, econtext=None):
"""
Run the command specified by `args` such that ``PATH`` searches for SSH by
the command will cause its attempt to use SSH to execute a remote program
to be redirected to use mitogen to execute that program using the context
`dest` instead.
:param list args:
Argument vector.
:param mitogen.core.Context dest:
The destination context to execute the SSH command line in.
:param mitogen.core.Router router:
:param list[str] args:
Command line arguments for local program, e.g.
``['rsync', '/tmp', 'remote:/tmp']``
:returns:
Exit status of the child process.
"""
if econtext is not None:
mitogen.parent.upgrade_router(econtext)
context_id = router.allocate_id()
fakessh = mitogen.parent.Context(router, context_id)
fakessh.name = u'fakessh.%d' % (context_id,)
sock1, sock2 = socket.socketpair()
stream = mitogen.core.Stream(router, context_id)
stream.name = u'fakessh'
stream.accept(sock1, sock1)
router.register(fakessh, stream)
# Held in socket buffer until process is booted.
fakessh.call_async(_fakessh_main, dest.context_id)
tmp_path = tempfile.mkdtemp(prefix='mitogen_fakessh')
try:
ssh_path = os.path.join(tmp_path, 'ssh')
fp = open(ssh_path, 'w')
try:
fp.write('#!%s\n' % (mitogen.parent.get_sys_executable(),))
fp.write(inspect.getsource(mitogen.core))
fp.write('\n')
fp.write('ExternalContext(%r).main()\n' % (
_get_econtext_config(context, sock2),
))
finally:
fp.close()
os.chmod(ssh_path, int('0755', 8))
env = os.environ.copy()
env.update({
'PATH': '%s:%s' % (tmp_path, env.get('PATH', '')),
'ARGV0': mitogen.parent.get_sys_executable(),
'SSH_PATH': ssh_path,
})
proc = subprocess.Popen(args, env=env)
return proc.wait()
finally:
shutil.rmtree(tmp_path)
# 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 errno
import logging
import os
import random
import sys
import threading
import traceback
import mitogen.core
import mitogen.parent
from mitogen.core import b
LOG = logging.getLogger(__name__)
# Python 2.4/2.5 cannot support fork+threads whatsoever, it doesn't even fix up
# interpreter state. So 2.4/2.5 interpreters start .local() contexts for
# isolation instead. Since we don't have any crazy memory sharing problems to
# avoid, there is no virginal fork parent either. The child is started directly
# from the login/become process. In future this will be default everywhere,
# fork is brainwrong from the stone age.
FORK_SUPPORTED = sys.version_info >= (2, 6)
class Error(mitogen.core.StreamError):
pass
def fixup_prngs():
"""
Add 256 bits of /dev/urandom to OpenSSL's PRNG in the child, and re-seed
the random package with the same data.
"""
s = os.urandom(256 // 8)
random.seed(s)
if 'ssl' in sys.modules:
sys.modules['ssl'].RAND_add(s, 75.0)
def reset_logging_framework():
"""
After fork, ensure any logging.Handler locks are recreated, as a variety of
threads in the parent may have been using the logging package at the moment
of fork.
It is not possible to solve this problem in general; see :gh:issue:`150`
for a full discussion.
"""
logging._lock = threading.RLock()
# The root logger does not appear in the loggerDict.
logging.Logger.manager.loggerDict = {}
logging.getLogger().handlers = []
def on_fork():
"""
Should be called by any program integrating Mitogen each time the process
is forked, in the context of the new child.
"""
reset_logging_framework() # Must be first!
fixup_prngs()
mitogen.core.Latch._on_fork()
mitogen.core.Side._on_fork()
mitogen.core.ExternalContext.service_stub_lock = threading.Lock()
mitogen__service = sys.modules.get('mitogen.service')
if mitogen__service:
mitogen__service._pool_lock = threading.Lock()
def handle_child_crash():
"""
Respond to _child_main() crashing by ensuring the relevant exception is
logged to /dev/tty.
"""
tty = open('/dev/tty', 'wb')
tty.write('\n\nFORKED CHILD PID %d CRASHED\n%s\n\n' % (
os.getpid(),
traceback.format_exc(),
))
tty.close()
os._exit(1)
def _convert_exit_status(status):
"""
Convert a :func:`os.waitpid`-style exit status to a :mod:`subprocess` style
exit status.
"""
if os.WIFEXITED(status):
return os.WEXITSTATUS(status)
elif os.WIFSIGNALED(status):
return -os.WTERMSIG(status)
elif os.WIFSTOPPED(status):
return -os.WSTOPSIG(status)
class Process(mitogen.parent.Process):
def poll(self):
try:
pid, status = os.waitpid(self.pid, os.WNOHANG)
except OSError:
e = sys.exc_info()[1]
if e.args[0] == errno.ECHILD:
LOG.warn('%r: waitpid(%r) produced ECHILD', self, self.pid)
return
raise
if not pid:
return
return _convert_exit_status(status)
class Options(mitogen.parent.Options):
#: Reference to the importer, if any, recovered from the parent.
importer = None
#: User-supplied function for cleaning up child process state.
on_fork = None
def __init__(self, old_router, max_message_size, on_fork=None, debug=False,
profiling=False, unidirectional=False, on_start=None,
name=None):
if not FORK_SUPPORTED:
raise Error(self.python_version_msg)
# fork method only supports a tiny subset of options.
super(Options, self).__init__(
max_message_size=max_message_size, debug=debug,
profiling=profiling, unidirectional=unidirectional, name=name,
)
self.on_fork = on_fork
self.on_start = on_start
responder = getattr(old_router, 'responder', None)
if isinstance(responder, mitogen.parent.ModuleForwarder):
self.importer = responder.importer
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = True
python_version_msg = (
"The mitogen.fork method is not supported on Python versions "
"prior to 2.6, since those versions made no attempt to repair "
"critical interpreter state following a fork. Please use the "
"local() method instead."
)
name_prefix = u'fork'
def start_child(self):
parentfp, childfp = mitogen.parent.create_socketpair()
pid = os.fork()
if pid:
childfp.close()
return Process(pid, stdin=parentfp, stdout=parentfp)
else:
parentfp.close()
self._wrap_child_main(childfp)
def _wrap_child_main(self, childfp):
try:
self._child_main(childfp)
except BaseException:
handle_child_crash()
def get_econtext_config(self):
config = super(Connection, self).get_econtext_config()
config['core_src_fd'] = None
config['importer'] = self.options.importer
config['send_ec2'] = False
config['setup_package'] = False
if self.options.on_start:
config['on_start'] = self.options.on_start
return config
def _child_main(self, childfp):
on_fork()
if self.options.on_fork:
self.options.on_fork()
mitogen.core.set_block(childfp.fileno())
childfp.send(b('MITO002\n'))
# Expected by the ExternalContext.main().
os.dup2(childfp.fileno(), 1)
os.dup2(childfp.fileno(), 100)
# Overwritten by ExternalContext.main(); we must replace the
# parent-inherited descriptors that were closed by Side._on_fork() to
# avoid ExternalContext.main() accidentally allocating new files over
# the standard handles.
os.dup2(childfp.fileno(), 0)
# Avoid corrupting the stream on fork crash by dupping /dev/null over
# stderr. Instead, handle_child_crash() uses /dev/tty to log errors.
devnull = os.open('/dev/null', os.O_WRONLY)
if devnull != 2:
os.dup2(devnull, 2)
os.close(devnull)
# If we're unlucky, childfp.fileno() may coincidentally be one of our
# desired FDs. In that case closing it breaks ExternalContext.main().
if childfp.fileno() not in (0, 1, 100):
childfp.close()
mitogen.core.IOLOG.setLevel(logging.INFO)
try:
try:
mitogen.core.ExternalContext(self.get_econtext_config()).main()
except Exception:
# TODO: report exception somehow.
os._exit(72)
finally:
# Don't trigger atexit handlers, they were copied from the parent.
os._exit(0)
# 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 mitogen.core
import mitogen.parent
class Options(mitogen.parent.Options):
container = None
username = None
jexec_path = u'/usr/sbin/jexec'
def __init__(self, container, jexec_path=None, username=None, **kwargs):
super(Options, self).__init__(**kwargs)
self.container = mitogen.core.to_text(container)
if username:
self.username = mitogen.core.to_text(username)
if jexec_path:
self.jexec_path = jexec_path
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = False
create_child_args = {
'merge_stdio': True
}
def _get_name(self):
return u'jail.' + self.options.container
def get_boot_command(self):
bits = [self.options.jexec_path]
if self.options.username:
bits += ['-U', self.options.username]
bits += [self.options.container]
return bits + super(Connection, self).get_boot_command()
# Copyright 2018, Yannig Perre
#
# 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 mitogen.core
import mitogen.parent
class Options(mitogen.parent.Options):
pod = None
kubectl_path = 'kubectl'
kubectl_args = None
def __init__(self, pod, kubectl_path=None, kubectl_args=None, **kwargs):
super(Options, self).__init__(**kwargs)
assert pod
self.pod = pod
if kubectl_path:
self.kubectl_path = kubectl_path
self.kubectl_args = kubectl_args or []
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = True
# TODO: better way of capturing errors such as "No such container."
create_child_args = {
'merge_stdio': True
}
def _get_name(self):
return u'kubectl.%s%s' % (self.options.pod, self.options.kubectl_args)
def get_boot_command(self):
bits = [
self.options.kubectl_path
] + self.options.kubectl_args + [
'exec', '-it', self.options.pod
]
return bits + ["--"] + super(Connection, self).get_boot_command()
# 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 mitogen.core
import mitogen.parent
class Options(mitogen.parent.Options):
container = None
lxc_attach_path = 'lxc-attach'
def __init__(self, container, lxc_attach_path=None, **kwargs):
super(Options, self).__init__(**kwargs)
self.container = container
if lxc_attach_path:
self.lxc_attach_path = lxc_attach_path
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = False
create_child_args = {
# If lxc-attach finds any of stdin, stdout, stderr connected to a TTY,
# to prevent input injection it creates a proxy pty, forcing all IO to
# be buffered in <4KiB chunks. So ensure stderr is also routed to the
# socketpair.
'merge_stdio': True
}
eof_error_hint = (
'Note: many versions of LXC do not report program execution failure '
'meaningfully. Please check the host logs (/var/log) for more '
'information.'
)
def _get_name(self):
return u'lxc.' + self.options.container
def get_boot_command(self):
bits = [
self.options.lxc_attach_path,
'--clear-env',
'--name', self.options.container,
'--',
]
return bits + super(Connection, self).get_boot_command()
# 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 mitogen.core
import mitogen.parent
class Options(mitogen.parent.Options):
container = None
lxc_path = 'lxc'
python_path = 'python'
def __init__(self, container, lxc_path=None, **kwargs):
super(Options, self).__init__(**kwargs)
self.container = container
if lxc_path:
self.lxc_path = lxc_path
class Connection(mitogen.parent.Connection):
options_class = Options
child_is_immediate_subprocess = False
create_child_args = {
# If lxc finds any of stdin, stdout, stderr connected to a TTY, to
# prevent input injection it creates a proxy pty, forcing all IO to be
# buffered in <4KiB chunks. So ensure stderr is also routed to the
# socketpair.
'merge_stdio': True
}
eof_error_hint = (
'Note: many versions of LXC do not report program execution failure '
'meaningfully. Please check the host logs (/var/log) for more '
'information.'
)
def _get_name(self):
return u'lxd.' + self.options.container
def get_boot_command(self):
bits = [
self.options.lxc_path,
'exec',
'--mode=noninteractive',
self.options.container,
'--',
]
return bits + super(Connection, self).get_boot_command()
# 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
"""
This module implements functionality required by master processes, such as
starting new contexts via SSH. Its size is also restricted, since it must
be sent to any context that will be used to establish additional child
contexts.
"""
import dis
import errno
import imp
import inspect
import itertools
import logging
import os
import pkgutil
import re
import string
import sys
import threading
import types
import zlib
try:
import sysconfig
except ImportError:
sysconfig = None
if not hasattr(pkgutil, 'find_loader'):
# find_loader() was new in >=2.5, but the modern pkgutil.py syntax has
# been kept intentionally 2.3 compatible so we can reuse it.
from mitogen.compat import pkgutil
import mitogen
import mitogen.core
import mitogen.minify
import mitogen.parent
from mitogen.core import b
from mitogen.core import IOLOG
from mitogen.core import LOG
from mitogen.core import str_partition
from mitogen.core import str_rpartition
from mitogen.core import to_text
imap = getattr(itertools, 'imap', map)
izip = getattr(itertools, 'izip', zip)
try:
any
except NameError:
from mitogen.core import any
try:
next
except NameError:
from mitogen.core import next
RLOG = logging.getLogger('mitogen.ctx')
def _stdlib_paths():
"""
Return a set of paths from which Python imports the standard library.
"""
attr_candidates = [
'prefix',
'real_prefix', # virtualenv: only set inside a virtual environment.
'base_prefix', # venv: always set, equal to prefix if outside.
]
prefixes = (getattr(sys, a, None) for a in attr_candidates)
version = 'python%s.%s' % sys.version_info[0:2]
s = set(os.path.abspath(os.path.join(p, 'lib', version))
for p in prefixes if p is not None)
# When running 'unit2 tests/module_finder_test.py' in a Py2 venv on Ubuntu
# 18.10, above is insufficient to catch the real directory.
if sysconfig is not None:
s.add(sysconfig.get_config_var('DESTLIB'))
return s
def is_stdlib_name(modname):
"""
Return :data:`True` if `modname` appears to come from the standard library.
"""
if imp.is_builtin(modname) != 0:
return True
module = sys.modules.get(modname)
if module is None:
return False
# six installs crap with no __file__
modpath = os.path.abspath(getattr(module, '__file__', ''))
return is_stdlib_path(modpath)
_STDLIB_PATHS = _stdlib_paths()
def is_stdlib_path(path):
return any(
os.path.commonprefix((libpath, path)) == libpath
and 'site-packages' not in path
and 'dist-packages' not in path
for libpath in _STDLIB_PATHS
)
def get_child_modules(path):
"""
Return the suffixes of submodules directly neated beneath of the package
directory at `path`.
:param str path:
Path to the module's source code on disk, or some PEP-302-recognized
equivalent. Usually this is the module's ``__file__`` attribute, but
is specified explicitly to avoid loading the module.
:return:
List of submodule name suffixes.
"""
it = pkgutil.iter_modules([os.path.dirname(path)])
return [to_text(name) for _, name, _ in it]
def _looks_like_script(path):
"""
Return :data:`True` if the (possibly extensionless) file at `path`
resembles a Python script. For now we simply verify the file contains
ASCII text.
"""
try:
fp = open(path, 'rb')
except IOError:
e = sys.exc_info()[1]
if e.args[0] == errno.EISDIR:
return False
raise
try:
sample = fp.read(512).decode('latin-1')
return not set(sample).difference(string.printable)
finally:
fp.close()
def _py_filename(path):
if not path:
return None
if path[-4:] in ('.pyc', '.pyo'):
path = path.rstrip('co')
if path.endswith('.py'):
return path
if os.path.exists(path) and _looks_like_script(path):
return path
def _get_core_source():
"""
Master version of parent.get_core_source().
"""
source = inspect.getsource(mitogen.core)
return mitogen.minify.minimize_source(source)
if mitogen.is_master:
# TODO: find a less surprising way of installing this.
mitogen.parent._get_core_source = _get_core_source
LOAD_CONST = dis.opname.index('LOAD_CONST')
IMPORT_NAME = dis.opname.index('IMPORT_NAME')
def _getarg(nextb, c):
if c >= dis.HAVE_ARGUMENT:
return nextb() | (nextb() << 8)
if sys.version_info < (3, 0):
def iter_opcodes(co):
# Yield `(op, oparg)` tuples from the code object `co`.
ordit = imap(ord, co.co_code)
nextb = ordit.next
return ((c, _getarg(nextb, c)) for c in ordit)
elif sys.version_info < (3, 6):
def iter_opcodes(co):
# Yield `(op, oparg)` tuples from the code object `co`.
ordit = iter(co.co_code)
nextb = ordit.__next__
return ((c, _getarg(nextb, c)) for c in ordit)
else:
def iter_opcodes(co):
# Yield `(op, oparg)` tuples from the code object `co`.
ordit = iter(co.co_code)
nextb = ordit.__next__
# https://github.com/abarnert/cpython/blob/c095a32f/Python/wordcode.md
return ((c, nextb()) for c in ordit)
def scan_code_imports(co):
"""
Given a code object `co`, scan its bytecode yielding any ``IMPORT_NAME``
and associated prior ``LOAD_CONST`` instructions representing an `Import`
statement or `ImportFrom` statement.
:return:
Generator producing `(level, modname, namelist)` tuples, where:
* `level`: -1 for normal import, 0, for absolute import, and >0 for
relative import.
* `modname`: Name of module to import, or from where `namelist` names
are imported.
* `namelist`: for `ImportFrom`, the list of names to be imported from
`modname`.
"""
opit = iter_opcodes(co)
opit, opit2, opit3 = itertools.tee(opit, 3)
try:
next(opit2)
next(opit3)
next(opit3)
except StopIteration:
return
if sys.version_info >= (2, 5):
for oparg1, oparg2, (op3, arg3) in izip(opit, opit2, opit3):
if op3 == IMPORT_NAME:
op2, arg2 = oparg2
op1, arg1 = oparg1
if op1 == op2 == LOAD_CONST:
yield (co.co_consts[arg1],
co.co_names[arg3],
co.co_consts[arg2] or ())
else:
# Python 2.4 did not yet have 'level', so stack format differs.
for oparg1, (op2, arg2) in izip(opit, opit2):
if op2 == IMPORT_NAME:
op1, arg1 = oparg1
if op1 == LOAD_CONST:
yield (-1, co.co_names[arg2], co.co_consts[arg1] or ())
class ThreadWatcher(object):
"""
Manage threads that wait for another thread to shut down, before invoking
`on_join()` for each associated ThreadWatcher.
In CPython it seems possible to use this method to ensure a non-main thread
is signalled when the main thread has exited, using a third thread as a
proxy.
"""
#: Protects remaining _cls_* members.
_cls_lock = threading.Lock()
#: PID of the process that last modified the class data. If the PID
#: changes, it means the thread watch dict refers to threads that no longer
#: exist in the current process (since it forked), and so must be reset.
_cls_pid = None
#: Map watched Thread -> list of ThreadWatcher instances.
_cls_instances_by_target = {}
#: Map watched Thread -> watcher Thread for each watched thread.
_cls_thread_by_target = {}
@classmethod
def _reset(cls):
"""
If we have forked since the watch dictionaries were initialized, all
that has is garbage, so clear it.
"""
if os.getpid() != cls._cls_pid:
cls._cls_pid = os.getpid()
cls._cls_instances_by_target.clear()
cls._cls_thread_by_target.clear()
def __init__(self, target, on_join):
self.target = target
self.on_join = on_join
@classmethod
def _watch(cls, target):
target.join()
for watcher in cls._cls_instances_by_target[target]:
watcher.on_join()
def install(self):
self._cls_lock.acquire()
try:
self._reset()
lst = self._cls_instances_by_target.setdefault(self.target, [])
lst.append(self)
if self.target not in self._cls_thread_by_target:
self._cls_thread_by_target[self.target] = threading.Thread(
name='mitogen.master.join_thread_async',
target=self._watch,
args=(self.target,)
)
self._cls_thread_by_target[self.target].start()
finally:
self._cls_lock.release()
def remove(self):
self._cls_lock.acquire()
try:
self._reset()
lst = self._cls_instances_by_target.get(self.target, [])
if self in lst:
lst.remove(self)
finally:
self._cls_lock.release()
@classmethod
def watch(cls, target, on_join):
watcher = cls(target, on_join)
watcher.install()
return watcher
class LogForwarder(object):
"""
Install a :data:`mitogen.core.FORWARD_LOG` handler that delivers forwarded
log events into the local logging framework. This is used by the master's
:class:`Router`.
The forwarded :class:`logging.LogRecord` objects are delivered to loggers
under ``mitogen.ctx.*`` corresponding to their
:attr:`mitogen.core.Context.name`, with the message prefixed with the
logger name used in the child. The records include some extra attributes:
* ``mitogen_message``: Unicode original message without the logger name
prepended.
* ``mitogen_context``: :class:`mitogen.parent.Context` reference to the
source context.
* ``mitogen_name``: Original logger name.
:param mitogen.master.Router router:
Router to install the handler on.
"""
def __init__(self, router):
self._router = router
self._cache = {}
router.add_handler(
fn=self._on_forward_log,
handle=mitogen.core.FORWARD_LOG,
)
def _on_forward_log(self, msg):
if msg.is_dead:
return
context = self._router.context_by_id(msg.src_id)
if context is None:
LOG.error('%s: dropping log from unknown context %d',
self, msg.src_id)
return
name, level_s, s = msg.data.decode('utf-8', 'replace').split('\x00', 2)
logger_name = '%s.[%s]' % (name, context.name)
logger = self._cache.get(logger_name)
if logger is None:
self._cache[logger_name] = logger = logging.getLogger(logger_name)
# See logging.Handler.makeRecord()
record = logging.LogRecord(
name=logger.name,
level=int(level_s),
pathname='(unknown file)',
lineno=0,
msg=s,
args=(),
exc_info=None,
)
record.mitogen_message = s
record.mitogen_context = self._router.context_by_id(msg.src_id)
record.mitogen_name = name
logger.handle(record)
def __repr__(self):
return 'LogForwarder(%r)' % (self._router,)
class FinderMethod(object):
"""
Interface to a method for locating a Python module or package given its
name according to the running Python interpreter. You'd think this was a
simple task, right? Naive young fellow, welcome to the real world.
"""
def __repr__(self):
return '%s()' % (type(self).__name__,)
def find(self, fullname):
"""
Accept a canonical module name as would be found in :data:`sys.modules`
and return a `(path, source, is_pkg)` tuple, where:
* `path`: Unicode string containing path to source file.
* `source`: Bytestring containing source file's content.
* `is_pkg`: :data:`True` if `fullname` is a package.
:returns:
:data:`None` if not found, or tuple as described above.
"""
raise NotImplementedError()
class DefectivePython3xMainMethod(FinderMethod):
"""
Recent versions of Python 3.x introduced an incomplete notion of
importer specs, and in doing so created permanent asymmetry in the
:mod:`pkgutil` interface handling for the :mod:`__main__` module. Therefore
we must handle :mod:`__main__` specially.
"""
def find(self, fullname):
"""
Find :mod:`__main__` using its :data:`__file__` attribute.
"""
if fullname != '__main__':
return None
mod = sys.modules.get(fullname)
if not mod:
return None
path = getattr(mod, '__file__', None)
if not (path is not None and os.path.exists(path) and _looks_like_script(path)):
return None
fp = open(path, 'rb')
try:
source = fp.read()
finally:
fp.close()
return path, source, False
class PkgutilMethod(FinderMethod):
"""
Attempt to fetch source code via pkgutil. In an ideal world, this would
be the only required implementation of get_module().
"""
def find(self, fullname):
"""
Find `fullname` using :func:`pkgutil.find_loader`.
"""
try:
# Pre-'import spec' this returned None, in Python3.6 it raises
# ImportError.
loader = pkgutil.find_loader(fullname)
except ImportError:
e = sys.exc_info()[1]
LOG.debug('%r._get_module_via_pkgutil(%r): %s',
self, fullname, e)
return None
IOLOG.debug('%r._get_module_via_pkgutil(%r) -> %r',
self, fullname, loader)
if not loader:
return
try:
path = _py_filename(loader.get_filename(fullname))
source = loader.get_source(fullname)
is_pkg = loader.is_package(fullname)
except (AttributeError, ImportError):
# - Per PEP-302, get_source() and is_package() are optional,
# calling them may throw AttributeError.
# - get_filename() may throw ImportError if pkgutil.find_loader()
# picks a "parent" package's loader for some crap that's been
# stuffed in sys.modules, for example in the case of urllib3:
# "loader for urllib3.contrib.pyopenssl cannot handle
# requests.packages.urllib3.contrib.pyopenssl"
e = sys.exc_info()[1]
LOG.debug('%r: loading %r using %r failed: %s',
self, fullname, loader, e)
return
if path is None or source is None:
return
if isinstance(source, mitogen.core.UnicodeType):
# get_source() returns "string" according to PEP-302, which was
# reinterpreted for Python 3 to mean a Unicode string.
source = source.encode('utf-8')
return path, source, is_pkg
class SysModulesMethod(FinderMethod):
"""
Attempt to fetch source code via :data:`sys.modules`. This was originally
specifically to support :mod:`__main__`, but it may catch a few more cases.
"""
def find(self, fullname):
"""
Find `fullname` using its :data:`__file__` attribute.
"""
module = sys.modules.get(fullname)
if not isinstance(module, types.ModuleType):
LOG.debug('%r: sys.modules[%r] absent or not a regular module',
self, fullname)
return
LOG.debug('_get_module_via_sys_modules(%r) -> %r', fullname, module)
alleged_name = getattr(module, '__name__', None)
if alleged_name != fullname:
LOG.debug('sys.modules[%r].__name__ is incorrect, assuming '
'this is a hacky module alias and ignoring it. '
'Got %r, module object: %r',
fullname, alleged_name, module)
return
path = _py_filename(getattr(module, '__file__', ''))
if not path:
return
LOG.debug('%r: sys.modules[%r]: found %s', self, fullname, path)
is_pkg = hasattr(module, '__path__')
try:
source = inspect.getsource(module)
except IOError:
# Work around inspect.getsourcelines() bug for 0-byte __init__.py
# files.
if not is_pkg:
raise
source = '\n'
if isinstance(source, mitogen.core.UnicodeType):
# get_source() returns "string" according to PEP-302, which was
# reinterpreted for Python 3 to mean a Unicode string.
source = source.encode('utf-8')
return path, source, is_pkg
class ParentEnumerationMethod(FinderMethod):
"""
Attempt to fetch source code by examining the module's (hopefully less
insane) parent package, and if no insane parents exist, simply use
:mod:`sys.path` to search for it from scratch on the filesystem using the
normal Python lookup mechanism.
This is required for older versions of :mod:`ansible.compat.six`,
:mod:`plumbum.colors`, Ansible 2.8 :mod:`ansible.module_utils.distro` and
its submodule :mod:`ansible.module_utils.distro._distro`.
When some package dynamically replaces itself in :data:`sys.modules`, but
only conditionally according to some program logic, it is possible that
children may attempt to load modules and subpackages from it that can no
longer be resolved by examining a (corrupted) parent.
For cases like :mod:`ansible.module_utils.distro`, this must handle cases
where a package transmuted itself into a totally unrelated module during
import and vice versa, where :data:`sys.modules` is replaced with junk that
makes it impossible to discover the loaded module using the in-memory
module object or any parent package's :data:`__path__`, since they have all
been overwritten. Some men just want to watch the world burn.
"""
def _find_sane_parent(self, fullname):
"""
Iteratively search :data:`sys.modules` for the least indirect parent of
`fullname` that is loaded and contains a :data:`__path__` attribute.
:return:
`(parent_name, path, modpath)` tuple, where:
* `modname`: canonical name of the found package, or the empty
string if none is found.
* `search_path`: :data:`__path__` attribute of the least
indirect parent found, or :data:`None` if no indirect parent
was found.
* `modpath`: list of module name components leading from `path`
to the target module.
"""
path = None
modpath = []
while True:
pkgname, _, modname = str_rpartition(to_text(fullname), u'.')
modpath.insert(0, modname)
if not pkgname:
return [], None, modpath
pkg = sys.modules.get(pkgname)
path = getattr(pkg, '__path__', None)
if pkg and path:
return pkgname.split('.'), path, modpath
LOG.debug('%r: %r lacks __path__ attribute', self, pkgname)
fullname = pkgname
def _found_package(self, fullname, path):
path = os.path.join(path, '__init__.py')
LOG.debug('%r: %r is PKG_DIRECTORY: %r', self, fullname, path)
return self._found_module(
fullname=fullname,
path=path,
fp=open(path, 'rb'),
is_pkg=True,
)
def _found_module(self, fullname, path, fp, is_pkg=False):
try:
path = _py_filename(path)
if not path:
return
source = fp.read()
finally:
if fp:
fp.close()
if isinstance(source, mitogen.core.UnicodeType):
# get_source() returns "string" according to PEP-302, which was
# reinterpreted for Python 3 to mean a Unicode string.
source = source.encode('utf-8')
return path, source, is_pkg
def _find_one_component(self, modname, search_path):
try:
#fp, path, (suffix, _, kind) = imp.find_module(modname, search_path)
return imp.find_module(modname, search_path)
except ImportError:
e = sys.exc_info()[1]
LOG.debug('%r: imp.find_module(%r, %r) -> %s',
self, modname, [search_path], e)
return None
def find(self, fullname):
"""
See implementation for a description of how this works.
"""
#if fullname not in sys.modules:
# Don't attempt this unless a module really exists in sys.modules,
# else we could return junk.
#return
fullname = to_text(fullname)
modname, search_path, modpath = self._find_sane_parent(fullname)
while True:
tup = self._find_one_component(modpath.pop(0), search_path)
if tup is None:
return None
fp, path, (suffix, _, kind) = tup
if modpath:
# Still more components to descent. Result must be a package
if fp:
fp.close()
if kind != imp.PKG_DIRECTORY:
LOG.debug('%r: %r appears to be child of non-package %r',
self, fullname, path)
return None
search_path = [path]
elif kind == imp.PKG_DIRECTORY:
return self._found_package(fullname, path)
else:
return self._found_module(fullname, path, fp)
class ModuleFinder(object):
"""
Given the name of a loaded module, make a best-effort attempt at finding
related modules likely needed by a child context requesting the original
module.
"""
def __init__(self):
#: Import machinery is expensive, keep :py:meth`:get_module_source`
#: results around.
self._found_cache = {}
#: Avoid repeated dependency scanning, which is expensive.
self._related_cache = {}
def __repr__(self):
return 'ModuleFinder()'
def add_source_override(self, fullname, path, source, is_pkg):
"""
Explicitly install a source cache entry, preventing usual lookup
methods from being used.
Beware the value of `path` is critical when `is_pkg` is specified,
since it directs where submodules are searched for.
:param str fullname:
Name of the module to override.
:param str path:
Module's path as it will appear in the cache.
:param bytes source:
Module source code as a bytestring.
:param bool is_pkg:
:data:`True` if the module is a package.
"""
self._found_cache[fullname] = (path, source, is_pkg)
get_module_methods = [
DefectivePython3xMainMethod(),
PkgutilMethod(),
SysModulesMethod(),
ParentEnumerationMethod(),
]
def get_module_source(self, fullname):
"""
Given the name of a loaded module `fullname`, attempt to find its
source code.
:returns:
Tuple of `(module path, source text, is package?)`, or :data:`None`
if the source cannot be found.
"""
tup = self._found_cache.get(fullname)
if tup:
return tup
for method in self.get_module_methods:
tup = method.find(fullname)
if tup:
#LOG.debug('%r returned %r', method, tup)
break
else:
tup = None, None, None
LOG.debug('get_module_source(%r): cannot find source', fullname)
self._found_cache[fullname] = tup
return tup
def resolve_relpath(self, fullname, level):
"""
Given an ImportFrom AST node, guess the prefix that should be tacked on
to an alias name to produce a canonical name. `fullname` is the name of
the module in which the ImportFrom appears.
"""
mod = sys.modules.get(fullname, None)
if hasattr(mod, '__path__'):
fullname += '.__init__'
if level == 0 or not fullname:
return ''
bits = fullname.split('.')
if len(bits) <= level:
# This would be an ImportError in real code.
return ''
return '.'.join(bits[:-level]) + '.'
def generate_parent_names(self, fullname):
while '.' in fullname:
fullname, _, _ = str_rpartition(to_text(fullname), u'.')
yield fullname
def find_related_imports(self, fullname):
"""
Return a list of non-stdlib modules that are directly imported by
`fullname`, plus their parents.
The list is determined by retrieving the source code of
`fullname`, compiling it, and examining all IMPORT_NAME ops.
:param fullname: Fully qualified name of an *already imported* module
for which source code can be retrieved
:type fullname: str
"""
related = self._related_cache.get(fullname)
if related is not None:
return related
modpath, src, _ = self.get_module_source(fullname)
if src is None:
return []
maybe_names = list(self.generate_parent_names(fullname))
co = compile(src, modpath, 'exec')
for level, modname, namelist in scan_code_imports(co):
if level == -1:
modnames = [modname, '%s.%s' % (fullname, modname)]
else:
modnames = [
'%s%s' % (self.resolve_relpath(fullname, level), modname)
]
maybe_names.extend(modnames)
maybe_names.extend(
'%s.%s' % (mname, name)
for mname in modnames
for name in namelist
)
return self._related_cache.setdefault(fullname, sorted(
set(
mitogen.core.to_text(name)
for name in maybe_names
if sys.modules.get(name) is not None
and not is_stdlib_name(name)
and u'six.moves' not in name # TODO: crap
)
))
def find_related(self, fullname):
"""
Return a list of non-stdlib modules that are imported directly or
indirectly by `fullname`, plus their parents.
This method is like :py:meth:`find_related_imports`, but also
recursively searches any modules which are imported by `fullname`.
:param fullname: Fully qualified name of an *already imported* module
for which source code can be retrieved
:type fullname: str
"""
stack = [fullname]
found = set()
while stack:
name = stack.pop(0)
names = self.find_related_imports(name)
stack.extend(set(names).difference(set(found).union(stack)))
found.update(names)
found.discard(fullname)
return sorted(found)
class ModuleResponder(object):
def __init__(self, router):
self._log = logging.getLogger('mitogen.responder')
self._router = router
self._finder = ModuleFinder()
self._cache = {} # fullname -> pickled
self.blacklist = []
self.whitelist = ['']
#: Context -> set([fullname, ..])
self._forwarded_by_context = {}
#: Number of GET_MODULE messages received.
self.get_module_count = 0
#: Total time spent in uncached GET_MODULE.
self.get_module_secs = 0.0
#: Total time spent minifying modules.
self.minify_secs = 0.0
#: Number of successful LOAD_MODULE messages sent.
self.good_load_module_count = 0
#: Total bytes in successful LOAD_MODULE payloads.
self.good_load_module_size = 0
#: Number of negative LOAD_MODULE messages sent.
self.bad_load_module_count = 0
router.add_handler(
fn=self._on_get_module,
handle=mitogen.core.GET_MODULE,
)
def __repr__(self):
return 'ModuleResponder'
def add_source_override(self, fullname, path, source, is_pkg):
"""
See :meth:`ModuleFinder.add_source_override`.
"""
self._finder.add_source_override(fullname, path, source, is_pkg)
MAIN_RE = re.compile(b(r'^if\s+__name__\s*==\s*.__main__.\s*:'), re.M)
main_guard_msg = (
"A child context attempted to import __main__, however the main "
"module present in the master process lacks an execution guard. "
"Update %r to prevent unintended execution, using a guard like:\n"
"\n"
" if __name__ == '__main__':\n"
" # your code here.\n"
)
def whitelist_prefix(self, fullname):
if self.whitelist == ['']:
self.whitelist = ['mitogen']
self.whitelist.append(fullname)
def blacklist_prefix(self, fullname):
self.blacklist.append(fullname)
def neutralize_main(self, path, src):
"""
Given the source for the __main__ module, try to find where it begins
conditional execution based on a "if __name__ == '__main__'" guard, and
remove any code after that point.
"""
match = self.MAIN_RE.search(src)
if match:
return src[:match.start()]
if b('mitogen.main(') in src:
return src
self._log.error(self.main_guard_msg, path)
raise ImportError('refused')
def _make_negative_response(self, fullname):
return (fullname, None, None, None, ())
minify_safe_re = re.compile(b(r'\s+#\s*!mitogen:\s*minify_safe'))
def _build_tuple(self, fullname):
if fullname in self._cache:
return self._cache[fullname]
if mitogen.core.is_blacklisted_import(self, fullname):
raise ImportError('blacklisted')
path, source, is_pkg = self._finder.get_module_source(fullname)
if path and is_stdlib_path(path):
# Prevent loading of 2.x<->3.x stdlib modules! This costs one
# RTT per hit, so a client-side solution is also required.
self._log.debug('refusing to serve stdlib module %r', fullname)
tup = self._make_negative_response(fullname)
self._cache[fullname] = tup
return tup
if source is None:
# TODO: make this .warning() or similar again once importer has its
# own logging category.
self._log.debug('could not find source for %r', fullname)
tup = self._make_negative_response(fullname)
self._cache[fullname] = tup
return tup
if self.minify_safe_re.search(source):
# If the module contains a magic marker, it's safe to minify.
t0 = mitogen.core.now()
source = mitogen.minify.minimize_source(source).encode('utf-8')
self.minify_secs += mitogen.core.now() - t0
if is_pkg:
pkg_present = get_child_modules(path)
self._log.debug('%s is a package at %s with submodules %r',
fullname, path, pkg_present)
else:
pkg_present = None
if fullname == '__main__':
source = self.neutralize_main(path, source)
compressed = mitogen.core.Blob(zlib.compress(source, 9))
related = [
to_text(name)
for name in self._finder.find_related(fullname)
if not mitogen.core.is_blacklisted_import(self, name)
]
# 0:fullname 1:pkg_present 2:path 3:compressed 4:related
tup = (
to_text(fullname),
pkg_present,
to_text(path),
compressed,
related
)
self._cache[fullname] = tup
return tup
def _send_load_module(self, stream, fullname):
if fullname not in stream.protocol.sent_modules:
tup = self._build_tuple(fullname)
msg = mitogen.core.Message.pickled(
tup,
dst_id=stream.protocol.remote_id,
handle=mitogen.core.LOAD_MODULE,
)
self._log.debug('sending %s (%.2f KiB) to %s',
fullname, len(msg.data) / 1024.0, stream.name)
self._router._async_route(msg)
stream.protocol.sent_modules.add(fullname)
if tup[2] is not None:
self.good_load_module_count += 1
self.good_load_module_size += len(msg.data)
else:
self.bad_load_module_count += 1
def _send_module_load_failed(self, stream, fullname):
self.bad_load_module_count += 1
stream.protocol.send(
mitogen.core.Message.pickled(
self._make_negative_response(fullname),
dst_id=stream.protocol.remote_id,
handle=mitogen.core.LOAD_MODULE,
)
)
def _send_module_and_related(self, stream, fullname):
if fullname in stream.protocol.sent_modules:
return
try:
tup = self._build_tuple(fullname)
for name in tup[4]: # related
parent, _, _ = str_partition(name, '.')
if parent != fullname and parent not in stream.protocol.sent_modules:
# Parent hasn't been sent, so don't load submodule yet.
continue
self._send_load_module(stream, name)
self._send_load_module(stream, fullname)
except Exception:
LOG.debug('While importing %r', fullname, exc_info=True)
self._send_module_load_failed(stream, fullname)
def _on_get_module(self, msg):
if msg.is_dead:
return
stream = self._router.stream_by_id(msg.src_id)
if stream is None:
return
fullname = msg.data.decode()
self._log.debug('%s requested module %s', stream.name, fullname)
self.get_module_count += 1
if fullname in stream.protocol.sent_modules:
LOG.warning('_on_get_module(): dup request for %r from %r',
fullname, stream)
t0 = mitogen.core.now()
try:
self._send_module_and_related(stream, fullname)
finally:
self.get_module_secs += mitogen.core.now() - t0
def _send_forward_module(self, stream, context, fullname):
if stream.protocol.remote_id != context.context_id:
stream.protocol._send(
mitogen.core.Message(
data=b('%s\x00%s' % (context.context_id, fullname)),
handle=mitogen.core.FORWARD_MODULE,
dst_id=stream.protocol.remote_id,
)
)
def _forward_one_module(self, context, fullname):
forwarded = self._forwarded_by_context.get(context)
if forwarded is None:
forwarded = set()
self._forwarded_by_context[context] = forwarded
if fullname in forwarded:
return
path = []
while fullname:
path.append(fullname)
fullname, _, _ = str_rpartition(fullname, u'.')
stream = self._router.stream_by_id(context.context_id)
if stream is None:
LOG.debug('%r: dropping forward of %s to no longer existent '
'%r', self, path[0], context)
return
for fullname in reversed(path):
self._send_module_and_related(stream, fullname)
self._send_forward_module(stream, context, fullname)
def _forward_modules(self, context, fullnames):
IOLOG.debug('%r._forward_modules(%r, %r)', self, context, fullnames)
for fullname in fullnames:
self._forward_one_module(context, mitogen.core.to_text(fullname))
def forward_modules(self, context, fullnames):
self._router.broker.defer(self._forward_modules, context, fullnames)
class Broker(mitogen.core.Broker):
"""
.. note::
You may construct as many brokers as desired, and use the same broker
for multiple routers, however usually only one broker need exist.
Multiple brokers may be useful when dealing with sets of children with
differing lifetimes. For example, a subscription service where
non-payment results in termination for one customer.
:param bool install_watcher:
If :data:`True`, an additional thread is started to monitor the
lifetime of the main thread, triggering :meth:`shutdown`
automatically in case the user forgets to call it, or their code
crashed.
You should not rely on this functionality in your program, it is only
intended as a fail-safe and to simplify the API for new users. In
particular, alternative Python implementations may not be able to
support watching the main thread.
"""
shutdown_timeout = 5.0
_watcher = None
poller_class = mitogen.parent.PREFERRED_POLLER
def __init__(self, install_watcher=True):
if install_watcher:
self._watcher = ThreadWatcher.watch(
target=threading.currentThread(),
on_join=self.shutdown,
)
super(Broker, self).__init__()
self.timers = mitogen.parent.TimerList()
def shutdown(self):
super(Broker, self).shutdown()
if self._watcher:
self._watcher.remove()
class Router(mitogen.parent.Router):
"""
Extend :class:`mitogen.core.Router` with functionality useful to masters,
and child contexts who later become masters. Currently when this class is
required, the target context's router is upgraded at runtime.
.. note::
You may construct as many routers as desired, and use the same broker
for multiple routers, however usually only one broker and router need
exist. Multiple routers may be useful when dealing with separate trust
domains, for example, manipulating infrastructure belonging to separate
customers or projects.
:param mitogen.master.Broker broker:
Broker to use. If not specified, a private :class:`Broker` is created.
:param int max_message_size:
Override the maximum message size this router is willing to receive or
transmit. Any value set here is automatically inherited by any children
created by the router.
This has a liberal default of 128 MiB, but may be set much lower.
Beware that setting it below 64KiB may encourage unexpected failures as
parents and children can no longer route large Python modules that may
be required by your application.
"""
broker_class = Broker
#: When :data:`True`, cause the broker thread and any subsequent broker and
#: main threads existing in any child to write
#: ``/tmp/mitogen.stats.<pid>.<thread_name>.log`` containing a
#: :mod:`cProfile` dump on graceful exit. Must be set prior to construction
#: of any :class:`Broker`, e.g. via::
#:
#: mitogen.master.Router.profiling = True
profiling = os.environ.get('MITOGEN_PROFILING') is not None
def __init__(self, broker=None, max_message_size=None):
if broker is None:
broker = self.broker_class()
if max_message_size:
self.max_message_size = max_message_size
super(Router, self).__init__(broker)
self.upgrade()
def upgrade(self):
self.id_allocator = IdAllocator(self)
self.responder = ModuleResponder(self)
self.log_forwarder = LogForwarder(self)
self.route_monitor = mitogen.parent.RouteMonitor(router=self)
self.add_handler( # TODO: cutpaste.
fn=self._on_detaching,
handle=mitogen.core.DETACHING,
persist=True,
)
def _on_broker_exit(self):
super(Router, self)._on_broker_exit()
dct = self.get_stats()
dct['self'] = self
dct['minify_ms'] = 1000 * dct['minify_secs']
dct['get_module_ms'] = 1000 * dct['get_module_secs']
dct['good_load_module_size_kb'] = dct['good_load_module_size'] / 1024.0
dct['good_load_module_size_avg'] = (
(
dct['good_load_module_size'] /
(float(dct['good_load_module_count']) or 1.0)
) / 1024.0
)
LOG.debug(
'%(self)r: stats: '
'%(get_module_count)d module requests in '
'%(get_module_ms)d ms, '
'%(good_load_module_count)d sent '
'(%(minify_ms)d ms minify time), '
'%(bad_load_module_count)d negative responses. '
'Sent %(good_load_module_size_kb).01f kb total, '
'%(good_load_module_size_avg).01f kb avg.'
% dct
)
def get_stats(self):
"""
Return performance data for the module responder.
:returns:
Dict containing keys:
* `get_module_count`: Integer count of
:data:`mitogen.core.GET_MODULE` messages received.
* `get_module_secs`: Floating point total seconds spent servicing
:data:`mitogen.core.GET_MODULE` requests.
* `good_load_module_count`: Integer count of successful
:data:`mitogen.core.LOAD_MODULE` messages sent.
* `good_load_module_size`: Integer total bytes sent in
:data:`mitogen.core.LOAD_MODULE` message payloads.
* `bad_load_module_count`: Integer count of negative
:data:`mitogen.core.LOAD_MODULE` messages sent.
* `minify_secs`: CPU seconds spent minifying modules marked
minify-safe.
"""
return {
'get_module_count': self.responder.get_module_count,
'get_module_secs': self.responder.get_module_secs,
'good_load_module_count': self.responder.good_load_module_count,
'good_load_module_size': self.responder.good_load_module_size,
'bad_load_module_count': self.responder.bad_load_module_count,
'minify_secs': self.responder.minify_secs,
}
def enable_debug(self):
"""
Cause this context and any descendant child contexts to write debug
logs to ``/tmp/mitogen.<pid>.log``.
"""
mitogen.core.enable_debug_logging()
self.debug = True
def __enter__(self):
return self
def __exit__(self, e_type, e_val, tb):
self.broker.shutdown()
self.broker.join()
def disconnect_stream(self, stream):
self.broker.defer(stream.on_disconnect, self.broker)
def disconnect_all(self):
for stream in self._stream_by_id.values():
self.disconnect_stream(stream)
class IdAllocator(object):
"""
Allocate IDs for new contexts constructed locally, and blocks of IDs for
children to allocate their own IDs using
:class:`mitogen.parent.ChildIdAllocator` without risk of conflict, and
without necessitating network round-trips for each new context.
This class responds to :data:`mitogen.core.ALLOCATE_ID` messages received
from children by replying with fresh block ID allocations.
The master's :class:`IdAllocator` instance can be accessed via
:attr:`mitogen.master.Router.id_allocator`.
"""
#: Block allocations are made in groups of 1000 by default.
BLOCK_SIZE = 1000
def __init__(self, router):
self.router = router
self.next_id = 1
self.lock = threading.Lock()
router.add_handler(
fn=self.on_allocate_id,
handle=mitogen.core.ALLOCATE_ID,
)
def __repr__(self):
return 'IdAllocator(%r)' % (self.router,)
def allocate(self):
"""
Allocate a context ID by directly incrementing an internal counter.
:returns:
The new context ID.
"""
self.lock.acquire()
try:
id_ = self.next_id
self.next_id += 1
return id_
finally:
self.lock.release()
def allocate_block(self):
"""
Allocate a block of IDs for use in a child context.
This function is safe to call from any thread.
:returns:
Tuple of the form `(id, end_id)` where `id` is the first usable ID
and `end_id` is the last usable ID.
"""
self.lock.acquire()
try:
id_ = self.next_id
self.next_id += self.BLOCK_SIZE
end_id = id_ + self.BLOCK_SIZE
LOG.debug('%r: allocating [%d..%d)', self, id_, end_id)
return id_, end_id
finally:
self.lock.release()
def on_allocate_id(self, msg):
if msg.is_dead:
return
id_, last_id = self.allocate_block()
requestee = self.router.context_by_id(msg.src_id)
LOG.debug('%r: allocating [%r..%r) to %r',
self, id_, last_id, requestee)
msg.reply((id_, last_id))
# Copyright 2017, Alex Willmer
#
# 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 sys
try:
from io import StringIO
except ImportError:
from StringIO import StringIO
import mitogen.core
if sys.version_info < (2, 7, 11):
from mitogen.compat import tokenize
else:
import tokenize
def minimize_source(source):
"""
Remove comments and docstrings from Python `source`, preserving line
numbers and syntax of empty blocks.
:param str source:
The source to minimize.
:returns str:
The minimized source.
"""
source = mitogen.core.to_text(source)
tokens = tokenize.generate_tokens(StringIO(source).readline)
tokens = strip_comments(tokens)
tokens = strip_docstrings(tokens)
tokens = reindent(tokens)
return tokenize.untokenize(tokens)
def strip_comments(tokens):
"""
Drop comment tokens from a `tokenize` stream.
Comments on lines 1-2 are kept, to preserve hashbang and encoding.
Trailing whitespace is remove from all lines.
"""
prev_typ = None
prev_end_col = 0
for typ, tok, (start_row, start_col), (end_row, end_col), line in tokens:
if typ in (tokenize.NL, tokenize.NEWLINE):
if prev_typ in (tokenize.NL, tokenize.NEWLINE):
start_col = 0
else:
start_col = prev_end_col
end_col = start_col + 1
elif typ == tokenize.COMMENT and start_row > 2:
continue
prev_typ = typ
prev_end_col = end_col
yield typ, tok, (start_row, start_col), (end_row, end_col), line
def strip_docstrings(tokens):
"""
Replace docstring tokens with NL tokens in a `tokenize` stream.
Any STRING token not part of an expression is deemed a docstring.
Indented docstrings are not yet recognised.
"""
stack = []
state = 'wait_string'
for t in tokens:
typ = t[0]
if state == 'wait_string':
if typ in (tokenize.NL, tokenize.COMMENT):
yield t
elif typ in (tokenize.DEDENT, tokenize.INDENT, tokenize.STRING):
stack.append(t)
elif typ == tokenize.NEWLINE:
stack.append(t)
start_line, end_line = stack[0][2][0], stack[-1][3][0]+1
for i in range(start_line, end_line):
yield tokenize.NL, '\n', (i, 0), (i,1), '\n'
for t in stack:
if t[0] in (tokenize.DEDENT, tokenize.INDENT):
yield t[0], t[1], (i+1, t[2][1]), (i+1, t[3][1]), t[4]
del stack[:]
else:
stack.append(t)
for t in stack: yield t
del stack[:]
state = 'wait_newline'
elif state == 'wait_newline':
if typ == tokenize.NEWLINE:
state = 'wait_string'
yield t
def reindent(tokens, indent=' '):
"""
Replace existing indentation in a token steam, with `indent`.
"""
old_levels = []
old_level = 0
new_level = 0
for typ, tok, (start_row, start_col), (end_row, end_col), line in tokens:
if typ == tokenize.INDENT:
old_levels.append(old_level)
old_level = len(tok)
new_level += 1
tok = indent * new_level
elif typ == tokenize.DEDENT:
old_level = old_levels.pop()
new_level -= 1
start_col = max(0, start_col - old_level + new_level)
if start_row == end_row:
end_col = start_col + len(tok)
yield typ, tok, (start_row, start_col), (end_row, end_col), line
# 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
"""
Support for operating in a mixed threading/forking environment.
"""
import os
import socket
import sys
import threading
import weakref
import mitogen.core
# List of weakrefs. On Python 2.4, mitogen.core registers its Broker on this
# list and mitogen.service registers its Pool too.
_brokers = weakref.WeakKeyDictionary()
_pools = weakref.WeakKeyDictionary()
def _notice_broker_or_pool(obj):
"""
Used by :mod:`mitogen.core` and :mod:`mitogen.service` to automatically
register every broker and pool on Python 2.4/2.5.
"""
if isinstance(obj, mitogen.core.Broker):
_brokers[obj] = True
else:
_pools[obj] = True
def wrap_os__fork():
corker = Corker(
brokers=list(_brokers),
pools=list(_pools),
)
try:
corker.cork()
return os__fork()
finally:
corker.uncork()
# If Python 2.4/2.5 where threading state is not fixed up, subprocess.Popen()
# may still deadlock due to the broker thread. In this case, pause os.fork() so
# that all active threads are paused during fork.
if sys.version_info < (2, 6):
os__fork = os.fork
os.fork = wrap_os__fork
class Corker(object):
"""
Arrange for :class:`mitogen.core.Broker` and optionally
:class:`mitogen.service.Pool` to be temporarily "corked" while fork
operations may occur.
In a mixed threading/forking environment, it is critical no threads are
active at the moment of fork, as they could hold mutexes whose state is
unrecoverably snapshotted in the locked state in the fork child, causing
deadlocks at random future moments.
To ensure a target thread has all locks dropped, it is made to write a
large string to a socket with a small buffer that has :data:`os.O_NONBLOCK`
disabled. CPython will drop the GIL and enter the ``write()`` system call,
where it will block until the socket buffer is drained, or the write side
is closed.
:class:`mitogen.core.Poller` is used to ensure the thread really has
blocked outside any Python locks, by checking if the socket buffer has
started to fill.
Since this necessarily involves posting a message to every existent thread
and verifying acknowledgement, it will never be a fast operation.
This does not yet handle the case of corking being initiated from within a
thread that is also a cork target.
:param brokers:
Sequence of :class:`mitogen.core.Broker` instances to cork.
:param pools:
Sequence of :class:`mitogen.core.Pool` instances to cork.
"""
def __init__(self, brokers=(), pools=()):
self.brokers = brokers
self.pools = pools
def _do_cork(self, s, wsock):
try:
try:
while True:
# at least EINTR is possible. Do our best to keep handling
# outside the GIL in this case using sendall().
wsock.sendall(s)
except socket.error:
pass
finally:
wsock.close()
def _cork_one(self, s, obj):
"""
Construct a socketpair, saving one side of it, and passing the other to
`obj` to be written to by one of its threads.
"""
rsock, wsock = mitogen.parent.create_socketpair(size=4096)
mitogen.core.set_cloexec(rsock.fileno())
mitogen.core.set_cloexec(wsock.fileno())
mitogen.core.set_block(wsock) # gevent
self._rsocks.append(rsock)
obj.defer(self._do_cork, s, wsock)
def _verify_one(self, rsock):
"""
Pause until the socket `rsock` indicates readability, due to
:meth:`_do_cork` triggering a blocking write on another thread.
"""
poller = mitogen.core.Poller()
poller.start_receive(rsock.fileno())
try:
while True:
for fd in poller.poll():
return
finally:
poller.close()
def cork(self):
"""
Arrange for any associated brokers and pools to be paused with no locks
held. This will not return until each thread acknowledges it has ceased
execution.
"""
current = threading.currentThread()
s = mitogen.core.b('CORK') * ((128 // 4) * 1024)
self._rsocks = []
# Pools must be paused first, as existing work may require the
# participation of a broker in order to complete.
for pool in self.pools:
if not pool.closed:
for th in pool._threads:
if th != current:
self._cork_one(s, pool)
for broker in self.brokers:
if broker._alive:
if broker._thread != current:
self._cork_one(s, broker)
# Pause until we can detect every thread has entered write().
for rsock in self._rsocks:
self._verify_one(rsock)
def uncork(self):
"""
Arrange for paused threads to resume operation.
"""
for rsock in self._rsocks:
rsock.close()
# 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
"""
This module defines functionality common to master and parent processes. It is
sent to any child context that is due to become a parent, due to recursive
connection.
"""
import codecs
import errno
import fcntl
import getpass
import heapq
import inspect
import logging
import os
import re
import signal
import socket
import struct
import subprocess
import sys
import termios
import textwrap
import threading
import zlib
# Absolute imports for <2.5.
select = __import__('select')
try:
import thread
except ImportError:
import threading as thread
import mitogen.core
from mitogen.core import b
from mitogen.core import bytes_partition
from mitogen.core import IOLOG
LOG = logging.getLogger(__name__)
# #410: we must avoid the use of socketpairs if SELinux is enabled.
try:
fp = open('/sys/fs/selinux/enforce', 'rb')
try:
SELINUX_ENABLED = bool(int(fp.read()))
finally:
fp.close()
except IOError:
SELINUX_ENABLED = False
try:
next
except NameError:
# Python 2.4/2.5
from mitogen.core import next
itervalues = getattr(dict, 'itervalues', dict.values)
if mitogen.core.PY3:
xrange = range
closure_attr = '__closure__'
IM_SELF_ATTR = '__self__'
else:
closure_attr = 'func_closure'
IM_SELF_ATTR = 'im_self'
try:
SC_OPEN_MAX = os.sysconf('SC_OPEN_MAX')
except ValueError:
SC_OPEN_MAX = 1024
BROKER_SHUTDOWN_MSG = (
'Connection cancelled because the associated Broker began to shut down.'
)
OPENPTY_MSG = (
"Failed to create a PTY: %s. It is likely the maximum number of PTYs has "
"been reached. Consider increasing the 'kern.tty.ptmx_max' sysctl on OS "
"X, the 'kernel.pty.max' sysctl on Linux, or modifying your configuration "
"to avoid PTY use."
)
SYS_EXECUTABLE_MSG = (
"The Python sys.executable variable is unset, indicating Python was "
"unable to determine its original program name. Unless explicitly "
"configured otherwise, child contexts will be started using "
"'/usr/bin/python'"
)
_sys_executable_warning_logged = False
def _ioctl_cast(n):
"""
Linux ioctl() request parameter is unsigned, whereas on BSD/Darwin it is
signed. Until 2.5 Python exclusively implemented the BSD behaviour,
preventing use of large unsigned int requests like the TTY layer uses
below. So on 2.4, we cast our unsigned to look like signed for Python.
"""
if sys.version_info < (2, 5):
n, = struct.unpack('i', struct.pack('I', n))
return n
# If not :data:`None`, called prior to exec() of any new child process. Used by
# :func:`mitogen.utils.reset_affinity` to allow the child to be freely
# scheduled.
_preexec_hook = None
# Get PTY number; asm-generic/ioctls.h
LINUX_TIOCGPTN = _ioctl_cast(2147767344)
# Lock/unlock PTY; asm-generic/ioctls.h
LINUX_TIOCSPTLCK = _ioctl_cast(1074025521)
IS_LINUX = os.uname()[0] == 'Linux'
SIGNAL_BY_NUM = dict(
(getattr(signal, name), name)
for name in sorted(vars(signal), reverse=True)
if name.startswith('SIG') and not name.startswith('SIG_')
)
_core_source_lock = threading.Lock()
_core_source_partial = None
def get_log_level():
return (LOG.getEffectiveLevel() or logging.INFO)
def get_sys_executable():
"""
Return :data:`sys.executable` if it is set, otherwise return
``"/usr/bin/python"`` and log a warning.
"""
if sys.executable:
return sys.executable
global _sys_executable_warning_logged
if not _sys_executable_warning_logged:
LOG.warn(SYS_EXECUTABLE_MSG)
_sys_executable_warning_logged = True
return '/usr/bin/python'
def _get_core_source():
"""
In non-masters, simply fetch the cached mitogen.core source code via the
import mechanism. In masters, this function is replaced with a version that
performs minification directly.
"""
return inspect.getsource(mitogen.core)
def get_core_source_partial():
"""
_get_core_source() is expensive, even with @lru_cache in minify.py, threads
can enter it simultaneously causing severe slowdowns.
"""
global _core_source_partial
if _core_source_partial is None:
_core_source_lock.acquire()
try:
if _core_source_partial is None:
_core_source_partial = PartialZlib(
_get_core_source().encode('utf-8')
)
finally:
_core_source_lock.release()
return _core_source_partial
def get_default_remote_name():
"""
Return the default name appearing in argv[0] of remote machines.
"""
s = u'%s@%s:%d'
s %= (getpass.getuser(), socket.gethostname(), os.getpid())
# In mixed UNIX/Windows environments, the username may contain slashes.
return s.translate({
ord(u'\\'): ord(u'_'),
ord(u'/'): ord(u'_')
})
def is_immediate_child(msg, stream):
"""
Handler policy that requires messages to arrive only from immediately
connected children.
"""
return msg.src_id == stream.protocol.remote_id
def flags(names):
"""
Return the result of ORing a set of (space separated) :py:mod:`termios`
module constants together.
"""
return sum(getattr(termios, name, 0)
for name in names.split())
def cfmakeraw(tflags):
"""
Given a list returned by :py:func:`termios.tcgetattr`, return a list
modified in a manner similar to the `cfmakeraw()` C library function, but
additionally disabling local echo.
"""
# BSD: github.com/freebsd/freebsd/blob/master/lib/libc/gen/termios.c#L162
# Linux: github.com/lattera/glibc/blob/master/termios/cfmakeraw.c#L20
iflag, oflag, cflag, lflag, ispeed, ospeed, cc = tflags
iflag &= ~flags('IMAXBEL IXOFF INPCK BRKINT PARMRK '
'ISTRIP INLCR ICRNL IXON IGNPAR')
iflag &= ~flags('IGNBRK BRKINT PARMRK')
oflag &= ~flags('OPOST')
lflag &= ~flags('ECHO ECHOE ECHOK ECHONL ICANON ISIG '
'IEXTEN NOFLSH TOSTOP PENDIN')
cflag &= ~flags('CSIZE PARENB')
cflag |= flags('CS8 CREAD')
return [iflag, oflag, cflag, lflag, ispeed, ospeed, cc]
def disable_echo(fd):
old = termios.tcgetattr(fd)
new = cfmakeraw(old)
flags = getattr(termios, 'TCSASOFT', 0)
if not mitogen.core.IS_WSL:
# issue #319: Windows Subsystem for Linux as of July 2018 throws EINVAL
# if TCSAFLUSH is specified.
flags |= termios.TCSAFLUSH
termios.tcsetattr(fd, flags, new)
def create_socketpair(size=None):
"""
Create a :func:`socket.socketpair` for use as a child's UNIX stdio
channels. As socketpairs are bidirectional, they are economical on file
descriptor usage as one descriptor can be used for ``stdin`` and
``stdout``. As they are sockets their buffers are tunable, allowing large
buffers to improve file transfer throughput and reduce IO loop iterations.
"""
if size is None:
size = mitogen.core.CHUNK_SIZE
parentfp, childfp = socket.socketpair()
for fp in parentfp, childfp:
fp.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, size)
return parentfp, childfp
def create_best_pipe(escalates_privilege=False):
"""
By default we prefer to communicate with children over a UNIX socket, as a
single file descriptor can represent bidirectional communication, and a
cross-platform API exists to align buffer sizes with the needs of the
library.
SELinux prevents us setting up a privileged process to inherit an AF_UNIX
socket, a facility explicitly designed as a better replacement for pipes,
because at some point in the mid 90s it might have been commonly possible
for AF_INET sockets to end up undesirably connected to a privileged
process, so let's make up arbitrary rules breaking all sockets instead.
If SELinux is detected, fall back to using pipes.
:param bool escalates_privilege:
If :data:`True`, the target program may escalate privileges, causing
SELinux to disconnect AF_UNIX sockets, so avoid those.
:returns:
`(parent_rfp, child_wfp, child_rfp, parent_wfp)`
"""
if (not escalates_privilege) or (not SELINUX_ENABLED):
parentfp, childfp = create_socketpair()
return parentfp, childfp, childfp, parentfp
parent_rfp, child_wfp = mitogen.core.pipe()
try:
child_rfp, parent_wfp = mitogen.core.pipe()
return parent_rfp, child_wfp, child_rfp, parent_wfp
except:
parent_rfp.close()
child_wfp.close()
raise
def popen(**kwargs):
"""
Wrap :class:`subprocess.Popen` to ensure any global :data:`_preexec_hook`
is invoked in the child.
"""
real_preexec_fn = kwargs.pop('preexec_fn', None)
def preexec_fn():
if _preexec_hook:
_preexec_hook()
if real_preexec_fn:
real_preexec_fn()
return subprocess.Popen(preexec_fn=preexec_fn, **kwargs)
def create_child(args, merge_stdio=False, stderr_pipe=False,
escalates_privilege=False, preexec_fn=None):
"""
Create a child process whose stdin/stdout is connected to a socket.
:param list args:
Program argument vector.
:param bool merge_stdio:
If :data:`True`, arrange for `stderr` to be connected to the `stdout`
socketpair, rather than inherited from the parent process. This may be
necessary to ensure that no TTY is connected to any stdio handle, for
instance when using LXC.
:param bool stderr_pipe:
If :data:`True` and `merge_stdio` is :data:`False`, arrange for
`stderr` to be connected to a separate pipe, to allow any ongoing debug
logs generated by e.g. SSH to be output as the session progresses,
without interfering with `stdout`.
:param bool escalates_privilege:
If :data:`True`, the target program may escalate privileges, causing
SELinux to disconnect AF_UNIX sockets, so avoid those.
:param function preexec_fn:
If not :data:`None`, a function to run within the post-fork child
before executing the target program.
:returns:
:class:`Process` instance.
"""
parent_rfp, child_wfp, child_rfp, parent_wfp = create_best_pipe(
escalates_privilege=escalates_privilege
)
stderr = None
stderr_r = None
if merge_stdio:
stderr = child_wfp
elif stderr_pipe:
stderr_r, stderr = mitogen.core.pipe()
mitogen.core.set_cloexec(stderr_r.fileno())
try:
proc = popen(
args=args,
stdin=child_rfp,
stdout=child_wfp,
stderr=stderr,
close_fds=True,
preexec_fn=preexec_fn,
)
except:
child_rfp.close()
child_wfp.close()
parent_rfp.close()
parent_wfp.close()
if stderr_pipe:
stderr.close()
stderr_r.close()
raise
child_rfp.close()
child_wfp.close()
if stderr_pipe:
stderr.close()
return PopenProcess(
proc=proc,
stdin=parent_wfp,
stdout=parent_rfp,
stderr=stderr_r,
)
def _acquire_controlling_tty():
os.setsid()
if sys.platform in ('linux', 'linux2'):
# On Linux, the controlling tty becomes the first tty opened by a
# process lacking any prior tty.
os.close(os.open(os.ttyname(2), os.O_RDWR))
if hasattr(termios, 'TIOCSCTTY') and not mitogen.core.IS_WSL:
# #550: prehistoric WSL does not like TIOCSCTTY.
# On BSD an explicit ioctl is required. For some inexplicable reason,
# Python 2.6 on Travis also requires it.
fcntl.ioctl(2, termios.TIOCSCTTY)
def _linux_broken_devpts_openpty():
"""
#462: On broken Linux hosts with mismatched configuration (e.g. old
/etc/fstab template installed), /dev/pts may be mounted without the gid=
mount option, causing new slave devices to be created with the group ID of
the calling process. This upsets glibc, whose openpty() is required by
specification to produce a slave owned by a special group ID (which is
always the 'tty' group).
Glibc attempts to use "pt_chown" to fix ownership. If that fails, it
chown()s the PTY directly, which fails due to non-root, causing openpty()
to fail with EPERM ("Operation not permitted"). Since we don't need the
magical TTY group to run sudo and su, open the PTY ourselves in this case.
"""
master_fd = None
try:
# Opening /dev/ptmx causes a PTY pair to be allocated, and the
# corresponding slave /dev/pts/* device to be created, owned by UID/GID
# matching this process.
master_fd = os.open('/dev/ptmx', os.O_RDWR)
# Clear the lock bit from the PTY. This a prehistoric feature from a
# time when slave device files were persistent.
fcntl.ioctl(master_fd, LINUX_TIOCSPTLCK, struct.pack('i', 0))
# Since v4.13 TIOCGPTPEER exists to open the slave in one step, but we
# must support older kernels. Ask for the PTY number.
pty_num_s = fcntl.ioctl(master_fd, LINUX_TIOCGPTN,
struct.pack('i', 0))
pty_num, = struct.unpack('i', pty_num_s)
pty_name = '/dev/pts/%d' % (pty_num,)
# Now open it with O_NOCTTY to ensure it doesn't change our controlling
# TTY. Otherwise when we close the FD we get killed by the kernel, and
# the child we spawn that should really attach to it will get EPERM
# during _acquire_controlling_tty().
slave_fd = os.open(pty_name, os.O_RDWR|os.O_NOCTTY)
return master_fd, slave_fd
except OSError:
if master_fd is not None:
os.close(master_fd)
e = sys.exc_info()[1]
raise mitogen.core.StreamError(OPENPTY_MSG, e)
def openpty():
"""
Call :func:`os.openpty`, raising a descriptive error if the call fails.
:raises mitogen.core.StreamError:
Creating a PTY failed.
:returns:
`(master_fp, slave_fp)` file-like objects.
"""
try:
master_fd, slave_fd = os.openpty()
except OSError:
e = sys.exc_info()[1]
if not (IS_LINUX and e.args[0] == errno.EPERM):
raise mitogen.core.StreamError(OPENPTY_MSG, e)
master_fd, slave_fd = _linux_broken_devpts_openpty()
master_fp = os.fdopen(master_fd, 'r+b', 0)
slave_fp = os.fdopen(slave_fd, 'r+b', 0)
disable_echo(master_fd)
disable_echo(slave_fd)
mitogen.core.set_block(slave_fd)
return master_fp, slave_fp
def tty_create_child(args):
"""
Return a file descriptor connected to the master end of a pseudo-terminal,
whose slave end is connected to stdin/stdout/stderr of a new child process.
The child is created such that the pseudo-terminal becomes its controlling
TTY, ensuring access to /dev/tty returns a new file descriptor open on the
slave end.
:param list args:
Program argument vector.
:returns:
:class:`Process` instance.
"""
master_fp, slave_fp = openpty()
try:
proc = popen(
args=args,
stdin=slave_fp,
stdout=slave_fp,
stderr=slave_fp,
preexec_fn=_acquire_controlling_tty,
close_fds=True,
)
except:
master_fp.close()
slave_fp.close()
raise
slave_fp.close()
return PopenProcess(
proc=proc,
stdin=master_fp,
stdout=master_fp,
)
def hybrid_tty_create_child(args, escalates_privilege=False):
"""
Like :func:`tty_create_child`, except attach stdin/stdout to a socketpair
like :func:`create_child`, but leave stderr and the controlling TTY
attached to a TTY.
This permits high throughput communication with programs that are reached
via some program that requires a TTY for password input, like many
configurations of sudo. The UNIX TTY layer tends to have tiny (no more than
14KiB) buffers, forcing many IO loop iterations when transferring bulk
data, causing significant performance loss.
:param bool escalates_privilege:
If :data:`True`, the target program may escalate privileges, causing
SELinux to disconnect AF_UNIX sockets, so avoid those.
:param list args:
Program argument vector.
:returns:
:class:`Process` instance.
"""
master_fp, slave_fp = openpty()
try:
parent_rfp, child_wfp, child_rfp, parent_wfp = create_best_pipe(
escalates_privilege=escalates_privilege,
)
try:
mitogen.core.set_block(child_rfp)
mitogen.core.set_block(child_wfp)
proc = popen(
args=args,
stdin=child_rfp,
stdout=child_wfp,
stderr=slave_fp,
preexec_fn=_acquire_controlling_tty,
close_fds=True,
)
except:
parent_rfp.close()
child_wfp.close()
parent_wfp.close()
child_rfp.close()
raise
except:
master_fp.close()
slave_fp.close()
raise
slave_fp.close()
child_rfp.close()
child_wfp.close()
return PopenProcess(
proc=proc,
stdin=parent_wfp,
stdout=parent_rfp,
stderr=master_fp,
)
class Timer(object):
"""
Represents a future event.
"""
#: Set to :data:`False` if :meth:`cancel` has been called, or immediately
#: prior to being executed by :meth:`TimerList.expire`.
active = True
def __init__(self, when, func):
self.when = when
self.func = func
def __repr__(self):
return 'Timer(%r, %r)' % (self.when, self.func)
def __eq__(self, other):
return self.when == other.when
def __lt__(self, other):
return self.when < other.when
def __le__(self, other):
return self.when <= other.when
def cancel(self):
"""
Cancel this event. If it has not yet executed, it will not execute
during any subsequent :meth:`TimerList.expire` call.
"""
self.active = False
class TimerList(object):
"""
Efficiently manage a list of cancellable future events relative to wall
clock time. An instance of this class is installed as
:attr:`mitogen.master.Broker.timers` by default, and as
:attr:`mitogen.core.Broker.timers` in children after a call to
:func:`mitogen.parent.upgrade_router`.
You can use :class:`TimerList` to cause the broker to wake at arbitrary
future moments, useful for implementing timeouts and polling in an
asynchronous context.
:class:`TimerList` methods can only be called from asynchronous context,
for example via :meth:`mitogen.core.Broker.defer`.
The broker automatically adjusts its sleep delay according to the installed
timer list, and arranges for timers to expire via automatic calls to
:meth:`expire`. The main user interface to :class:`TimerList` is
:meth:`schedule`.
"""
_now = mitogen.core.now
def __init__(self):
self._lst = []
def get_timeout(self):
"""
Return the floating point seconds until the next event is due.
:returns:
Floating point delay, or 0.0, or :data:`None` if no events are
scheduled.
"""
while self._lst and not self._lst[0].active:
heapq.heappop(self._lst)
if self._lst:
return max(0, self._lst[0].when - self._now())
def schedule(self, when, func):
"""
Schedule a future event.
:param float when:
UNIX time in seconds when event should occur.
:param callable func:
Callable to invoke on expiry.
:returns:
A :class:`Timer` instance, exposing :meth:`Timer.cancel`, which may
be used to cancel the future invocation.
"""
timer = Timer(when, func)
heapq.heappush(self._lst, timer)
return timer
def expire(self):
"""
Invoke callbacks for any events in the past.
"""
now = self._now()
while self._lst and self._lst[0].when <= now:
timer = heapq.heappop(self._lst)
if timer.active:
timer.active = False
timer.func()
class PartialZlib(object):
"""
Because the mitogen.core source has a line appended to it during bootstrap,
it must be recompressed for each connection. This is not a problem for a
small number of connections, but it amounts to 30 seconds CPU time by the
time 500 targets are in use.
For that reason, build a compressor containing mitogen.core and flush as
much of it as possible into an initial buffer. Then to append the custom
line, clone the compressor and compress just that line.
A full compression costs ~6ms on a modern machine, this method costs ~35
usec.
"""
def __init__(self, s):
self.s = s
if sys.version_info > (2, 5):
self._compressor = zlib.compressobj(9)
self._out = self._compressor.compress(s)
self._out += self._compressor.flush(zlib.Z_SYNC_FLUSH)
else:
self._compressor = None
def append(self, s):
"""
Append the bytestring `s` to the compressor state and return the
final compressed output.
"""
if self._compressor is None:
return zlib.compress(self.s + s, 9)
else:
compressor = self._compressor.copy()
out = self._out
out += compressor.compress(s)
return out + compressor.flush()
def _upgrade_broker(broker):
"""
Extract the poller state from Broker and replace it with the industrial
strength poller for this OS. Must run on the Broker thread.
"""
# This function is deadly! The act of calling start_receive() generates log
# messages which must be silenced as the upgrade progresses, otherwise the
# poller state will change as it is copied, resulting in write fds that are
# lost. (Due to LogHandler->Router->Stream->Protocol->Broker->Poller, where
# Stream only calls start_transmit() when transitioning from empty to
# non-empty buffer. If the start_transmit() is lost, writes from the child
# hang permanently).
root = logging.getLogger()
old_level = root.level
root.setLevel(logging.CRITICAL)
try:
old = broker.poller
new = PREFERRED_POLLER()
for fd, data in old.readers:
new.start_receive(fd, data)
for fd, data in old.writers:
new.start_transmit(fd, data)
old.close()
broker.poller = new
finally:
root.setLevel(old_level)
broker.timers = TimerList()
LOG.debug('upgraded %r with %r (new: %d readers, %d writers; '
'old: %d readers, %d writers)', old, new,
len(new.readers), len(new.writers),
len(old.readers), len(old.writers))
@mitogen.core.takes_econtext
def upgrade_router(econtext):
if not isinstance(econtext.router, Router): # TODO
econtext.broker.defer(_upgrade_broker, econtext.broker)
econtext.router.__class__ = Router # TODO
econtext.router.upgrade(
importer=econtext.importer,
parent=econtext.parent,
)
def get_connection_class(name):
"""
Given the name of a Mitogen connection method, import its implementation
module and return its Stream subclass.
"""
if name == u'local':
name = u'parent'
module = mitogen.core.import_module(u'mitogen.' + name)
return module.Connection
@mitogen.core.takes_econtext
def _proxy_connect(name, method_name, kwargs, econtext):
"""
Implements the target portion of Router._proxy_connect() by upgrading the
local process to a parent if it was not already, then calling back into
Router._connect() using the arguments passed to the parent's
Router.connect().
:returns:
Dict containing:
* ``id``: :data:`None`, or integer new context ID.
* ``name``: :data:`None`, or string name attribute of new Context.
* ``msg``: :data:`None`, or StreamError exception text.
"""
upgrade_router(econtext)
try:
context = econtext.router._connect(
klass=get_connection_class(method_name),
name=name,
**kwargs
)
except mitogen.core.StreamError:
return {
u'id': None,
u'name': None,
u'msg': 'error occurred on host %s: %s' % (
socket.gethostname(),
sys.exc_info()[1],
),
}
return {
u'id': context.context_id,
u'name': context.name,
u'msg': None,
}
def returncode_to_str(n):
"""
Parse and format a :func:`os.waitpid` exit status.
"""
if n < 0:
return 'exited due to signal %d (%s)' % (-n, SIGNAL_BY_NUM.get(-n))
return 'exited with return code %d' % (n,)
class EofError(mitogen.core.StreamError):
"""
Raised by :class:`Connection` when an empty read is detected from the
remote process before bootstrap completes.
"""
# inherits from StreamError to maintain compatibility.
pass
class CancelledError(mitogen.core.StreamError):
"""
Raised by :class:`Connection` when :meth:`mitogen.core.Broker.shutdown` is
called before bootstrap completes.
"""
pass
class Argv(object):
"""
Wrapper to defer argv formatting when debug logging is disabled.
"""
def __init__(self, argv):
self.argv = argv
must_escape = frozenset('\\$"`!')
must_escape_or_space = must_escape | frozenset(' ')
def escape(self, x):
if not self.must_escape_or_space.intersection(x):
return x
s = '"'
for c in x:
if c in self.must_escape:
s += '\\'
s += c
s += '"'
return s
def __str__(self):
return ' '.join(map(self.escape, self.argv))
class CallSpec(object):
"""
Wrapper to defer call argument formatting when debug logging is disabled.
"""
def __init__(self, func, args, kwargs):
self.func = func
self.args = args
self.kwargs = kwargs
def _get_name(self):
bits = [self.func.__module__]
if inspect.ismethod(self.func):
im_self = getattr(self.func, IM_SELF_ATTR)
bits.append(getattr(im_self, '__name__', None) or
getattr(type(im_self), '__name__', None))
bits.append(self.func.__name__)
return u'.'.join(bits)
def _get_args(self):
return u', '.join(repr(a) for a in self.args)
def _get_kwargs(self):
s = u''
if self.kwargs:
s = u', '.join('%s=%r' % (k, v) for k, v in self.kwargs.items())
if self.args:
s = u', ' + s
return s
def __repr__(self):
return '%s(%s%s)' % (
self._get_name(),
self._get_args(),
self._get_kwargs(),
)
class PollPoller(mitogen.core.Poller):
"""
Poller based on the POSIX :linux:man2:`poll` interface. Not available on
some versions of OS X, otherwise it is the preferred poller for small FD
counts, as there is no setup/teardown/configuration system call overhead.
"""
SUPPORTED = hasattr(select, 'poll')
_repr = 'PollPoller()'
def __init__(self):
super(PollPoller, self).__init__()
self._pollobj = select.poll()
# TODO: no proof we dont need writemask too
_readmask = (
getattr(select, 'POLLIN', 0) |
getattr(select, 'POLLHUP', 0)
)
def _update(self, fd):
mask = (((fd in self._rfds) and self._readmask) |
((fd in self._wfds) and select.POLLOUT))
if mask:
self._pollobj.register(fd, mask)
else:
try:
self._pollobj.unregister(fd)
except KeyError:
pass
def _poll(self, timeout):
if timeout:
timeout *= 1000
events, _ = mitogen.core.io_op(self._pollobj.poll, timeout)
for fd, event in events:
if event & self._readmask:
IOLOG.debug('%r: POLLIN|POLLHUP for %r', self, fd)
data, gen = self._rfds.get(fd, (None, None))
if gen and gen < self._generation:
yield data
if event & select.POLLOUT:
IOLOG.debug('%r: POLLOUT for %r', self, fd)
data, gen = self._wfds.get(fd, (None, None))
if gen and gen < self._generation:
yield data
class KqueuePoller(mitogen.core.Poller):
"""
Poller based on the FreeBSD/Darwin :freebsd:man2:`kqueue` interface.
"""
SUPPORTED = hasattr(select, 'kqueue')
_repr = 'KqueuePoller()'
def __init__(self):
super(KqueuePoller, self).__init__()
self._kqueue = select.kqueue()
self._changelist = []
def close(self):
super(KqueuePoller, self).close()
self._kqueue.close()
def _control(self, fd, filters, flags):
mitogen.core._vv and IOLOG.debug(
'%r._control(%r, %r, %r)', self, fd, filters, flags)
# TODO: at shutdown it is currently possible for KQ_EV_ADD/KQ_EV_DEL
# pairs to be pending after the associated file descriptor has already
# been closed. Fixing this requires maintaining extra state, or perhaps
# making fd closure the poller's responsibility. In the meantime,
# simply apply changes immediately.
# self._changelist.append(select.kevent(fd, filters, flags))
changelist = [select.kevent(fd, filters, flags)]
events, _ = mitogen.core.io_op(self._kqueue.control, changelist, 0, 0)
assert not events
def start_receive(self, fd, data=None):
mitogen.core._vv and IOLOG.debug('%r.start_receive(%r, %r)',
self, fd, data)
if fd not in self._rfds:
self._control(fd, select.KQ_FILTER_READ, select.KQ_EV_ADD)
self._rfds[fd] = (data or fd, self._generation)
def stop_receive(self, fd):
mitogen.core._vv and IOLOG.debug('%r.stop_receive(%r)', self, fd)
if fd in self._rfds:
self._control(fd, select.KQ_FILTER_READ, select.KQ_EV_DELETE)
del self._rfds[fd]
def start_transmit(self, fd, data=None):
mitogen.core._vv and IOLOG.debug('%r.start_transmit(%r, %r)',
self, fd, data)
if fd not in self._wfds:
self._control(fd, select.KQ_FILTER_WRITE, select.KQ_EV_ADD)
self._wfds[fd] = (data or fd, self._generation)
def stop_transmit(self, fd):
mitogen.core._vv and IOLOG.debug('%r.stop_transmit(%r)', self, fd)
if fd in self._wfds:
self._control(fd, select.KQ_FILTER_WRITE, select.KQ_EV_DELETE)
del self._wfds[fd]
def _poll(self, timeout):
changelist = self._changelist
self._changelist = []
events, _ = mitogen.core.io_op(self._kqueue.control,
changelist, 32, timeout)
for event in events:
fd = event.ident
if event.flags & select.KQ_EV_ERROR:
LOG.debug('ignoring stale event for fd %r: errno=%d: %s',
fd, event.data, errno.errorcode.get(event.data))
elif event.filter == select.KQ_FILTER_READ:
data, gen = self._rfds.get(fd, (None, None))
# Events can still be read for an already-discarded fd.
if gen and gen < self._generation:
mitogen.core._vv and IOLOG.debug('%r: POLLIN: %r', self, fd)
yield data
elif event.filter == select.KQ_FILTER_WRITE and fd in self._wfds:
data, gen = self._wfds.get(fd, (None, None))
if gen and gen < self._generation:
mitogen.core._vv and IOLOG.debug('%r: POLLOUT: %r', self, fd)
yield data
class EpollPoller(mitogen.core.Poller):
"""
Poller based on the Linux :linux:man2:`epoll` interface.
"""
SUPPORTED = hasattr(select, 'epoll')
_repr = 'EpollPoller()'
def __init__(self):
super(EpollPoller, self).__init__()
self._epoll = select.epoll(32)
self._registered_fds = set()
def close(self):
super(EpollPoller, self).close()
self._epoll.close()
def _control(self, fd):
mitogen.core._vv and IOLOG.debug('%r._control(%r)', self, fd)
mask = (((fd in self._rfds) and select.EPOLLIN) |
((fd in self._wfds) and select.EPOLLOUT))
if mask:
if fd in self._registered_fds:
self._epoll.modify(fd, mask)
else:
self._epoll.register(fd, mask)
self._registered_fds.add(fd)
elif fd in self._registered_fds:
self._epoll.unregister(fd)
self._registered_fds.remove(fd)
def start_receive(self, fd, data=None):
mitogen.core._vv and IOLOG.debug('%r.start_receive(%r, %r)',
self, fd, data)
self._rfds[fd] = (data or fd, self._generation)
self._control(fd)
def stop_receive(self, fd):
mitogen.core._vv and IOLOG.debug('%r.stop_receive(%r)', self, fd)
self._rfds.pop(fd, None)
self._control(fd)
def start_transmit(self, fd, data=None):
mitogen.core._vv and IOLOG.debug('%r.start_transmit(%r, %r)',
self, fd, data)
self._wfds[fd] = (data or fd, self._generation)
self._control(fd)
def stop_transmit(self, fd):
mitogen.core._vv and IOLOG.debug('%r.stop_transmit(%r)', self, fd)
self._wfds.pop(fd, None)
self._control(fd)
_inmask = (getattr(select, 'EPOLLIN', 0) |
getattr(select, 'EPOLLHUP', 0))
def _poll(self, timeout):
the_timeout = -1
if timeout is not None:
the_timeout = timeout
events, _ = mitogen.core.io_op(self._epoll.poll, the_timeout, 32)
for fd, event in events:
if event & self._inmask:
data, gen = self._rfds.get(fd, (None, None))
if gen and gen < self._generation:
# Events can still be read for an already-discarded fd.
mitogen.core._vv and IOLOG.debug('%r: POLLIN: %r', self, fd)
yield data
if event & select.EPOLLOUT:
data, gen = self._wfds.get(fd, (None, None))
if gen and gen < self._generation:
mitogen.core._vv and IOLOG.debug('%r: POLLOUT: %r', self, fd)
yield data
# 2.4 and 2.5 only had select.select() and select.poll().
for _klass in mitogen.core.Poller, PollPoller, KqueuePoller, EpollPoller:
if _klass.SUPPORTED:
PREFERRED_POLLER = _klass
# For processes that start many threads or connections, it's possible Latch
# will also get high-numbered FDs, and so select() becomes useless there too.
# So swap in our favourite poller.
if PollPoller.SUPPORTED:
mitogen.core.Latch.poller_class = PollPoller
else:
mitogen.core.Latch.poller_class = PREFERRED_POLLER
class LineLoggingProtocolMixin(object):
def __init__(self, **kwargs):
super(LineLoggingProtocolMixin, self).__init__(**kwargs)
self.logged_lines = []
self.logged_partial = None
def on_line_received(self, line):
self.logged_partial = None
self.logged_lines.append((mitogen.core.now(), line))
self.logged_lines[:] = self.logged_lines[-100:]
return super(LineLoggingProtocolMixin, self).on_line_received(line)
def on_partial_line_received(self, line):
self.logged_partial = line
return super(LineLoggingProtocolMixin, self).on_partial_line_received(line)
def on_disconnect(self, broker):
if self.logged_partial:
self.logged_lines.append((mitogen.core.now(), self.logged_partial))
self.logged_partial = None
super(LineLoggingProtocolMixin, self).on_disconnect(broker)
def get_history(streams):
history = []
for stream in streams:
if stream:
history.extend(getattr(stream.protocol, 'logged_lines', []))
history.sort()
s = b('\n').join(h[1] for h in history)
return mitogen.core.to_text(s)
class RegexProtocol(LineLoggingProtocolMixin, mitogen.core.DelimitedProtocol):
"""
Implement a delimited protocol where messages matching a set of regular
expressions are dispatched to individual handler methods. Input is
dispatches using :attr:`PATTERNS` and :attr:`PARTIAL_PATTERNS`, before
falling back to :meth:`on_unrecognized_line_received` and
:meth:`on_unrecognized_partial_line_received`.
"""
#: A sequence of 2-tuples of the form `(compiled pattern, method)` for
#: patterns that should be matched against complete (delimited) messages,
#: i.e. full lines.
PATTERNS = []
#: Like :attr:`PATTERNS`, but patterns that are matched against incomplete
#: lines.
PARTIAL_PATTERNS = []
def on_line_received(self, line):
super(RegexProtocol, self).on_line_received(line)
for pattern, func in self.PATTERNS:
match = pattern.search(line)
if match is not None:
return func(self, line, match)
return self.on_unrecognized_line_received(line)
def on_unrecognized_line_received(self, line):
LOG.debug('%s: (unrecognized): %s',
self.stream.name, line.decode('utf-8', 'replace'))
def on_partial_line_received(self, line):
super(RegexProtocol, self).on_partial_line_received(line)
LOG.debug('%s: (partial): %s',
self.stream.name, line.decode('utf-8', 'replace'))
for pattern, func in self.PARTIAL_PATTERNS:
match = pattern.search(line)
if match is not None:
return func(self, line, match)
return self.on_unrecognized_partial_line_received(line)
def on_unrecognized_partial_line_received(self, line):
LOG.debug('%s: (unrecognized partial): %s',
self.stream.name, line.decode('utf-8', 'replace'))
class BootstrapProtocol(RegexProtocol):
"""
Respond to stdout of a child during bootstrap. Wait for :attr:`EC0_MARKER`
to be written by the first stage to indicate it can receive the bootstrap,
then await :attr:`EC1_MARKER` to indicate success, and
:class:`MitogenProtocol` can be enabled.
"""
#: Sentinel value emitted by the first stage to indicate it is ready to
#: receive the compressed bootstrap. For :mod:`mitogen.ssh` this must have
#: length of at least `max(len('password'), len('debug1:'))`
EC0_MARKER = b('MITO000')
EC1_MARKER = b('MITO001')
EC2_MARKER = b('MITO002')
def __init__(self, broker):
super(BootstrapProtocol, self).__init__()
self._writer = mitogen.core.BufferedWriter(broker, self)
def on_transmit(self, broker):
self._writer.on_transmit(broker)
def _on_ec0_received(self, line, match):
LOG.debug('%r: first stage started succcessfully', self)
self._writer.write(self.stream.conn.get_preamble())
def _on_ec1_received(self, line, match):
LOG.debug('%r: first stage received mitogen.core source', self)
def _on_ec2_received(self, line, match):
LOG.debug('%r: new child booted successfully', self)
self.stream.conn._complete_connection()
return False
def on_unrecognized_line_received(self, line):
LOG.debug('%s: stdout: %s', self.stream.name,
line.decode('utf-8', 'replace'))
PATTERNS = [
(re.compile(EC0_MARKER), _on_ec0_received),
(re.compile(EC1_MARKER), _on_ec1_received),
(re.compile(EC2_MARKER), _on_ec2_received),
]
class LogProtocol(LineLoggingProtocolMixin, mitogen.core.DelimitedProtocol):
"""
For "hybrid TTY/socketpair" mode, after connection setup a spare TTY master
FD exists that cannot be closed, and to which SSH or sudo may continue
writing log messages.
The descriptor cannot be closed since the UNIX TTY layer sends SIGHUP to
processes whose controlling TTY is the slave whose master side was closed.
LogProtocol takes over this FD and creates log messages for anything
written to it.
"""
def on_line_received(self, line):
"""
Read a line, decode it as UTF-8, and log it.
"""
super(LogProtocol, self).on_line_received(line)
LOG.info(u'%s: %s', self.stream.name, line.decode('utf-8', 'replace'))
class MitogenProtocol(mitogen.core.MitogenProtocol):
"""
Extend core.MitogenProtocol to cause SHUTDOWN to be sent to the child
during graceful shutdown.
"""
def on_shutdown(self, broker):
"""
Respond to the broker's request for the stream to shut down by sending
SHUTDOWN to the child.
"""
LOG.debug('%r: requesting child shutdown', self)
self._send(
mitogen.core.Message(
src_id=mitogen.context_id,
dst_id=self.remote_id,
handle=mitogen.core.SHUTDOWN,
)
)
class Options(object):
name = None
#: The path to the remote Python interpreter.
python_path = get_sys_executable()
#: Maximum time to wait for a connection attempt.
connect_timeout = 30.0
#: True to cause context to write verbose /tmp/mitogen.<pid>.log.
debug = False
#: True to cause context to write /tmp/mitogen.stats.<pid>.<thread>.log.
profiling = False
#: True if unidirectional routing is enabled in the new child.
unidirectional = False
#: Passed via Router wrapper methods, must eventually be passed to
#: ExternalContext.main().
max_message_size = None
#: Remote name.
remote_name = None
#: Derived from :py:attr:`connect_timeout`; absolute floating point
#: UNIX timestamp after which the connection attempt should be abandoned.
connect_deadline = None
def __init__(self, max_message_size, name=None, remote_name=None,
python_path=None, debug=False, connect_timeout=None,
profiling=False, unidirectional=False, old_router=None):
self.name = name
self.max_message_size = max_message_size
if python_path:
self.python_path = python_path
if connect_timeout:
self.connect_timeout = connect_timeout
if remote_name is None:
remote_name = get_default_remote_name()
if '/' in remote_name or '\\' in remote_name:
raise ValueError('remote_name= cannot contain slashes')
if remote_name:
self.remote_name = mitogen.core.to_text(remote_name)
self.debug = debug
self.profiling = profiling
self.unidirectional = unidirectional
self.max_message_size = max_message_size
self.connect_deadline = mitogen.core.now() + self.connect_timeout
class Connection(object):
"""
Manage the lifetime of a set of :class:`Streams <Stream>` connecting to a
remote Python interpreter, including bootstrap, disconnection, and external
tool integration.
Base for streams capable of starting children.
"""
options_class = Options
#: The protocol attached to stdio of the child.
stream_protocol_class = BootstrapProtocol
#: The protocol attached to stderr of the child.
diag_protocol_class = LogProtocol
#: :class:`Process`
proc = None
#: :class:`mitogen.core.Stream` with sides connected to stdin/stdout.
stdio_stream = None
#: If `proc.stderr` is set, referencing either a plain pipe or the
#: controlling TTY, this references the corresponding
#: :class:`LogProtocol`'s stream, allowing it to be disconnected when this
#: stream is disconnected.
stderr_stream = None
#: Function with the semantics of :func:`create_child` used to create the
#: child process.
create_child = staticmethod(create_child)
#: Dictionary of extra kwargs passed to :attr:`create_child`.
create_child_args = {}
#: :data:`True` if the remote has indicated that it intends to detach, and
#: should not be killed on disconnect.
detached = False
#: If :data:`True`, indicates the child should not be killed during
#: graceful detachment, as it the actual process implementing the child
#: context. In all other cases, the subprocess is SSH, sudo, or a similar
#: tool that should be reminded to quit during disconnection.
child_is_immediate_subprocess = True
#: Prefix given to default names generated by :meth:`connect`.
name_prefix = u'local'
#: :class:`Timer` that runs :meth:`_on_timer_expired` when connection
#: timeout occurs.
_timer = None
#: When disconnection completes, instance of :class:`Reaper` used to wait
#: on the exit status of the subprocess.
_reaper = None
#: On failure, the exception object that should be propagated back to the
#: user.
exception = None
#: Extra text appended to :class:`EofError` if that exception is raised on
#: a failed connection attempt. May be used in subclasses to hint at common
#: problems with a particular connection method.
eof_error_hint = None
def __init__(self, options, router):
#: :class:`Options`
self.options = options
self._router = router
def __repr__(self):
return 'Connection(%r)' % (self.stdio_stream,)
# Minimised, gzipped, base64'd and passed to 'python -c'. It forks, dups
# file descriptor 0 as 100, creates a pipe, then execs a new interpreter
# with a custom argv.
# * Optimized for minimum byte count after minification & compression.
# * 'CONTEXT_NAME' and 'PREAMBLE_COMPRESSED_LEN' are substituted with
# their respective values.
# * CONTEXT_NAME must be prefixed with the name of the Python binary in
# order to allow virtualenvs to detect their install prefix.
# * For Darwin, OS X installs a craptacular argv0-introspecting Python
# version switcher as /usr/bin/python. Override attempts to call it
# with an explicit call to python2.7
#
# Locals:
# R: read side of interpreter stdin.
# W: write side of interpreter stdin.
# r: read side of core_src FD.
# w: write side of core_src FD.
# C: the decompressed core source.
# Final os.close(2) to avoid --py-debug build from corrupting stream with
# "[1234 refs]" during exit.
@staticmethod
def _first_stage():
R,W=os.pipe()
r,w=os.pipe()
if os.fork():
os.dup2(0,100)
os.dup2(R,0)
os.dup2(r,101)
os.close(R)
os.close(r)
os.close(W)
os.close(w)
if sys.platform == 'darwin' and sys.executable == '/usr/bin/python':
sys.executable += sys.version[:3]
os.environ['ARGV0']=sys.executable
os.execl(sys.executable,sys.executable+'(mitogen:CONTEXT_NAME)')
os.write(1,'MITO000\n'.encode())
C=_(os.fdopen(0,'rb').read(PREAMBLE_COMPRESSED_LEN),'zip')
fp=os.fdopen(W,'wb',0)
fp.write(C)
fp.close()
fp=os.fdopen(w,'wb',0)
fp.write(C)
fp.close()
os.write(1,'MITO001\n'.encode())
os.close(2)
def get_python_argv(self):
"""
Return the initial argument vector elements necessary to invoke Python,
by returning a 1-element list containing :attr:`python_path` if it is a
string, or simply returning it if it is already a list.
This allows emulation of existing tools where the Python invocation may
be set to e.g. `['/usr/bin/env', 'python']`.
"""
if isinstance(self.options.python_path, list):
return self.options.python_path
return [self.options.python_path]
def get_boot_command(self):
source = inspect.getsource(self._first_stage)
source = textwrap.dedent('\n'.join(source.strip().split('\n')[2:]))
source = source.replace(' ', '\t')
source = source.replace('CONTEXT_NAME', self.options.remote_name)
preamble_compressed = self.get_preamble()
source = source.replace('PREAMBLE_COMPRESSED_LEN',
str(len(preamble_compressed)))
compressed = zlib.compress(source.encode(), 9)
encoded = codecs.encode(compressed, 'base64').replace(b('\n'), b(''))
# We can't use bytes.decode() in 3.x since it was restricted to always
# return unicode, so codecs.decode() is used instead. In 3.x
# codecs.decode() requires a bytes object. Since we must be compatible
# with 2.4 (no bytes literal), an extra .encode() either returns the
# same str (2.x) or an equivalent bytes (3.x).
return self.get_python_argv() + [
'-c',
'import codecs,os,sys;_=codecs.decode;'
'exec(_(_("%s".encode(),"base64"),"zip"))' % (encoded.decode(),)
]
def get_econtext_config(self):
assert self.options.max_message_size is not None
parent_ids = mitogen.parent_ids[:]
parent_ids.insert(0, mitogen.context_id)
return {
'parent_ids': parent_ids,
'context_id': self.context.context_id,
'debug': self.options.debug,
'profiling': self.options.profiling,
'unidirectional': self.options.unidirectional,
'log_level': get_log_level(),
'whitelist': self._router.get_module_whitelist(),
'blacklist': self._router.get_module_blacklist(),
'max_message_size': self.options.max_message_size,
'version': mitogen.__version__,
}
def get_preamble(self):
suffix = (
'\nExternalContext(%r).main()\n' %\
(self.get_econtext_config(),)
)
partial = get_core_source_partial()
return partial.append(suffix.encode('utf-8'))
def _get_name(self):
"""
Called by :meth:`connect` after :attr:`pid` is known. Subclasses can
override it to specify a default stream name, or set
:attr:`name_prefix` to generate a default format.
"""
return u'%s.%s' % (self.name_prefix, self.proc.pid)
def start_child(self):
args = self.get_boot_command()
LOG.debug('command line for %r: %s', self, Argv(args))
try:
return self.create_child(args=args, **self.create_child_args)
except OSError:
e = sys.exc_info()[1]
msg = 'Child start failed: %s. Command was: %s' % (e, Argv(args))
raise mitogen.core.StreamError(msg)
def _adorn_eof_error(self, e):
"""
Subclasses may provide additional information in the case of a failed
connection.
"""
if self.eof_error_hint:
e.args = ('%s\n\n%s' % (e.args[0], self.eof_error_hint),)
def _complete_connection(self):
self._timer.cancel()
if not self.exception:
mitogen.core.unlisten(self._router.broker, 'shutdown',
self._on_broker_shutdown)
self._router.register(self.context, self.stdio_stream)
self.stdio_stream.set_protocol(
MitogenProtocol(
router=self._router,
remote_id=self.context.context_id,
)
)
self._router.route_monitor.notice_stream(self.stdio_stream)
self.latch.put()
def _fail_connection(self, exc):
"""
Fail the connection attempt.
"""
LOG.debug('failing connection %s due to %r',
self.stdio_stream and self.stdio_stream.name, exc)
if self.exception is None:
self._adorn_eof_error(exc)
self.exception = exc
mitogen.core.unlisten(self._router.broker, 'shutdown',
self._on_broker_shutdown)
for stream in self.stdio_stream, self.stderr_stream:
if stream and not stream.receive_side.closed:
stream.on_disconnect(self._router.broker)
self._complete_connection()
eof_error_msg = 'EOF on stream; last 100 lines received:\n'
def on_stdio_disconnect(self):
"""
Handle stdio stream disconnection by failing the Connection if the
stderr stream has already been closed. Otherwise, wait for it to close
(or timeout), to allow buffered diagnostic logs to be consumed.
It is normal that when a subprocess aborts, stdio has nothing buffered
when it is closed, thus signalling readability, causing an empty read
(interpreted as indicating disconnection) on the next loop iteration,
even if its stderr pipe has lots of diagnostic logs still buffered in
the kernel. Therefore we must wait for both pipes to indicate they are
empty before triggering connection failure.
"""
stderr = self.stderr_stream
if stderr is None or stderr.receive_side.closed:
self._on_streams_disconnected()
def on_stderr_disconnect(self):
"""
Inverse of :func:`on_stdio_disconnect`.
"""
if self.stdio_stream.receive_side.closed:
self._on_streams_disconnected()
def _on_streams_disconnected(self):
"""
When disconnection has been detected for both streams, cancel the
connection timer, mark the connection failed, and reap the child
process. Do nothing if the timer has already been cancelled, indicating
some existing failure has already been noticed.
"""
if self._timer.active:
self._timer.cancel()
self._fail_connection(EofError(
self.eof_error_msg + get_history(
[self.stdio_stream, self.stderr_stream]
)
))
if self._reaper:
return
self._reaper = Reaper(
broker=self._router.broker,
proc=self.proc,
kill=not (
(self.detached and self.child_is_immediate_subprocess) or
# Avoid killing so child has chance to write cProfile data
self._router.profiling
),
# Don't delay shutdown waiting for a detached child, since the
# detached child may expect to live indefinitely after its parent
# exited.
wait_on_shutdown=(not self.detached),
)
self._reaper.reap()
def _on_broker_shutdown(self):
"""
Respond to broker.shutdown() being called by failing the connection
attempt.
"""
self._fail_connection(CancelledError(BROKER_SHUTDOWN_MSG))
def stream_factory(self):
return self.stream_protocol_class.build_stream(
broker=self._router.broker,
)
def stderr_stream_factory(self):
return self.diag_protocol_class.build_stream()
def _setup_stdio_stream(self):
stream = self.stream_factory()
stream.conn = self
stream.name = self.options.name or self._get_name()
stream.accept(self.proc.stdout, self.proc.stdin)
mitogen.core.listen(stream, 'disconnect', self.on_stdio_disconnect)
self._router.broker.start_receive(stream)
return stream
def _setup_stderr_stream(self):
stream = self.stderr_stream_factory()
stream.conn = self
stream.name = self.options.name or self._get_name()
stream.accept(self.proc.stderr, self.proc.stderr)
mitogen.core.listen(stream, 'disconnect', self.on_stderr_disconnect)
self._router.broker.start_receive(stream)
return stream
def _on_timer_expired(self):
self._fail_connection(
mitogen.core.TimeoutError(
'Failed to setup connection after %.2f seconds',
self.options.connect_timeout,
)
)
def _async_connect(self):
LOG.debug('creating connection to context %d using %s',
self.context.context_id, self.__class__.__module__)
mitogen.core.listen(self._router.broker, 'shutdown',
self._on_broker_shutdown)
self._timer = self._router.broker.timers.schedule(
when=self.options.connect_deadline,
func=self._on_timer_expired,
)
try:
self.proc = self.start_child()
except Exception:
LOG.debug('failed to start child', exc_info=True)
self._fail_connection(sys.exc_info()[1])
return
LOG.debug('child for %r started: pid:%r stdin:%r stdout:%r stderr:%r',
self, self.proc.pid,
self.proc.stdin.fileno(),
self.proc.stdout.fileno(),
self.proc.stderr and self.proc.stderr.fileno())
self.stdio_stream = self._setup_stdio_stream()
if self.context.name is None:
self.context.name = self.stdio_stream.name
self.proc.name = self.stdio_stream.name
if self.proc.stderr:
self.stderr_stream = self._setup_stderr_stream()
def connect(self, context):
self.context = context
self.latch = mitogen.core.Latch()
self._router.broker.defer(self._async_connect)
self.latch.get()
if self.exception:
raise self.exception
class ChildIdAllocator(object):
"""
Allocate new context IDs from a block of unique context IDs allocated by
the master process.
"""
def __init__(self, router):
self.router = router
self.lock = threading.Lock()
self.it = iter(xrange(0))
def allocate(self):
"""
Allocate an ID, requesting a fresh block from the master if the
existing block is exhausted.
:returns:
The new context ID.
.. warning::
This method is not safe to call from the :class:`Broker` thread, as
it may block on IO of its own.
"""
self.lock.acquire()
try:
for id_ in self.it:
return id_
master = self.router.context_by_id(0)
start, end = master.send_await(
mitogen.core.Message(dst_id=0, handle=mitogen.core.ALLOCATE_ID)
)
self.it = iter(xrange(start, end))
finally:
self.lock.release()
return self.allocate()
class CallChain(object):
"""
Deliver :data:`mitogen.core.CALL_FUNCTION` messages to a target context,
optionally threading related calls so an exception in an earlier call
cancels subsequent calls.
:param mitogen.core.Context context:
Target context.
:param bool pipelined:
Enable pipelining.
:meth:`call`, :meth:`call_no_reply` and :meth:`call_async`
normally issue calls and produce responses with no memory of prior
exceptions. If a call made with :meth:`call_no_reply` fails, the exception
is logged to the target context's logging framework.
**Pipelining**
When pipelining is enabled, if an exception occurs during a call,
subsequent calls made by the same :class:`CallChain` fail with the same
exception, including those already in-flight on the network, and no further
calls execute until :meth:`reset` is invoked.
No exception is logged for calls made with :meth:`call_no_reply`, instead
the exception is saved and reported as the result of subsequent
:meth:`call` or :meth:`call_async` calls.
Sequences of asynchronous calls can be made without wasting network
round-trips to discover if prior calls succeed, and chains originating from
multiple unrelated source contexts may overlap concurrently at a target
context without interference.
In this example, 4 calls complete in one round-trip::
chain = mitogen.parent.CallChain(context, pipelined=True)
chain.call_no_reply(os.mkdir, '/tmp/foo')
# If previous mkdir() failed, this never runs:
chain.call_no_reply(os.mkdir, '/tmp/foo/bar')
# If either mkdir() failed, this never runs, and the exception is
# asynchronously delivered to the receiver.
recv = chain.call_async(subprocess.check_output, '/tmp/foo')
# If anything so far failed, this never runs, and raises the exception.
chain.call(do_something)
# If this code was executed, the exception would also be raised.
if recv.get().unpickle() == 'baz':
pass
When pipelining is enabled, :meth:`reset` must be invoked to ensure any
exception is discarded, otherwise unbounded memory usage is possible in
long-running programs. The context manager protocol is supported to ensure
:meth:`reset` is always invoked::
with mitogen.parent.CallChain(context, pipelined=True) as chain:
chain.call_no_reply(...)
chain.call_no_reply(...)
chain.call_no_reply(...)
chain.call(...)
# chain.reset() automatically invoked.
"""
def __init__(self, context, pipelined=False):
self.context = context
if pipelined:
self.chain_id = self.make_chain_id()
else:
self.chain_id = None
@classmethod
def make_chain_id(cls):
return '%s-%s-%x-%x' % (
socket.gethostname(),
os.getpid(),
thread.get_ident(),
int(1e6 * mitogen.core.now()),
)
def __repr__(self):
return '%s(%s)' % (self.__class__.__name__, self.context)
def __enter__(self):
return self
def __exit__(self, _1, _2, _3):
self.reset()
def reset(self):
"""
Instruct the target to forget any related exception.
"""
if not self.chain_id:
return
saved, self.chain_id = self.chain_id, None
try:
self.call_no_reply(mitogen.core.Dispatcher.forget_chain, saved)
finally:
self.chain_id = saved
closures_msg = (
'Mitogen cannot invoke closures, as doing so would require '
'serializing arbitrary program state, and no universal '
'method exists to recover a reference to them.'
)
lambda_msg = (
'Mitogen cannot invoke anonymous functions, as no universal method '
'exists to recover a reference to an anonymous function.'
)
method_msg = (
'Mitogen cannot invoke instance methods, as doing so would require '
'serializing arbitrary program state.'
)
def make_msg(self, fn, *args, **kwargs):
if getattr(fn, closure_attr, None) is not None:
raise TypeError(self.closures_msg)
if fn.__name__ == '<lambda>':
raise TypeError(self.lambda_msg)
if inspect.ismethod(fn):
im_self = getattr(fn, IM_SELF_ATTR)
if not inspect.isclass(im_self):
raise TypeError(self.method_msg)
klass = mitogen.core.to_text(im_self.__name__)
else:
klass = None
tup = (
self.chain_id,
mitogen.core.to_text(fn.__module__),
klass,
mitogen.core.to_text(fn.__name__),
args,
mitogen.core.Kwargs(kwargs)
)
return mitogen.core.Message.pickled(tup,
handle=mitogen.core.CALL_FUNCTION)
def call_no_reply(self, fn, *args, **kwargs):
"""
Like :meth:`call_async`, but do not wait for a return value, and inform
the target context no reply is expected. If the call fails and
pipelining is disabled, the exception will be logged to the target
context's logging framework.
"""
LOG.debug('starting no-reply function call to %r: %r',
self.context.name or self.context.context_id,
CallSpec(fn, args, kwargs))
self.context.send(self.make_msg(fn, *args, **kwargs))
def call_async(self, fn, *args, **kwargs):
"""
Arrange for `fn(*args, **kwargs)` to be invoked on the context's main
thread.
:param fn:
A free function in module scope or a class method of a class
directly reachable from module scope:
.. code-block:: python
# mymodule.py
def my_func():
'''A free function reachable as mymodule.my_func'''
class MyClass:
@classmethod
def my_classmethod(cls):
'''Reachable as mymodule.MyClass.my_classmethod'''
def my_instancemethod(self):
'''Unreachable: requires a class instance!'''
class MyEmbeddedClass:
@classmethod
def my_classmethod(cls):
'''Not directly reachable from module scope!'''
:param tuple args:
Function arguments, if any. See :ref:`serialization-rules` for
permitted types.
:param dict kwargs:
Function keyword arguments, if any. See :ref:`serialization-rules`
for permitted types.
:returns:
:class:`mitogen.core.Receiver` configured to receive the result of
the invocation:
.. code-block:: python
recv = context.call_async(os.check_output, 'ls /tmp/')
try:
# Prints output once it is received.
msg = recv.get()
print(msg.unpickle())
except mitogen.core.CallError, e:
print('Call failed:', str(e))
Asynchronous calls may be dispatched in parallel to multiple
contexts and consumed as they complete using
:class:`mitogen.select.Select`.
"""
LOG.debug('starting function call to %s: %r',
self.context.name or self.context.context_id,
CallSpec(fn, args, kwargs))
return self.context.send_async(self.make_msg(fn, *args, **kwargs))
def call(self, fn, *args, **kwargs):
"""
Like :meth:`call_async`, but block until the return value is available.
Equivalent to::
call_async(fn, *args, **kwargs).get().unpickle()
:returns:
The function's return value.
:raises mitogen.core.CallError:
An exception was raised in the remote context during execution.
"""
receiver = self.call_async(fn, *args, **kwargs)
return receiver.get().unpickle(throw_dead=False)
class Context(mitogen.core.Context):
"""
Extend :class:`mitogen.core.Context` with functionality useful to masters,
and child contexts who later become parents. Currently when this class is
required, the target context's router is upgraded at runtime.
"""
#: A :class:`CallChain` instance constructed by default, with pipelining
#: disabled. :meth:`call`, :meth:`call_async` and :meth:`call_no_reply` use
#: this instance.
call_chain_class = CallChain
via = None
def __init__(self, *args, **kwargs):
super(Context, self).__init__(*args, **kwargs)
self.default_call_chain = self.call_chain_class(self)
def __ne__(self, other):
return not (self == other)
def __eq__(self, other):
return (
isinstance(other, mitogen.core.Context) and
(other.context_id == self.context_id) and
(other.router == self.router)
)
def __hash__(self):
return hash((self.router, self.context_id))
def call_async(self, fn, *args, **kwargs):
"""
See :meth:`CallChain.call_async`.
"""
return self.default_call_chain.call_async(fn, *args, **kwargs)
def call(self, fn, *args, **kwargs):
"""
See :meth:`CallChain.call`.
"""
return self.default_call_chain.call(fn, *args, **kwargs)
def call_no_reply(self, fn, *args, **kwargs):
"""
See :meth:`CallChain.call_no_reply`.
"""
self.default_call_chain.call_no_reply(fn, *args, **kwargs)
def shutdown(self, wait=False):
"""
Arrange for the context to receive a ``SHUTDOWN`` message, triggering
graceful shutdown.
Due to a lack of support for timers, no attempt is made yet to force
terminate a hung context using this method. This will be fixed shortly.
:param bool wait:
If :data:`True`, block the calling thread until the context has
completely terminated.
:returns:
If `wait` is :data:`False`, returns a :class:`mitogen.core.Latch`
whose :meth:`get() <mitogen.core.Latch.get>` method returns
:data:`None` when shutdown completes. The `timeout` parameter may
be used to implement graceful timeouts.
"""
LOG.debug('%r.shutdown() sending SHUTDOWN', self)
latch = mitogen.core.Latch()
mitogen.core.listen(self, 'disconnect', lambda: latch.put(None))
self.send(
mitogen.core.Message(
handle=mitogen.core.SHUTDOWN,
)
)
if wait:
latch.get()
else:
return latch
class RouteMonitor(object):
"""
Generate and respond to :data:`mitogen.core.ADD_ROUTE` and
:data:`mitogen.core.DEL_ROUTE` messages sent to the local context by
maintaining a table of available routes, and propagating messages towards
parents and siblings as appropriate.
:class:`RouteMonitor` is responsible for generating routing messages for
directly attached children. It learns of new children via
:meth:`notice_stream` called by :class:`Router`, and subscribes to their
``disconnect`` event to learn when they disappear.
In children, constructing this class overwrites the stub
:data:`mitogen.core.DEL_ROUTE` handler installed by
:class:`mitogen.core.ExternalContext`, which is expected behaviour when a
child is beging upgraded in preparation to become a parent of children of
its own.
By virtue of only being active while responding to messages from a handler,
RouteMonitor lives entirely on the broker thread, so its data requires no
locking.
:param mitogen.master.Router router:
Router to install handlers on.
:param mitogen.core.Context parent:
:data:`None` in the master process, or reference to the parent context
we should propagate route updates towards.
"""
def __init__(self, router, parent=None):
self.router = router
self.parent = parent
self._log = logging.getLogger('mitogen.route_monitor')
#: Mapping of Stream instance to integer context IDs reachable via the
#: stream; used to cleanup routes during disconnection.
self._routes_by_stream = {}
self.router.add_handler(
fn=self._on_add_route,
handle=mitogen.core.ADD_ROUTE,
persist=True,
policy=is_immediate_child,
overwrite=True,
)
self.router.add_handler(
fn=self._on_del_route,
handle=mitogen.core.DEL_ROUTE,
persist=True,
policy=is_immediate_child,
overwrite=True,
)
def __repr__(self):
return 'RouteMonitor()'
def _send_one(self, stream, handle, target_id, name):
"""
Compose and send an update message on a stream.
:param mitogen.core.Stream stream:
Stream to send it on.
:param int handle:
:data:`mitogen.core.ADD_ROUTE` or :data:`mitogen.core.DEL_ROUTE`
:param int target_id:
ID of the connecting or disconnecting context.
:param str name:
Context name or :data:`None`.
"""
if not stream:
# We may not have a stream during shutdown.
return
data = str(target_id)
if name:
data = '%s:%s' % (target_id, name)
stream.protocol.send(
mitogen.core.Message(
handle=handle,
data=data.encode('utf-8'),
dst_id=stream.protocol.remote_id,
)
)
def _propagate_up(self, handle, target_id, name=None):
"""
In a non-master context, propagate an update towards the master.
:param int handle:
:data:`mitogen.core.ADD_ROUTE` or :data:`mitogen.core.DEL_ROUTE`
:param int target_id:
ID of the connecting or disconnecting context.
:param str name:
For :data:`mitogen.core.ADD_ROUTE`, the name of the new context
assigned by its parent. This is used by parents to assign the
:attr:`mitogen.core.Context.name` attribute.
"""
if self.parent:
stream = self.router.stream_by_id(self.parent.context_id)
self._send_one(stream, handle, target_id, name)
def _propagate_down(self, handle, target_id):
"""
For DEL_ROUTE, we additionally want to broadcast the message to any
stream that has ever communicated with the disconnecting ID, so
core.py's :meth:`mitogen.core.Router._on_del_route` can turn the
message into a disconnect event.
:param int handle:
:data:`mitogen.core.ADD_ROUTE` or :data:`mitogen.core.DEL_ROUTE`
:param int target_id:
ID of the connecting or disconnecting context.
"""
for stream in self.router.get_streams():
if target_id in stream.protocol.egress_ids and (
(self.parent is None) or
(self.parent.context_id != stream.protocol.remote_id)
):
self._send_one(stream, mitogen.core.DEL_ROUTE, target_id, None)
def notice_stream(self, stream):
"""
When this parent is responsible for a new directly connected child
stream, we're also responsible for broadcasting
:data:`mitogen.core.DEL_ROUTE` upstream when that child disconnects.
"""
self._routes_by_stream[stream] = set([stream.protocol.remote_id])
self._propagate_up(mitogen.core.ADD_ROUTE, stream.protocol.remote_id,
stream.name)
mitogen.core.listen(
obj=stream,
name='disconnect',
func=lambda: self._on_stream_disconnect(stream),
)
def get_routes(self, stream):
"""
Return the set of context IDs reachable on a stream.
:param mitogen.core.Stream stream:
:returns: set([int])
"""
return self._routes_by_stream.get(stream) or set()
def _on_stream_disconnect(self, stream):
"""
Respond to disconnection of a local stream by propagating DEL_ROUTE for
any contexts we know were attached to it.
"""
# During a stream crash it is possible for disconnect signal to fire
# twice, in which case ignore the second instance.
routes = self._routes_by_stream.pop(stream, None)
if routes is None:
return
self._log.debug('stream %s is gone; propagating DEL_ROUTE for %r',
stream.name, routes)
for target_id in routes:
self.router.del_route(target_id)
self._propagate_up(mitogen.core.DEL_ROUTE, target_id)
self._propagate_down(mitogen.core.DEL_ROUTE, target_id)
context = self.router.context_by_id(target_id, create=False)
if context:
mitogen.core.fire(context, 'disconnect')
def _on_add_route(self, msg):
"""
Respond to :data:`mitogen.core.ADD_ROUTE` by validating the source of
the message, updating the local table, and propagating the message
upwards.
"""
if msg.is_dead:
return
target_id_s, _, target_name = bytes_partition(msg.data, b(':'))
target_name = target_name.decode()
target_id = int(target_id_s)
self.router.context_by_id(target_id).name = target_name
stream = self.router.stream_by_id(msg.src_id)
current = self.router.stream_by_id(target_id)
if current and current.protocol.remote_id != mitogen.parent_id:
self._log.error('Cannot add duplicate route to %r via %r, '
'already have existing route via %r',
target_id, stream, current)
return
self._log.debug('Adding route to %d via %r', target_id, stream)
self._routes_by_stream[stream].add(target_id)
self.router.add_route(target_id, stream)
self._propagate_up(mitogen.core.ADD_ROUTE, target_id, target_name)
def _on_del_route(self, msg):
"""
Respond to :data:`mitogen.core.DEL_ROUTE` by validating the source of
the message, updating the local table, propagating the message
upwards, and downwards towards any stream that every had a message
forwarded from it towards the disconnecting context.
"""
if msg.is_dead:
return
target_id = int(msg.data)
registered_stream = self.router.stream_by_id(target_id)
if registered_stream is None:
return
stream = self.router.stream_by_id(msg.src_id)
if registered_stream != stream:
self._log.error('received DEL_ROUTE for %d from %r, expected %r',
target_id, stream, registered_stream)
return
context = self.router.context_by_id(target_id, create=False)
if context:
self._log.debug('firing local disconnect signal for %r', context)
mitogen.core.fire(context, 'disconnect')
self._log.debug('deleting route to %d via %r', target_id, stream)
routes = self._routes_by_stream.get(stream)
if routes:
routes.discard(target_id)
self.router.del_route(target_id)
if stream.protocol.remote_id != mitogen.parent_id:
self._propagate_up(mitogen.core.DEL_ROUTE, target_id)
self._propagate_down(mitogen.core.DEL_ROUTE, target_id)
class Router(mitogen.core.Router):
context_class = Context
debug = False
profiling = False
id_allocator = None
responder = None
log_forwarder = None
route_monitor = None
def upgrade(self, importer, parent):
LOG.debug('upgrading %r with capabilities to start new children', self)
self.id_allocator = ChildIdAllocator(router=self)
self.responder = ModuleForwarder(
router=self,
parent_context=parent,
importer=importer,
)
self.route_monitor = RouteMonitor(self, parent)
self.add_handler(
fn=self._on_detaching,
handle=mitogen.core.DETACHING,
persist=True,
)
def _on_detaching(self, msg):
if msg.is_dead:
return
stream = self.stream_by_id(msg.src_id)
if stream.protocol.remote_id != msg.src_id or stream.conn.detached:
LOG.warning('bad DETACHING received on %r: %r', stream, msg)
return
LOG.debug('%r: marking as detached', stream)
stream.conn.detached = True
msg.reply(None)
def get_streams(self):
"""
Return an atomic snapshot of all streams in existence at time of call.
This is safe to call from any thread.
"""
self._write_lock.acquire()
try:
return itervalues(self._stream_by_id)
finally:
self._write_lock.release()
def disconnect(self, context):
"""
Disconnect a context and forget its stream, assuming the context is
directly connected.
"""
stream = self.stream_by_id(context)
if stream is None or stream.protocol.remote_id != context.context_id:
return
l = mitogen.core.Latch()
mitogen.core.listen(stream, 'disconnect', l.put)
def disconnect():
LOG.debug('Starting disconnect of %r', stream)
stream.on_disconnect(self.broker)
self.broker.defer(disconnect)
l.get()
def add_route(self, target_id, stream):
"""
Arrange for messages whose `dst_id` is `target_id` to be forwarded on a
directly connected :class:`Stream`. Safe to call from any thread.
This is called automatically by :class:`RouteMonitor` in response to
:data:`mitogen.core.ADD_ROUTE` messages, but remains public while the
design has not yet settled, and situations may arise where routing is
not fully automatic.
:param int target_id:
Target context ID to add a route for.
:param mitogen.core.Stream stream:
Stream over which messages to the target should be routed.
"""
LOG.debug('%r: adding route to context %r via %r',
self, target_id, stream)
assert isinstance(target_id, int)
assert isinstance(stream, mitogen.core.Stream)
self._write_lock.acquire()
try:
self._stream_by_id[target_id] = stream
finally:
self._write_lock.release()
def del_route(self, target_id):
"""
Delete any route that exists for `target_id`. It is not an error to
delete a route that does not currently exist. Safe to call from any
thread.
This is called automatically by :class:`RouteMonitor` in response to
:data:`mitogen.core.DEL_ROUTE` messages, but remains public while the
design has not yet settled, and situations may arise where routing is
not fully automatic.
:param int target_id:
Target context ID to delete route for.
"""
LOG.debug('%r: deleting route to %r', self, target_id)
# DEL_ROUTE may be sent by a parent if it knows this context sent
# messages to a peer that has now disconnected, to let us raise
# 'disconnect' event on the appropriate Context instance. In that case,
# we won't a matching _stream_by_id entry for the disappearing route,
# so don't raise an error for a missing key here.
self._write_lock.acquire()
try:
self._stream_by_id.pop(target_id, None)
finally:
self._write_lock.release()
def get_module_blacklist(self):
if mitogen.context_id == 0:
return self.responder.blacklist
return self.importer.master_blacklist
def get_module_whitelist(self):
if mitogen.context_id == 0:
return self.responder.whitelist
return self.importer.master_whitelist
def allocate_id(self):
return self.id_allocator.allocate()
connection_timeout_msg = u"Connection timed out."
def _connect(self, klass, **kwargs):
context_id = self.allocate_id()
context = self.context_class(self, context_id)
context.name = kwargs.get('name')
kwargs['old_router'] = self
kwargs['max_message_size'] = self.max_message_size
conn = klass(klass.options_class(**kwargs), self)
try:
conn.connect(context=context)
except mitogen.core.TimeoutError:
raise mitogen.core.StreamError(self.connection_timeout_msg)
return context
def connect(self, method_name, name=None, **kwargs):
if name:
name = mitogen.core.to_text(name)
klass = get_connection_class(method_name)
kwargs.setdefault(u'debug', self.debug)
kwargs.setdefault(u'profiling', self.profiling)
kwargs.setdefault(u'unidirectional', self.unidirectional)
kwargs.setdefault(u'name', name)
via = kwargs.pop(u'via', None)
if via is not None:
return self.proxy_connect(via, method_name,
**mitogen.core.Kwargs(kwargs))
return self._connect(klass, **mitogen.core.Kwargs(kwargs))
def proxy_connect(self, via_context, method_name, name=None, **kwargs):
resp = via_context.call(_proxy_connect,
name=name,
method_name=method_name,
kwargs=mitogen.core.Kwargs(kwargs),
)
if resp['msg'] is not None:
raise mitogen.core.StreamError(resp['msg'])
name = u'%s.%s' % (via_context.name, resp['name'])
context = self.context_class(self, resp['id'], name=name)
context.via = via_context
self._write_lock.acquire()
try:
self._context_by_id[context.context_id] = context
finally:
self._write_lock.release()
return context
def buildah(self, **kwargs):
return self.connect(u'buildah', **kwargs)
def doas(self, **kwargs):
return self.connect(u'doas', **kwargs)
def docker(self, **kwargs):
return self.connect(u'docker', **kwargs)
def kubectl(self, **kwargs):
return self.connect(u'kubectl', **kwargs)
def fork(self, **kwargs):
return self.connect(u'fork', **kwargs)
def jail(self, **kwargs):
return self.connect(u'jail', **kwargs)
def local(self, **kwargs):
return self.connect(u'local', **kwargs)
def lxc(self, **kwargs):
return self.connect(u'lxc', **kwargs)
def lxd(self, **kwargs):
return self.connect(u'lxd', **kwargs)
def setns(self, **kwargs):
return self.connect(u'setns', **kwargs)
def su(self, **kwargs):
return self.connect(u'su', **kwargs)
def sudo(self, **kwargs):
return self.connect(u'sudo', **kwargs)
def ssh(self, **kwargs):
return self.connect(u'ssh', **kwargs)
class Reaper(object):
"""
Asynchronous logic for reaping :class:`Process` objects. This is necessary
to prevent uncontrolled buildup of zombie processes in long-lived parents
that will eventually reach an OS limit, preventing creation of new threads
and processes, and to log the exit status of the child in the case of an
error.
To avoid modifying process-global state such as with
:func:`signal.set_wakeup_fd` or installing a :data:`signal.SIGCHLD` handler
that might interfere with the user's ability to use those facilities,
Reaper polls for exit with backoff using timers installed on an associated
:class:`Broker`.
:param mitogen.core.Broker broker:
The :class:`Broker` on which to install timers
:param mitogen.parent.Process proc:
The process to reap.
:param bool kill:
If :data:`True`, send ``SIGTERM`` and ``SIGKILL`` to the process.
:param bool wait_on_shutdown:
If :data:`True`, delay :class:`Broker` shutdown if child has not yet
exited. If :data:`False` simply forget the child.
"""
#: :class:`Timer` that invokes :meth:`reap` after some polling delay.
_timer = None
def __init__(self, broker, proc, kill, wait_on_shutdown):
self.broker = broker
self.proc = proc
self.kill = kill
self.wait_on_shutdown = wait_on_shutdown
self._tries = 0
def _signal_child(self, signum):
# For processes like sudo we cannot actually send sudo a signal,
# because it is setuid, so this is best-effort only.
LOG.debug('%r: sending %s', self.proc, SIGNAL_BY_NUM[signum])
try:
os.kill(self.proc.pid, signum)
except OSError:
e = sys.exc_info()[1]
if e.args[0] != errno.EPERM:
raise
def _calc_delay(self, count):
"""
Calculate a poll delay given `count` attempts have already been made.
These constants have no principle, they just produce rapid but still
relatively conservative retries.
"""
delay = 0.05
for _ in xrange(count):
delay *= 1.72
return delay
def _on_broker_shutdown(self):
"""
Respond to :class:`Broker` shutdown by cancelling the reap timer if
:attr:`Router.await_children_at_shutdown` is disabled. Otherwise
shutdown is delayed for up to :attr:`Broker.shutdown_timeout` for
subprocesses may have no intention of exiting any time soon.
"""
if not self.wait_on_shutdown:
self._timer.cancel()
def _install_timer(self, delay):
new = self._timer is None
self._timer = self.broker.timers.schedule(
when=mitogen.core.now() + delay,
func=self.reap,
)
if new:
mitogen.core.listen(self.broker, 'shutdown',
self._on_broker_shutdown)
def _remove_timer(self):
if self._timer and self._timer.active:
self._timer.cancel()
mitogen.core.unlisten(self.broker, 'shutdown',
self._on_broker_shutdown)
def reap(self):
"""
Reap the child process during disconnection.
"""
status = self.proc.poll()
if status is not None:
LOG.debug('%r: %s', self.proc, returncode_to_str(status))
mitogen.core.fire(self.proc, 'exit')
self._remove_timer()
return
self._tries += 1
if self._tries > 20:
LOG.warning('%r: child will not exit, giving up', self)
self._remove_timer()
return
delay = self._calc_delay(self._tries - 1)
LOG.debug('%r still running after IO disconnect, recheck in %.03fs',
self.proc, delay)
self._install_timer(delay)
if not self.kill:
pass
elif self._tries == 2:
self._signal_child(signal.SIGTERM)
elif self._tries == 6: # roughly 4 seconds
self._signal_child(signal.SIGKILL)
class Process(object):
"""
Process objects provide a uniform interface to the :mod:`subprocess` and
:mod:`mitogen.fork`. This class is extended by :class:`PopenProcess` and
:class:`mitogen.fork.Process`.
:param int pid:
The process ID.
:param file stdin:
File object attached to standard input.
:param file stdout:
File object attached to standard output.
:param file stderr:
File object attached to standard error, or :data:`None`.
"""
#: Name of the process used in logs. Set to the stream/context name by
#: :class:`Connection`.
name = None
def __init__(self, pid, stdin, stdout, stderr=None):
#: The process ID.
self.pid = pid
#: File object attached to standard input.
self.stdin = stdin
#: File object attached to standard output.
self.stdout = stdout
#: File object attached to standard error.
self.stderr = stderr
def __repr__(self):
return '%s %s pid %d' % (
type(self).__name__,
self.name,
self.pid,
)
def poll(self):
"""
Fetch the child process exit status, or :data:`None` if it is still
running. This should be overridden by subclasses.
:returns:
Exit status in the style of the :attr:`subprocess.Popen.returncode`
attribute, i.e. with signals represented by a negative integer.
"""
raise NotImplementedError()
class PopenProcess(Process):
"""
:class:`Process` subclass wrapping a :class:`subprocess.Popen` object.
:param subprocess.Popen proc:
The subprocess.
"""
def __init__(self, proc, stdin, stdout, stderr=None):
super(PopenProcess, self).__init__(proc.pid, stdin, stdout, stderr)
#: The subprocess.
self.proc = proc
def poll(self):
return self.proc.poll()
class ModuleForwarder(object):
"""
Respond to :data:`mitogen.core.GET_MODULE` requests in a child by
forwarding the request to our parent context, or satisfying the request
from our local Importer cache.
"""
def __init__(self, router, parent_context, importer):
self.router = router
self.parent_context = parent_context
self.importer = importer
router.add_handler(
fn=self._on_forward_module,
handle=mitogen.core.FORWARD_MODULE,
persist=True,
policy=mitogen.core.has_parent_authority,
)
router.add_handler(
fn=self._on_get_module,
handle=mitogen.core.GET_MODULE,
persist=True,
policy=is_immediate_child,
)
def __repr__(self):
return 'ModuleForwarder'
def _on_forward_module(self, msg):
if msg.is_dead:
return
context_id_s, _, fullname = bytes_partition(msg.data, b('\x00'))
fullname = mitogen.core.to_text(fullname)
context_id = int(context_id_s)
stream = self.router.stream_by_id(context_id)
if stream.protocol.remote_id == mitogen.parent_id:
LOG.error('%r: dropping FORWARD_MODULE(%d, %r): no route to child',
self, context_id, fullname)
return
if fullname in stream.protocol.sent_modules:
return
LOG.debug('%r._on_forward_module() sending %r to %r via %r',
self, fullname, context_id, stream.protocol.remote_id)
self._send_module_and_related(stream, fullname)
if stream.protocol.remote_id != context_id:
stream.protocol._send(
mitogen.core.Message(
data=msg.data,
handle=mitogen.core.FORWARD_MODULE,
dst_id=stream.protocol.remote_id,
)
)
def _on_get_module(self, msg):
if msg.is_dead:
return
fullname = msg.data.decode('utf-8')
LOG.debug('%r: %s requested by context %d', self, fullname, msg.src_id)
callback = lambda: self._on_cache_callback(msg, fullname)
self.importer._request_module(fullname, callback)
def _on_cache_callback(self, msg, fullname):
stream = self.router.stream_by_id(msg.src_id)
LOG.debug('%r: sending %s to %r', self, fullname, stream)
self._send_module_and_related(stream, fullname)
def _send_module_and_related(self, stream, fullname):
tup = self.importer._cache[fullname]
for related in tup[4]:
rtup = self.importer._cache.get(related)
if rtup:
self._send_one_module(stream, rtup)
else:
LOG.debug('%r: %s not in cache (for %s)',
self, related, fullname)
self._send_one_module(stream, tup)
def _send_one_module(self, stream, tup):
if tup[0] not in stream.protocol.sent_modules:
stream.protocol.sent_modules.add(tup[0])
self.router._async_route(
mitogen.core.Message.pickled(
tup,
dst_id=stream.protocol.remote_id,
handle=mitogen.core.LOAD_MODULE,
)
)
# 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
"""
mitogen.profiler
Record and report cProfile statistics from a run. Creates one aggregated
output file, one aggregate containing only workers, and one for the
top-level process.
Usage:
mitogen.profiler record <dest_path> <tool> [args ..]
mitogen.profiler report <dest_path> [sort_mode]
mitogen.profiler stat <sort_mode> <tool> [args ..]
Mode:
record: Record a trace.
report: Report on a previously recorded trace.
stat: Record and report in a single step.
Where:
dest_path: Filesystem prefix to write .pstats files to.
sort_mode: Sorting mode; defaults to "cumulative". See:
https://docs.python.org/2/library/profile.html#pstats.Stats.sort_stats
Example:
mitogen.profiler record /tmp/mypatch ansible-playbook foo.yml
mitogen.profiler dump /tmp/mypatch-worker.pstats
"""
from __future__ import print_function
import os
import pstats
import shutil
import subprocess
import sys
import tempfile
import time
def try_merge(stats, path):
try:
stats.add(path)
return True
except Exception as e:
print('%s failed. Will retry. %s' % (path, e))
return False
def merge_stats(outpath, inpaths):
first, rest = inpaths[0], inpaths[1:]
for x in range(1):
try:
stats = pstats.Stats(first)
except EOFError:
time.sleep(0.2)
continue
print("Writing %r..." % (outpath,))
for path in rest:
#print("Merging %r into %r.." % (os.path.basename(path), outpath))
for x in range(5):
if try_merge(stats, path):
break
time.sleep(0.2)
stats.dump_stats(outpath)
def generate_stats(outpath, tmpdir):
print('Generating stats..')
all_paths = []
paths_by_ident = {}
for name in os.listdir(tmpdir):
if name.endswith('-dump.pstats'):
ident, _, pid = name.partition('-')
path = os.path.join(tmpdir, name)
all_paths.append(path)
paths_by_ident.setdefault(ident, []).append(path)
merge_stats('%s-all.pstat' % (outpath,), all_paths)
for ident, paths in paths_by_ident.items():
merge_stats('%s-%s.pstat' % (outpath, ident), paths)
def do_record(tmpdir, path, *args):
env = os.environ.copy()
fmt = '%(identity)s-%(pid)s.%(now)s-dump.%(ext)s'
env['MITOGEN_PROFILING'] = '1'
env['MITOGEN_PROFILE_FMT'] = os.path.join(tmpdir, fmt)
rc = subprocess.call(args, env=env)
generate_stats(path, tmpdir)
return rc
def do_report(tmpdir, path, sort='cumulative'):
stats = pstats.Stats(path).sort_stats(sort)
stats.print_stats(100)
def do_stat(tmpdir, sort, *args):
valid_sorts = pstats.Stats.sort_arg_dict_default
if sort not in valid_sorts:
sys.stderr.write('Invalid sort %r, must be one of %s\n' %
(sort, ', '.join(sorted(valid_sorts))))
sys.exit(1)
outfile = os.path.join(tmpdir, 'combined')
do_record(tmpdir, outfile, *args)
aggs = ('app.main', 'mitogen.broker', 'mitogen.child_main',
'mitogen.service.pool', 'Strategy', 'WorkerProcess',
'all')
for agg in aggs:
path = '%s-%s.pstat' % (outfile, agg)
if os.path.exists(path):
print()
print()
print('------ Aggregation %r ------' % (agg,))
print()
do_report(tmpdir, path, sort)
print()
def main():
if len(sys.argv) < 2 or sys.argv[1] not in ('record', 'report', 'stat'):
sys.stderr.write(__doc__.lstrip())
sys.exit(1)
func = globals()['do_' + sys.argv[1]]
tmpdir = tempfile.mkdtemp(prefix='mitogen.profiler')
try:
sys.exit(func(tmpdir, *sys.argv[2:]) or 0)
finally:
shutil.rmtree(tmpdir)
if __name__ == '__main__':
main()