Always stop monitoring fds in wait_{,conn_}async()

psycopg · Nov 12, 2021 · c017623 · c017623
1 parent 2150ac2
commit c017623
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/psycopg/psycopg/waiting.py b/psycopg/psycopg/waiting.py
@@ -135,11 +135,13 @@ def wakeup(state: Ready) -> None:
                 loop.add_reader(fileno, wakeup, Ready.R)
             if writer:
                 loop.add_writer(fileno, wakeup, Ready.W)
-            await ev.wait()
-            if reader:
-                loop.remove_reader(fileno)
-            if writer:
-                loop.remove_writer(fileno)
+            try:
+                await ev.wait()
+            finally:
+                if reader:
+                    loop.remove_reader(fileno)
+                if writer:
+                    loop.remove_writer(fileno)
             s = gen.send(ready)
 
     except StopIteration as ex:
@@ -188,11 +190,13 @@ def wakeup(state: Ready) -> None:
                 loop.add_reader(fileno, wakeup, Ready.R)
             if writer:
                 loop.add_writer(fileno, wakeup, Ready.W)
-            await wait_for(ev.wait(), timeout)
-            if reader:
-                loop.remove_reader(fileno)
-            if writer:
-                loop.remove_writer(fileno)
+            try:
+                await wait_for(ev.wait(), timeout)
+            finally:
+                if reader:
+                    loop.remove_reader(fileno)
+                if writer:
+                    loop.remove_writer(fileno)
             fileno, s = gen.send(ready)
 
     except TimeoutError: