[pypy-commit] pypy default: issue #2064: trying to call socket.close() on all sockets at exit, on Windows

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r78222:ddbc9f97d312
Date: 2015-06-21 11:39 +0200
http://bitbucket.org/pypy/pypy/changeset/ddbc9f97d312/

Log:	issue #2064: trying to call socket.close() on all sockets at exit,
	on Windows

diff --git a/pypy/module/_socket/__init__.py b/pypy/module/_socket/__init__.py
--- a/pypy/module/_socket/__init__.py
+++ b/pypy/module/_socket/__init__.py
@@ -18,6 +18,10 @@
         from rpython.rlib.rsocket import rsocket_startup
         rsocket_startup()
 
+    def shutdown(self, space):
+        from pypy.module._socket.interp_socket import close_all_sockets
+        close_all_sockets(space)
+
     def buildloaders(cls):
         from rpython.rlib import rsocket
         for name in """
diff --git a/pypy/module/_socket/interp_func.py b/pypy/module/_socket/interp_func.py
--- a/pypy/module/_socket/interp_func.py
+++ b/pypy/module/_socket/interp_func.py
@@ -142,7 +142,7 @@
         sock = rsocket.fromfd(fd, family, type, proto)
     except SocketError, e:
         raise converted_error(space, e)
-    return space.wrap(W_Socket(sock))
+    return space.wrap(W_Socket(space, sock))
 
 @unwrap_spec(family=int, type=int, proto=int)
 def socketpair(space, family=rsocket.socketpair_default_family,
@@ -160,8 +160,8 @@
     except SocketError, e:
         raise converted_error(space, e)
     return space.newtuple([
-        space.wrap(W_Socket(sock1)),
-        space.wrap(W_Socket(sock2))
+        space.wrap(W_Socket(space, sock1)),
+        space.wrap(W_Socket(space, sock2))
     ])
 
 # The following 4 functions refuse all negative numbers, like CPython 2.6.
diff --git a/pypy/module/_socket/interp_socket.py b/pypy/module/_socket/interp_socket.py
--- a/pypy/module/_socket/interp_socket.py
+++ b/pypy/module/_socket/interp_socket.py
@@ -1,4 +1,5 @@
-from rpython.rlib import rsocket
+import sys
+from rpython.rlib import rsocket, rweaklist
 from rpython.rlib.rarithmetic import intmask
 from rpython.rlib.rsocket import (
     RSocket, AF_INET, SOCK_STREAM, SocketError, SocketErrorWithErrno,
@@ -153,8 +154,9 @@
 
 
 class W_Socket(W_Root):
-    def __init__(self, sock):
+    def __init__(self, space, sock):
         self.sock = sock
+        register_socket(space, sock)
 
     def get_type_w(self, space):
         return space.wrap(self.sock.type)
@@ -183,7 +185,7 @@
             fd, addr = self.sock.accept()
             sock = rsocket.make_socket(
                 fd, self.sock.family, self.sock.type, self.sock.proto)
-            return space.newtuple([space.wrap(W_Socket(sock)),
+            return space.newtuple([space.wrap(W_Socket(space, sock)),
                                    addr_as_object(addr, sock.fd, space)])
         except SocketError as e:
             raise converted_error(space, e)
@@ -248,7 +250,7 @@
     def dup_w(self, space):
         try:
             sock = self.sock.dup()
-            return W_Socket(sock)
+            return W_Socket(space, sock)
         except SocketError as e:
             raise converted_error(space, e)
 
@@ -592,10 +594,50 @@
         sock = RSocket(family, type, proto)
     except SocketError as e:
         raise converted_error(space, e)
-    W_Socket.__init__(self, sock)
+    W_Socket.__init__(self, space, sock)
     return space.wrap(self)
 descr_socket_new = interp2app(newsocket)
 
+
+# ____________________________________________________________
+# Automatic shutdown()/close()
+
+# On some systems, the C library does not guarantee that when the program
+# finishes, all data sent so far is really sent even if the socket is not
+# explicitly closed.  This behavior has been observed on Windows but not
+# on Linux, so far.
+NEED_EXPLICIT_CLOSE = (sys.platform == 'win32')
+
+class OpenRSockets(rweaklist.RWeakListMixin):
+    pass
+class OpenRSocketsState:
+    def __init__(self, space):
+        self.openrsockets = OpenRSockets()
+        self.openrsockets.initialize()
+
+def getopenrsockets(space):
+    if NEED_EXPLICIT_CLOSE and space.config.translation.rweakref:
+        return space.fromcache(OpenRSocketsState).openrsockets
+    else:
+        return None
+
+def register_socket(space, socket):
+    openrsockets = getopenrsockets(space)
+    if openrsockets is not None:
+        openrsockets.add_handle(socket)
+
+def close_all_sockets(space):
+    openrsockets = getopenrsockets(space)
+    if openrsockets is not None:
+        for sock_wref in openrsockets.get_all_handles():
+            sock = sock_wref()
+            if sock is not None:
+                try:
+                    sock.close()
+                except SocketError:
+                    pass
+
+
 # ____________________________________________________________
 # Error handling
 
diff --git a/pypy/module/_socket/test/test_sock_app.py b/pypy/module/_socket/test/test_sock_app.py
--- a/pypy/module/_socket/test/test_sock_app.py
+++ b/pypy/module/_socket/test/test_sock_app.py
@@ -309,10 +309,15 @@
 
 
 class AppTestSocket:
+    spaceconfig = dict(usemodules=['_socket', '_weakref', 'struct'])
+
     def setup_class(cls):
         cls.space = space
         cls.w_udir = space.wrap(str(udir))
 
+    def teardown_class(cls):
+        cls.space.sys.getmodule('_socket').shutdown(cls.space)
+
     def test_module(self):
         import _socket
         assert _socket.socket.__name__ == 'socket'
@@ -614,6 +619,12 @@
         finally:
             os.chdir(oldcwd)
 
+    def test_automatic_shutdown(self):
+        # doesn't really test anything, but at least should not explode
+        # in close_all_sockets()
+        import _socket
+        self.foo = _socket.socket()
+
 
 class AppTestPacket:
     def setup_class(cls):