Fix remove_reader and remove_writer to only decref prev added sockets

Author: 1st1

tests/test_sockets.py

@@ -469,6 +469,42 @@ async def client(sock, addr):
                 self.loop.close()
                 self.assertEqual(sock.fileno(), -1)
 
+    def test_socket_close_remove_reader(self):
+        s = socket.socket()
+        with s:
+            s.setblocking(False)
+            self.loop.add_reader(s, lambda: None)
+            self.loop.remove_reader(s.fileno())
+            s.close()
+            self.assertEqual(s.fileno(), -1)
+
+        s = socket.socket()
+        with s:
+            s.setblocking(False)
+            self.loop.add_reader(s.fileno(), lambda: None)
+            self.loop.remove_reader(s)
+            self.assertNotEqual(s.fileno(), -1)
+            s.close()
+            self.assertEqual(s.fileno(), -1)
+
+    def test_socket_close_remove_writer(self):
+        s = socket.socket()
+        with s:
+            s.setblocking(False)
+            self.loop.add_writer(s, lambda: None)
+            self.loop.remove_writer(s.fileno())
+            s.close()
+            self.assertEqual(s.fileno(), -1)
+
+        s = socket.socket()
+        with s:
+            s.setblocking(False)
+            self.loop.add_writer(s.fileno(), lambda: None)
+            self.loop.remove_writer(s)
+            self.assertNotEqual(s.fileno(), -1)
+            s.close()
+            self.assertEqual(s.fileno(), -1)
+
 
 class TestAIOSockets(_TestSockets, tb.AIOTestCase):
     pass

uvloop/loop.pyx

@@ -664,8 +664,9 @@ cdef class Loop:
             poll = UVPoll.new(self, fd)
             self._polls[fd] = poll
 
-        self._fd_to_reader_fileobj[fd] = fileobj
         poll.start_reading(handle)
+
+        self._fd_to_reader_fileobj[fd] = fileobj
         socket_inc_io_ref(fileobj)
 
     cdef _remove_reader(self, fileobj):
@@ -674,8 +675,10 @@ cdef class Loop:
 
         fd = self._fileobj_to_fd(fileobj)
         self._ensure_fd_no_transport(fd)
-        self._fd_to_reader_fileobj.pop(fd, None)
-        socket_dec_io_ref(fileobj)
+
+        mapped_fileobj = self._fd_to_reader_fileobj.pop(fd, None)
+        if mapped_fileobj is not None:
+            socket_dec_io_ref(mapped_fileobj)
 
         if self._closed == 1:
             return False
@@ -706,8 +709,9 @@ cdef class Loop:
             poll = UVPoll.new(self, fd)
             self._polls[fd] = poll
 
-        self._fd_to_writer_fileobj[fd] = fileobj
         poll.start_writing(handle)
+
+        self._fd_to_writer_fileobj[fd] = fileobj
         socket_inc_io_ref(fileobj)
 
     cdef _remove_writer(self, fileobj):
@@ -716,8 +720,10 @@ cdef class Loop:
 
         fd = self._fileobj_to_fd(fileobj)
         self._ensure_fd_no_transport(fd)
-        self._fd_to_writer_fileobj.pop(fd, None)
-        socket_dec_io_ref(fileobj)
+
+        mapped_fileobj = self._fd_to_writer_fileobj.pop(fd, None)
+        if mapped_fileobj is not None:
+            socket_dec_io_ref(mapped_fileobj)
 
         if self._closed == 1:
             return False