patcher.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. import sys
  2. import imp
  3. __all__ = ['inject', 'import_patched', 'monkey_patch', 'is_monkey_patched']
  4. __exclude = set(('__builtins__', '__file__', '__name__'))
  5. class SysModulesSaver(object):
  6. """Class that captures some subset of the current state of
  7. sys.modules. Pass in an iterator of module names to the
  8. constructor."""
  9. def __init__(self, module_names=()):
  10. self._saved = {}
  11. imp.acquire_lock()
  12. self.save(*module_names)
  13. def save(self, *module_names):
  14. """Saves the named modules to the object."""
  15. for modname in module_names:
  16. self._saved[modname] = sys.modules.get(modname, None)
  17. def restore(self):
  18. """Restores the modules that the saver knows about into
  19. sys.modules.
  20. """
  21. try:
  22. for modname, mod in self._saved.iteritems():
  23. if mod is not None:
  24. sys.modules[modname] = mod
  25. else:
  26. try:
  27. del sys.modules[modname]
  28. except KeyError:
  29. pass
  30. finally:
  31. imp.release_lock()
  32. def inject(module_name, new_globals, *additional_modules):
  33. """Base method for "injecting" greened modules into an imported module. It
  34. imports the module specified in *module_name*, arranging things so
  35. that the already-imported modules in *additional_modules* are used when
  36. *module_name* makes its imports.
  37. *new_globals* is either None or a globals dictionary that gets populated
  38. with the contents of the *module_name* module. This is useful when creating
  39. a "green" version of some other module.
  40. *additional_modules* should be a collection of two-element tuples, of the
  41. form (<name>, <module>). If it's not specified, a default selection of
  42. name/module pairs is used, which should cover all use cases but may be
  43. slower because there are inevitably redundant or unnecessary imports.
  44. """
  45. patched_name = '__patched_module_' + module_name
  46. if patched_name in sys.modules:
  47. # returning already-patched module so as not to destroy existing
  48. # references to patched modules
  49. return sys.modules[patched_name]
  50. if not additional_modules:
  51. # supply some defaults
  52. additional_modules = (
  53. _green_os_modules() +
  54. _green_select_modules() +
  55. _green_socket_modules() +
  56. _green_thread_modules() +
  57. _green_time_modules())
  58. #_green_MySQLdb()) # enable this after a short baking-in period
  59. # after this we are gonna screw with sys.modules, so capture the
  60. # state of all the modules we're going to mess with, and lock
  61. saver = SysModulesSaver([name for name, m in additional_modules])
  62. saver.save(module_name)
  63. # Cover the target modules so that when you import the module it
  64. # sees only the patched versions
  65. for name, mod in additional_modules:
  66. sys.modules[name] = mod
  67. ## Remove the old module from sys.modules and reimport it while
  68. ## the specified modules are in place
  69. sys.modules.pop(module_name, None)
  70. try:
  71. module = __import__(module_name, {}, {}, module_name.split('.')[:-1])
  72. if new_globals is not None:
  73. ## Update the given globals dictionary with everything from this new module
  74. for name in dir(module):
  75. if name not in __exclude:
  76. new_globals[name] = getattr(module, name)
  77. ## Keep a reference to the new module to prevent it from dying
  78. sys.modules[patched_name] = module
  79. finally:
  80. saver.restore() ## Put the original modules back
  81. return module
  82. def import_patched(module_name, *additional_modules, **kw_additional_modules):
  83. """Imports a module in a way that ensures that the module uses "green"
  84. versions of the standard library modules, so that everything works
  85. nonblockingly.
  86. The only required argument is the name of the module to be imported.
  87. """
  88. return inject(
  89. module_name,
  90. None,
  91. *additional_modules + tuple(kw_additional_modules.items()))
  92. def patch_function(func, *additional_modules):
  93. """Decorator that returns a version of the function that patches
  94. some modules for the duration of the function call. This is
  95. deeply gross and should only be used for functions that import
  96. network libraries within their function bodies that there is no
  97. way of getting around."""
  98. if not additional_modules:
  99. # supply some defaults
  100. additional_modules = (
  101. _green_os_modules() +
  102. _green_select_modules() +
  103. _green_socket_modules() +
  104. _green_thread_modules() +
  105. _green_time_modules())
  106. def patched(*args, **kw):
  107. saver = SysModulesSaver()
  108. for name, mod in additional_modules:
  109. saver.save(name)
  110. sys.modules[name] = mod
  111. try:
  112. return func(*args, **kw)
  113. finally:
  114. saver.restore()
  115. return patched
  116. def _original_patch_function(func, *module_names):
  117. """Kind of the contrapositive of patch_function: decorates a
  118. function such that when it's called, sys.modules is populated only
  119. with the unpatched versions of the specified modules. Unlike
  120. patch_function, only the names of the modules need be supplied,
  121. and there are no defaults. This is a gross hack; tell your kids not
  122. to import inside function bodies!"""
  123. def patched(*args, **kw):
  124. saver = SysModulesSaver(module_names)
  125. for name in module_names:
  126. sys.modules[name] = original(name)
  127. try:
  128. return func(*args, **kw)
  129. finally:
  130. saver.restore()
  131. return patched
  132. def original(modname):
  133. """ This returns an unpatched version of a module; this is useful for
  134. Eventlet itself (i.e. tpool)."""
  135. # note that it's not necessary to temporarily install unpatched
  136. # versions of all patchable modules during the import of the
  137. # module; this is because none of them import each other, except
  138. # for threading which imports thread
  139. original_name = '__original_module_' + modname
  140. if original_name in sys.modules:
  141. return sys.modules.get(original_name)
  142. # re-import the "pure" module and store it in the global _originals
  143. # dict; be sure to restore whatever module had that name already
  144. saver = SysModulesSaver((modname,))
  145. sys.modules.pop(modname, None)
  146. # some rudimentary dependency checking -- fortunately the modules
  147. # we're working on don't have many dependencies so we can just do
  148. # some special-casing here
  149. deps = {'threading':'thread', 'Queue':'threading'}
  150. if modname in deps:
  151. dependency = deps[modname]
  152. saver.save(dependency)
  153. sys.modules[dependency] = original(dependency)
  154. try:
  155. real_mod = __import__(modname, {}, {}, modname.split('.')[:-1])
  156. if modname == 'Queue' and not hasattr(real_mod, '_threading'):
  157. # tricky hack: Queue's constructor in <2.7 imports
  158. # threading on every instantiation; therefore we wrap
  159. # it so that it always gets the original threading
  160. real_mod.Queue.__init__ = _original_patch_function(
  161. real_mod.Queue.__init__,
  162. 'threading')
  163. # save a reference to the unpatched module so it doesn't get lost
  164. sys.modules[original_name] = real_mod
  165. finally:
  166. saver.restore()
  167. return sys.modules[original_name]
  168. already_patched = {}
  169. def monkey_patch(**on):
  170. """Globally patches certain system modules to be greenthread-friendly.
  171. The keyword arguments afford some control over which modules are patched.
  172. If no keyword arguments are supplied, all possible modules are patched.
  173. If keywords are set to True, only the specified modules are patched. E.g.,
  174. ``monkey_patch(socket=True, select=True)`` patches only the select and
  175. socket modules. Most arguments patch the single module of the same name
  176. (os, time, select). The exceptions are socket, which also patches the ssl
  177. module if present; and thread, which patches thread, threading, and Queue.
  178. It's safe to call monkey_patch multiple times.
  179. """
  180. accepted_args = set(('os', 'select', 'socket',
  181. 'thread', 'time', 'psycopg', 'MySQLdb'))
  182. default_on = on.pop("all",None)
  183. for k in on.iterkeys():
  184. if k not in accepted_args:
  185. raise TypeError("monkey_patch() got an unexpected "\
  186. "keyword argument %r" % k)
  187. if default_on is None:
  188. default_on = not (True in on.values())
  189. for modname in accepted_args:
  190. if modname == 'MySQLdb':
  191. # MySQLdb is only on when explicitly patched for the moment
  192. on.setdefault(modname, False)
  193. on.setdefault(modname, default_on)
  194. modules_to_patch = []
  195. patched_thread = False
  196. if on['os'] and not already_patched.get('os'):
  197. modules_to_patch += _green_os_modules()
  198. already_patched['os'] = True
  199. if on['select'] and not already_patched.get('select'):
  200. modules_to_patch += _green_select_modules()
  201. already_patched['select'] = True
  202. if on['socket'] and not already_patched.get('socket'):
  203. modules_to_patch += _green_socket_modules()
  204. already_patched['socket'] = True
  205. if on['thread'] and not already_patched.get('thread'):
  206. patched_thread = True
  207. modules_to_patch += _green_thread_modules()
  208. already_patched['thread'] = True
  209. if on['time'] and not already_patched.get('time'):
  210. modules_to_patch += _green_time_modules()
  211. already_patched['time'] = True
  212. if on.get('MySQLdb') and not already_patched.get('MySQLdb'):
  213. modules_to_patch += _green_MySQLdb()
  214. already_patched['MySQLdb'] = True
  215. if on['psycopg'] and not already_patched.get('psycopg'):
  216. try:
  217. from eventlet.support import psycopg2_patcher
  218. psycopg2_patcher.make_psycopg_green()
  219. already_patched['psycopg'] = True
  220. except ImportError:
  221. # note that if we get an importerror from trying to
  222. # monkeypatch psycopg, we will continually retry it
  223. # whenever monkey_patch is called; this should not be a
  224. # performance problem but it allows is_monkey_patched to
  225. # tell us whether or not we succeeded
  226. pass
  227. imp.acquire_lock()
  228. try:
  229. for name, mod in modules_to_patch:
  230. orig_mod = sys.modules.get(name)
  231. if orig_mod is None:
  232. orig_mod = __import__(name)
  233. for attr_name in mod.__patched__:
  234. patched_attr = getattr(mod, attr_name, None)
  235. if patched_attr is not None:
  236. setattr(orig_mod, attr_name, patched_attr)
  237. # hacks ahead; this is necessary to prevent a KeyError on program exit
  238. if patched_thread:
  239. _patch_main_thread(sys.modules['threading'])
  240. finally:
  241. imp.release_lock()
  242. def _patch_main_thread(mod):
  243. """This is some gnarly patching specific to the threading module;
  244. threading will always be initialized prior to monkeypatching, and
  245. its _active dict will have the wrong key (it uses the real thread
  246. id but once it's patched it will use the greenlet ids); so what we
  247. do is rekey the _active dict so that the main thread's entry uses
  248. the greenthread key. Other threads' keys are ignored."""
  249. thread = original('thread')
  250. curthread = mod._active.pop(thread.get_ident(), None)
  251. if curthread:
  252. import eventlet.green.thread
  253. mod._active[eventlet.green.thread.get_ident()] = curthread
  254. def is_monkey_patched(module):
  255. """Returns True if the given module is monkeypatched currently, False if
  256. not. *module* can be either the module itself or its name.
  257. Based entirely off the name of the module, so if you import a
  258. module some other way than with the import keyword (including
  259. import_patched), this might not be correct about that particular
  260. module."""
  261. return module in already_patched or \
  262. getattr(module, '__name__', None) in already_patched
  263. def _green_os_modules():
  264. from eventlet.green import os
  265. return [('os', os)]
  266. def _green_select_modules():
  267. from eventlet.green import select
  268. return [('select', select)]
  269. def _green_socket_modules():
  270. from eventlet.green import socket
  271. try:
  272. from eventlet.green import ssl
  273. return [('socket', socket), ('ssl', ssl)]
  274. except ImportError:
  275. return [('socket', socket)]
  276. def _green_thread_modules():
  277. from eventlet.green import Queue
  278. from eventlet.green import thread
  279. from eventlet.green import threading
  280. return [('Queue', Queue), ('thread', thread), ('threading', threading)]
  281. def _green_time_modules():
  282. from eventlet.green import time
  283. return [('time', time)]
  284. def _green_MySQLdb():
  285. try:
  286. from eventlet.green import MySQLdb
  287. return [('MySQLdb', MySQLdb)]
  288. except ImportError:
  289. return []
  290. if __name__ == "__main__":
  291. import sys
  292. sys.argv.pop(0)
  293. monkey_patch()
  294. execfile(sys.argv[0])