-
Notifications
You must be signed in to change notification settings - Fork 197
/
utils.py
225 lines (185 loc) · 7.12 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
# Copyright 2017, David Wilson
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
import datetime
import logging
import os
import sys
import mitogen
import mitogen.core
import mitogen.master
import mitogen.parent
LOG = logging.getLogger('mitogen')
iteritems = getattr(dict, 'iteritems', dict.items)
if mitogen.core.PY3:
iteritems = dict.items
else:
iteritems = dict.iteritems
def setup_gil():
"""
Set extremely long GIL release interval to let threads naturally progress
through CPU-heavy sequences without forcing the wake of another thread that
may contend trying to run the same CPU-heavy code. For the new-style
Ansible work, this drops runtime ~33% and involuntary context switches by
>80%, essentially making threads cooperatively scheduled.
"""
try:
# Python 2.
sys.setcheckinterval(100000)
except AttributeError:
pass
try:
# Python 3.
sys.setswitchinterval(10)
except AttributeError:
pass
def disable_site_packages():
"""
Remove all entries mentioning ``site-packages`` or ``Extras`` from
:attr:sys.path. Used primarily for testing on OS X within a virtualenv,
where OS X bundles some ancient version of the :mod:`six` module.
"""
for entry in sys.path[:]:
if 'site-packages' in entry or 'Extras' in entry:
sys.path.remove(entry)
def _formatTime(record, datefmt=None):
dt = datetime.datetime.fromtimestamp(record.created)
return dt.strftime(datefmt)
def log_get_formatter():
datefmt = '%H:%M:%S'
if sys.version_info > (2, 6):
datefmt += '.%f'
fmt = '%(asctime)s %(levelname).1s %(name)s: %(message)s'
formatter = logging.Formatter(fmt, datefmt)
formatter.formatTime = _formatTime
return formatter
def log_to_file(path=None, io=False, level='INFO'):
"""
Install a new :class:`logging.Handler` writing applications logs to the
filesystem. Useful when debugging slave IO problems.
Parameters to this function may be overridden at runtime using environment
variables. See :ref:`logging-env-vars`.
:param str path:
If not :data:`None`, a filesystem path to write logs to. Otherwise,
logs are written to :data:`sys.stderr`.
:param bool io:
If :data:`True`, include extremely verbose IO logs in the output.
Useful for debugging hangs, less useful for debugging application code.
:param str level:
Name of the :mod:`logging` package constant that is the minimum level
to log at. Useful levels are ``DEBUG``, ``INFO``, ``WARNING``, and
``ERROR``.
"""
log = logging.getLogger('')
if path:
fp = open(path, 'w', 1)
mitogen.core.set_cloexec(fp.fileno())
else:
fp = sys.stderr
level = os.environ.get('MITOGEN_LOG_LEVEL', level).upper()
io = level == 'IO'
if io:
level = 'DEBUG'
logging.getLogger('mitogen.io').setLevel(level)
level = getattr(logging, level, logging.INFO)
log.setLevel(level)
# Prevent accidental duplicate log_to_file() calls from generating
# duplicate output.
for handler_ in reversed(log.handlers):
if getattr(handler_, 'is_mitogen', None):
log.handlers.remove(handler_)
handler = logging.StreamHandler(fp)
handler.is_mitogen = True
handler.formatter = log_get_formatter()
log.handlers.insert(0, handler)
def run_with_router(func, *args, **kwargs):
"""
Arrange for `func(router, *args, **kwargs)` to run with a temporary
:class:`mitogen.master.Router`, ensuring the Router and Broker are
correctly shut down during normal or exceptional return.
:returns:
`func`'s return value.
"""
broker = mitogen.master.Broker()
router = mitogen.master.Router(broker)
try:
return func(router, *args, **kwargs)
finally:
broker.shutdown()
broker.join()
def with_router(func):
"""
Decorator version of :func:`run_with_router`. Example:
.. code-block:: python
@with_router
def do_stuff(router, arg):
pass
do_stuff(blah, 123)
"""
def wrapper(*args, **kwargs):
return run_with_router(func, *args, **kwargs)
if mitogen.core.PY3:
wrapper.func_name = func.__name__
else:
wrapper.func_name = func.func_name
return wrapper
PASSTHROUGH = (
int, float, bool,
type(None),
mitogen.core.Context,
mitogen.core.CallError,
mitogen.core.Blob,
mitogen.core.Secret,
)
def cast(obj):
"""
Many tools love to subclass built-in types in order to implement useful
functionality, such as annotating the safety of a Unicode string, or adding
additional methods to a dict. However, cPickle loves to preserve those
subtypes during serialization, resulting in CallError during :meth:`call
<mitogen.parent.Context.call>` in the target when it tries to deserialize
the data.
This function walks the object graph `obj`, producing a copy with any
custom sub-types removed. The functionality is not default since the
resulting walk may be computationally expensive given a large enough graph.
See :ref:`serialization-rules` for a list of supported types.
:param obj:
Object to undecorate.
:returns:
Undecorated object.
"""
if isinstance(obj, dict):
return dict((cast(k), cast(v)) for k, v in iteritems(obj))
if isinstance(obj, (list, tuple)):
return [cast(v) for v in obj]
if isinstance(obj, PASSTHROUGH):
return obj
if isinstance(obj, mitogen.core.UnicodeType):
return mitogen.core.UnicodeType(obj)
if isinstance(obj, mitogen.core.BytesType):
return mitogen.core.BytesType(obj)
raise TypeError("Cannot serialize: %r: %r" % (type(obj), obj))