Undoing Gevent's monkey-patching
Update
I'm a genius: simply executing reload(socket)
undoes Gevent's
patch_socket()
. Obviously, this only applies to new sockets created
after executing reload
, but that's good enough for my unittests. The
dumb solution below is preserved for hysterical porpoises.
Prior
I ran into an odd problem while testing the next release of PyMongo, the Python driver for MongoDB which I help develop. We're improving its support for Gevent, so we're of course doing additional tests that begin with:
from gevent import monkey; monkey.patch_socket()
Now, some tests rely on this patching, and some rely on not being
patched. Gevent doesn't provide an unpatch_socket
, so I had a clever
idea: I'll fork a subprocess with
multiprocessing,
do the test there, and return its result to the parent process in a
multiprocessing.Value
. Then subsequent tests won't be affected by the
patching.
SUCCESS = 1
FAILURE = 0
def my_test(outcome):
from gevent import monkey; monkey.patch_socket()
# do the test ....
outcome.value = SUCCESS
class Test(unittest.TestCase):
def test(self):
outcome = multiprocessing.Value('i', FAILURE)
multiprocessing.Process(
target=my_test,
args=(outcome,)
).start().join()
self.assertEqual(SUCCESS, outcome.value)
Nice and straightforward, right? In sane operating systems this worked
great. On Windows it broke horribly. When I did python setup.py test
,
instead of executing my_test()
, multiprocessing on Windows restarted
the whole test suite, which started another whole test suite, ...
Apparently, since Windows can't fork()
, multiprocessing re-imports
your script and attempts to execute the proper function within it. If
the test suite is begun with python setup.py test
, then everything
goes haywire. This problem with multiprocessing and unittests on
Windows
was discussed on the Python mailing list last February.
After some gloomy minutes, I decided to look at what patch_socket()
is
doing. Turns out it's simple, so I wrote a version which allows
unpatching:
def patch_socket(aggressive=True):
"""Like gevent.monkey.patch_socket(), but stores old socket attributes
for unpatching.
"""
from gevent import socket
_socket = __import__('socket')
old_attrs = {}
for attr in (
'socket', 'SocketType', 'create_connection', 'socketpair', 'fromfd'
):
if hasattr(_socket, attr):
old_attrs[attr] = getattr(_socket, attr)
setattr(_socket, attr, getattr(socket, attr))
try:
from gevent.socket import ssl, sslerror
old_attrs['ssl'] = _socket.ssl
_socket.ssl = ssl
old_attrs['sslerror'] = _socket.sslerror
_socket.sslerror = sslerror
except ImportError:
if aggressive:
try:
del _socket.ssl
except AttributeError:
pass
return old_attrs
def unpatch_socket(old_attrs):
"""Take output of patch_socket() and undo patching."""
_socket = __import__('socket')
for attr in old_attrs:
if hasattr(_socket, attr):
setattr(_socket, attr, old_attrs[attr])
def patch_dns():
"""Like gevent.monkey.patch_dns(), but stores old socket attributes
for unpatching.
"""
from gevent.socket import gethostbyname, getaddrinfo
_socket = __import__('socket')
old_attrs = {}
old_attrs['getaddrinfo'] = _socket.getaddrinfo
_socket.getaddrinfo = getaddrinfo
old_attrs['gethostbyname'] = _socket.gethostbyname
_socket.gethostbyname = gethostbyname
return old_attrs
def unpatch_dns(old_attrs):
"""Take output of patch_dns() and undo patching."""
_socket = __import__('socket')
for attr in old_attrs:
setattr(_socket, attr, old_attrs[attr])
In Gevent's version, calling patch_socket()
calls patch_dns()
implicitly, in mine you must call both:
class Test(unittest.TestCase):
def test(self):
old_socket_attrs = patch_socket()
old_dns_attrs = patch_dns()
try:
# do test ...
finally:
unpatch_dns(old_dns_attrs)
unpatch_socket(old_socket_attrs)
Now I don't need multiprocessing at all.