Merge pull request #845 from rthalley/asyncio-timeout-fix

Asyncio sockets should work after a timeout [#843].

1 files changed, 7 insertions, 5 deletions

diff --git a/dns/_asyncio_backend.py b/dns/_asyncio_backend.py
index 50bde1d..736539b 100644
--- a/dns/_asyncio_backend.py
+++ b/dns/_asyncio_backend.py
@@ -31,7 +31,6 @@ class _DatagramProtocol:
     def datagram_received(self, data, addr):
         if self.recvfrom and not self.recvfrom.done():
             self.recvfrom.set_result((data, addr))
-            self.recvfrom = None
 
     def error_received(self, exc):  # pragma: no cover
         if self.recvfrom and not self.recvfrom.done():
@@ -68,10 +67,13 @@ class DatagramSocket(dns._asyncbackend.DatagramSocket):
     async def recvfrom(self, size, timeout):
         # ignore size as there's no way I know to tell protocol about it
         done = _get_running_loop().create_future()
-        assert self.protocol.recvfrom is None
-        self.protocol.recvfrom = done
-        await _maybe_wait_for(done, timeout)
-        return done.result()
+        try:
+            assert self.protocol.recvfrom is None
+            self.protocol.recvfrom = done
+            await _maybe_wait_for(done, timeout)
+            return done.result()
+        finally:
+            self.protocol.recvfrom = None
 
     async def close(self):
         self.protocol.close()