patcher.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527
  1. import imp
  2. import sys
  3. try:
  4. # Only for this purpose, it's irrelevant if `os` was already patched.
  5. # https://github.com/eventlet/eventlet/pull/661
  6. from os import register_at_fork
  7. except ImportError:
  8. register_at_fork = None
  9. import eventlet
  10. import six
  11. __all__ = ['inject', 'import_patched', 'monkey_patch', 'is_monkey_patched']
  12. __exclude = set(('__builtins__', '__file__', '__name__'))
  13. class SysModulesSaver(object):
  14. """Class that captures some subset of the current state of
  15. sys.modules. Pass in an iterator of module names to the
  16. constructor."""
  17. def __init__(self, module_names=()):
  18. self._saved = {}
  19. imp.acquire_lock()
  20. self.save(*module_names)
  21. def save(self, *module_names):
  22. """Saves the named modules to the object."""
  23. for modname in module_names:
  24. self._saved[modname] = sys.modules.get(modname, None)
  25. def restore(self):
  26. """Restores the modules that the saver knows about into
  27. sys.modules.
  28. """
  29. try:
  30. for modname, mod in six.iteritems(self._saved):
  31. if mod is not None:
  32. sys.modules[modname] = mod
  33. else:
  34. try:
  35. del sys.modules[modname]
  36. except KeyError:
  37. pass
  38. finally:
  39. imp.release_lock()
  40. def inject(module_name, new_globals, *additional_modules):
  41. """Base method for "injecting" greened modules into an imported module. It
  42. imports the module specified in *module_name*, arranging things so
  43. that the already-imported modules in *additional_modules* are used when
  44. *module_name* makes its imports.
  45. **Note:** This function does not create or change any sys.modules item, so
  46. if your greened module use code like 'sys.modules["your_module_name"]', you
  47. need to update sys.modules by yourself.
  48. *new_globals* is either None or a globals dictionary that gets populated
  49. with the contents of the *module_name* module. This is useful when creating
  50. a "green" version of some other module.
  51. *additional_modules* should be a collection of two-element tuples, of the
  52. form (<name>, <module>). If it's not specified, a default selection of
  53. name/module pairs is used, which should cover all use cases but may be
  54. slower because there are inevitably redundant or unnecessary imports.
  55. """
  56. patched_name = '__patched_module_' + module_name
  57. if patched_name in sys.modules:
  58. # returning already-patched module so as not to destroy existing
  59. # references to patched modules
  60. return sys.modules[patched_name]
  61. if not additional_modules:
  62. # supply some defaults
  63. additional_modules = (
  64. _green_os_modules() +
  65. _green_select_modules() +
  66. _green_socket_modules() +
  67. _green_thread_modules() +
  68. _green_time_modules())
  69. # _green_MySQLdb()) # enable this after a short baking-in period
  70. # after this we are gonna screw with sys.modules, so capture the
  71. # state of all the modules we're going to mess with, and lock
  72. saver = SysModulesSaver([name for name, m in additional_modules])
  73. saver.save(module_name)
  74. # Cover the target modules so that when you import the module it
  75. # sees only the patched versions
  76. for name, mod in additional_modules:
  77. sys.modules[name] = mod
  78. # Remove the old module from sys.modules and reimport it while
  79. # the specified modules are in place
  80. sys.modules.pop(module_name, None)
  81. # Also remove sub modules and reimport. Use copy the keys to list
  82. # because of the pop operations will change the content of sys.modules
  83. # within th loop
  84. for imported_module_name in list(sys.modules.keys()):
  85. if imported_module_name.startswith(module_name + '.'):
  86. sys.modules.pop(imported_module_name, None)
  87. try:
  88. module = __import__(module_name, {}, {}, module_name.split('.')[:-1])
  89. if new_globals is not None:
  90. # Update the given globals dictionary with everything from this new module
  91. for name in dir(module):
  92. if name not in __exclude:
  93. new_globals[name] = getattr(module, name)
  94. # Keep a reference to the new module to prevent it from dying
  95. sys.modules[patched_name] = module
  96. finally:
  97. saver.restore() # Put the original modules back
  98. return module
  99. def import_patched(module_name, *additional_modules, **kw_additional_modules):
  100. """Imports a module in a way that ensures that the module uses "green"
  101. versions of the standard library modules, so that everything works
  102. nonblockingly.
  103. The only required argument is the name of the module to be imported.
  104. """
  105. return inject(
  106. module_name,
  107. None,
  108. *additional_modules + tuple(kw_additional_modules.items()))
  109. def patch_function(func, *additional_modules):
  110. """Decorator that returns a version of the function that patches
  111. some modules for the duration of the function call. This is
  112. deeply gross and should only be used for functions that import
  113. network libraries within their function bodies that there is no
  114. way of getting around."""
  115. if not additional_modules:
  116. # supply some defaults
  117. additional_modules = (
  118. _green_os_modules() +
  119. _green_select_modules() +
  120. _green_socket_modules() +
  121. _green_thread_modules() +
  122. _green_time_modules())
  123. def patched(*args, **kw):
  124. saver = SysModulesSaver()
  125. for name, mod in additional_modules:
  126. saver.save(name)
  127. sys.modules[name] = mod
  128. try:
  129. return func(*args, **kw)
  130. finally:
  131. saver.restore()
  132. return patched
  133. def _original_patch_function(func, *module_names):
  134. """Kind of the contrapositive of patch_function: decorates a
  135. function such that when it's called, sys.modules is populated only
  136. with the unpatched versions of the specified modules. Unlike
  137. patch_function, only the names of the modules need be supplied,
  138. and there are no defaults. This is a gross hack; tell your kids not
  139. to import inside function bodies!"""
  140. def patched(*args, **kw):
  141. saver = SysModulesSaver(module_names)
  142. for name in module_names:
  143. sys.modules[name] = original(name)
  144. try:
  145. return func(*args, **kw)
  146. finally:
  147. saver.restore()
  148. return patched
  149. def original(modname):
  150. """ This returns an unpatched version of a module; this is useful for
  151. Eventlet itself (i.e. tpool)."""
  152. # note that it's not necessary to temporarily install unpatched
  153. # versions of all patchable modules during the import of the
  154. # module; this is because none of them import each other, except
  155. # for threading which imports thread
  156. original_name = '__original_module_' + modname
  157. if original_name in sys.modules:
  158. return sys.modules.get(original_name)
  159. # re-import the "pure" module and store it in the global _originals
  160. # dict; be sure to restore whatever module had that name already
  161. saver = SysModulesSaver((modname,))
  162. sys.modules.pop(modname, None)
  163. # some rudimentary dependency checking -- fortunately the modules
  164. # we're working on don't have many dependencies so we can just do
  165. # some special-casing here
  166. if six.PY2:
  167. deps = {'threading': 'thread', 'Queue': 'threading'}
  168. if six.PY3:
  169. deps = {'threading': '_thread', 'queue': 'threading'}
  170. if modname in deps:
  171. dependency = deps[modname]
  172. saver.save(dependency)
  173. sys.modules[dependency] = original(dependency)
  174. try:
  175. real_mod = __import__(modname, {}, {}, modname.split('.')[:-1])
  176. if modname in ('Queue', 'queue') and not hasattr(real_mod, '_threading'):
  177. # tricky hack: Queue's constructor in <2.7 imports
  178. # threading on every instantiation; therefore we wrap
  179. # it so that it always gets the original threading
  180. real_mod.Queue.__init__ = _original_patch_function(
  181. real_mod.Queue.__init__,
  182. 'threading')
  183. # save a reference to the unpatched module so it doesn't get lost
  184. sys.modules[original_name] = real_mod
  185. finally:
  186. saver.restore()
  187. return sys.modules[original_name]
  188. already_patched = {}
  189. def monkey_patch(**on):
  190. """Globally patches certain system modules to be greenthread-friendly.
  191. The keyword arguments afford some control over which modules are patched.
  192. If no keyword arguments are supplied, all possible modules are patched.
  193. If keywords are set to True, only the specified modules are patched. E.g.,
  194. ``monkey_patch(socket=True, select=True)`` patches only the select and
  195. socket modules. Most arguments patch the single module of the same name
  196. (os, time, select). The exceptions are socket, which also patches the ssl
  197. module if present; and thread, which patches thread, threading, and Queue.
  198. It's safe to call monkey_patch multiple times.
  199. """
  200. # Workaround for import cycle observed as following in monotonic
  201. # RuntimeError: no suitable implementation for this system
  202. # see https://github.com/eventlet/eventlet/issues/401#issuecomment-325015989
  203. #
  204. # Make sure the hub is completely imported before any
  205. # monkey-patching, or we risk recursion if the process of importing
  206. # the hub calls into monkey-patched modules.
  207. eventlet.hubs.get_hub()
  208. accepted_args = set(('os', 'select', 'socket',
  209. 'thread', 'time', 'psycopg', 'MySQLdb',
  210. 'builtins', 'subprocess'))
  211. # To make sure only one of them is passed here
  212. assert not ('__builtin__' in on and 'builtins' in on)
  213. try:
  214. b = on.pop('__builtin__')
  215. except KeyError:
  216. pass
  217. else:
  218. on['builtins'] = b
  219. default_on = on.pop("all", None)
  220. for k in six.iterkeys(on):
  221. if k not in accepted_args:
  222. raise TypeError("monkey_patch() got an unexpected "
  223. "keyword argument %r" % k)
  224. if default_on is None:
  225. default_on = not (True in on.values())
  226. for modname in accepted_args:
  227. if modname == 'MySQLdb':
  228. # MySQLdb is only on when explicitly patched for the moment
  229. on.setdefault(modname, False)
  230. if modname == 'builtins':
  231. on.setdefault(modname, False)
  232. on.setdefault(modname, default_on)
  233. if on['thread'] and not already_patched.get('thread'):
  234. _green_existing_locks()
  235. modules_to_patch = []
  236. for name, modules_function in [
  237. ('os', _green_os_modules),
  238. ('select', _green_select_modules),
  239. ('socket', _green_socket_modules),
  240. ('thread', _green_thread_modules),
  241. ('time', _green_time_modules),
  242. ('MySQLdb', _green_MySQLdb),
  243. ('builtins', _green_builtins),
  244. ('subprocess', _green_subprocess_modules),
  245. ]:
  246. if on[name] and not already_patched.get(name):
  247. modules_to_patch += modules_function()
  248. already_patched[name] = True
  249. if on['psycopg'] and not already_patched.get('psycopg'):
  250. try:
  251. from eventlet.support import psycopg2_patcher
  252. psycopg2_patcher.make_psycopg_green()
  253. already_patched['psycopg'] = True
  254. except ImportError:
  255. # note that if we get an importerror from trying to
  256. # monkeypatch psycopg, we will continually retry it
  257. # whenever monkey_patch is called; this should not be a
  258. # performance problem but it allows is_monkey_patched to
  259. # tell us whether or not we succeeded
  260. pass
  261. _threading = original('threading')
  262. imp.acquire_lock()
  263. try:
  264. for name, mod in modules_to_patch:
  265. orig_mod = sys.modules.get(name)
  266. if orig_mod is None:
  267. orig_mod = __import__(name)
  268. for attr_name in mod.__patched__:
  269. patched_attr = getattr(mod, attr_name, None)
  270. if patched_attr is not None:
  271. setattr(orig_mod, attr_name, patched_attr)
  272. deleted = getattr(mod, '__deleted__', [])
  273. for attr_name in deleted:
  274. if hasattr(orig_mod, attr_name):
  275. delattr(orig_mod, attr_name)
  276. # https://github.com/eventlet/eventlet/issues/592
  277. if name == 'threading' and register_at_fork:
  278. def fix_threading_active(
  279. _global_dict=_threading.current_thread.__globals__,
  280. # alias orig_mod as patched to reflect its new state
  281. # https://github.com/eventlet/eventlet/pull/661#discussion_r509877481
  282. _patched=orig_mod,
  283. ):
  284. _prefork_active = [None]
  285. def before_fork():
  286. _prefork_active[0] = _global_dict['_active']
  287. _global_dict['_active'] = _patched._active
  288. def after_fork():
  289. _global_dict['_active'] = _prefork_active[0]
  290. register_at_fork(
  291. before=before_fork,
  292. after_in_parent=after_fork)
  293. fix_threading_active()
  294. finally:
  295. imp.release_lock()
  296. if sys.version_info >= (3, 3):
  297. import importlib._bootstrap
  298. thread = original('_thread')
  299. # importlib must use real thread locks, not eventlet.Semaphore
  300. importlib._bootstrap._thread = thread
  301. # Issue #185: Since Python 3.3, threading.RLock is implemented in C and
  302. # so call a C function to get the thread identifier, instead of calling
  303. # threading.get_ident(). Force the Python implementation of RLock which
  304. # calls threading.get_ident() and so is compatible with eventlet.
  305. import threading
  306. threading.RLock = threading._PyRLock
  307. # Issue #508: Since Python 3.7 queue.SimpleQueue is implemented in C,
  308. # causing a deadlock. Replace the C implementation with the Python one.
  309. if sys.version_info >= (3, 7):
  310. import queue
  311. queue.SimpleQueue = queue._PySimpleQueue
  312. def is_monkey_patched(module):
  313. """Returns True if the given module is monkeypatched currently, False if
  314. not. *module* can be either the module itself or its name.
  315. Based entirely off the name of the module, so if you import a
  316. module some other way than with the import keyword (including
  317. import_patched), this might not be correct about that particular
  318. module."""
  319. return module in already_patched or \
  320. getattr(module, '__name__', None) in already_patched
  321. def _green_existing_locks():
  322. """Make locks created before monkey-patching safe.
  323. RLocks rely on a Lock and on Python 2, if an unpatched Lock blocks, it
  324. blocks the native thread. We need to replace these with green Locks.
  325. This was originally noticed in the stdlib logging module."""
  326. import gc
  327. import threading
  328. import eventlet.green.thread
  329. lock_type = type(threading.Lock())
  330. rlock_type = type(threading.RLock())
  331. if hasattr(threading, '_PyRLock'):
  332. # this happens on CPython3 and PyPy >= 7.0.0: "py3-style" rlocks, they
  333. # are implemented natively in C and RPython respectively
  334. py3_style = True
  335. pyrlock_type = type(threading._PyRLock())
  336. else:
  337. # this happens on CPython2.7 and PyPy < 7.0.0: "py2-style" rlocks,
  338. # they are implemented in pure-python
  339. py3_style = False
  340. pyrlock_type = None
  341. # We're monkey-patching so there can't be any greenlets yet, ergo our thread
  342. # ID is the only valid owner possible.
  343. tid = eventlet.green.thread.get_ident()
  344. for obj in gc.get_objects():
  345. if isinstance(obj, rlock_type):
  346. if not py3_style and isinstance(obj._RLock__block, lock_type):
  347. _fix_py2_rlock(obj, tid)
  348. elif py3_style and not isinstance(obj, pyrlock_type):
  349. _fix_py3_rlock(obj)
  350. def _fix_py2_rlock(rlock, tid):
  351. import eventlet.green.threading
  352. old = rlock._RLock__block
  353. new = eventlet.green.threading.Lock()
  354. rlock._RLock__block = new
  355. if old.locked():
  356. new.acquire()
  357. rlock._RLock__owner = tid
  358. def _fix_py3_rlock(old):
  359. import gc
  360. import threading
  361. new = threading._PyRLock()
  362. while old._is_owned():
  363. old.release()
  364. new.acquire()
  365. if old._is_owned():
  366. new.acquire()
  367. gc.collect()
  368. for ref in gc.get_referrers(old):
  369. try:
  370. ref_vars = vars(ref)
  371. except TypeError:
  372. pass
  373. else:
  374. for k, v in ref_vars.items():
  375. if v == old:
  376. setattr(ref, k, new)
  377. def _green_os_modules():
  378. from eventlet.green import os
  379. return [('os', os)]
  380. def _green_select_modules():
  381. from eventlet.green import select
  382. modules = [('select', select)]
  383. if sys.version_info >= (3, 4):
  384. from eventlet.green import selectors
  385. modules.append(('selectors', selectors))
  386. return modules
  387. def _green_socket_modules():
  388. from eventlet.green import socket
  389. try:
  390. from eventlet.green import ssl
  391. return [('socket', socket), ('ssl', ssl)]
  392. except ImportError:
  393. return [('socket', socket)]
  394. def _green_subprocess_modules():
  395. from eventlet.green import subprocess
  396. return [('subprocess', subprocess)]
  397. def _green_thread_modules():
  398. from eventlet.green import Queue
  399. from eventlet.green import thread
  400. from eventlet.green import threading
  401. if six.PY2:
  402. return [('Queue', Queue), ('thread', thread), ('threading', threading)]
  403. if six.PY3:
  404. return [('queue', Queue), ('_thread', thread), ('threading', threading)]
  405. def _green_time_modules():
  406. from eventlet.green import time
  407. return [('time', time)]
  408. def _green_MySQLdb():
  409. try:
  410. from eventlet.green import MySQLdb
  411. return [('MySQLdb', MySQLdb)]
  412. except ImportError:
  413. return []
  414. def _green_builtins():
  415. try:
  416. from eventlet.green import builtin
  417. return [('__builtin__' if six.PY2 else 'builtins', builtin)]
  418. except ImportError:
  419. return []
  420. def slurp_properties(source, destination, ignore=[], srckeys=None):
  421. """Copy properties from *source* (assumed to be a module) to
  422. *destination* (assumed to be a dict).
  423. *ignore* lists properties that should not be thusly copied.
  424. *srckeys* is a list of keys to copy, if the source's __all__ is
  425. untrustworthy.
  426. """
  427. if srckeys is None:
  428. srckeys = source.__all__
  429. destination.update(dict([
  430. (name, getattr(source, name))
  431. for name in srckeys
  432. if not (name.startswith('__') or name in ignore)
  433. ]))
  434. if __name__ == "__main__":
  435. sys.argv.pop(0)
  436. monkey_patch()
  437. with open(sys.argv[0]) as f:
  438. code = compile(f.read(), sys.argv[0], 'exec')
  439. exec(code)