Keep receiving if there are still pending requests

lib/BasicResolver.php

@@ -4,7 +4,6 @@ namespace Amp\Dns;
 
 use Amp\Cache\ArrayCache;
 use Amp\Cache\Cache;
-use Amp\Coroutine;
 use Amp\MultiReasonException;
 use Amp\Promise;
 use LibDNS\Messages\Message;
@@ -111,7 +110,7 @@ class BasicResolver implements Resolver {
                 }
             }
 
-            return array_merge(...$records);
+            return \array_merge(...$records);
         });
     }
 
@@ -135,10 +134,7 @@ class BasicResolver implements Resolver {
 
     /** @inheritdoc */
     public function query(string $name, int $type): Promise {
-        return new Coroutine($this->doQuery($name, $type));
+        return call(function () use ($name, $type) {
-    }
-
-    public function doQuery(string $name, int $type): \Generator {
             if (!$this->config) {
                 $this->config = yield $this->configLoader->loadConfig();
             }
@@ -206,12 +202,13 @@ class BasicResolver implements Resolver {
                     throw new NoRecordException("No records returned for {$name}");
                 }
 
-            return array_map(function ($data) use ($type, $ttls) {
+                return \array_map(function ($data) use ($type, $ttls) {
                     return new Record($data, $type, $ttls[$type]);
                 }, $result[$type]);
             }
 
             throw new ResolutionException("No response from any nameserver after {$attempts} attempts");
+        });
     }
 
     /**
@@ -269,9 +266,9 @@ class BasicResolver implements Resolver {
         if ($type === Record::PTR) {
             if (($packedIp = @inet_pton($name)) !== false) {
                 if (isset($packedIp[4])) { // IPv6
-                    $name = wordwrap(strrev(bin2hex($packedIp)), 1, ".", true) . ".ip6.arpa";
+                    $name = \wordwrap(\strrev(\bin2hex($packedIp)), 1, ".", true) . ".ip6.arpa";
                 } else { // IPv4
-                    $name = inet_ntop(strrev($packedIp)) . ".in-addr.arpa";
+                    $name = \inet_ntop(\strrev($packedIp)) . ".in-addr.arpa";
                 }
             }
         } else if (\in_array($type, [Record::A, Record::AAAA])) {
@@ -286,7 +283,7 @@ class BasicResolver implements Resolver {
             return $this->servers[$uri];
         }
 
-        if (substr($uri, 0, 3) === "udp") {
+        if (\substr($uri, 0, 3) === "udp") {
             $server = UdpServer::connect($uri);
         } else {
             $server = TcpServer::connect($uri);

lib/Server.php

@@ -4,13 +4,14 @@ namespace Amp\Dns;
 
 use Amp\ByteStream\ResourceInputStream;
 use Amp\ByteStream\ResourceOutputStream;
-use Amp\Coroutine;
+use Amp\ByteStream\StreamException;
 use Amp\Deferred;
 use Amp\Promise;
 use LibDNS\Messages\Message;
 use LibDNS\Messages\MessageFactory;
 use LibDNS\Messages\MessageTypes;
 use LibDNS\Records\Question;
+use function Amp\call;
 
 abstract class Server {
     /** @var \Amp\ByteStream\ResourceInputStream */
@@ -61,31 +62,47 @@ abstract class Server {
 
         $this->onResolve = function (\Throwable $exception = null, Message $message = null) {
             if ($exception) {
+                $questions = $this->questions;
+                $this->questions = [];
+                foreach ($questions as $deferred) {
+                    $deferred->fail($exception);
+                }
                 return;
             }
 
             $id = $message->getId();
 
             if (!isset($this->questions[$id])) {
-                return;
+                return; // Ignore duplicate response.
             }
 
             $deferred = $this->questions[$id];
             unset($this->questions[$id]);
+
+            $empty = empty($this->questions);
+
             $deferred->resolve($message);
+
+            if (!$empty) {
+                $this->receive()->onResolve($this->onResolve);
+            }
         };
     }
 
+    /**
+     * @param \LibDNS\Records\Question $question
+     *
+     * @return \Amp\Promise<\LibDNS\Messages\Message>
+     */
     public function ask(Question $question): Promise {
-        return new Coroutine($this->doAsk($question));
+        return call(function () use ($question) {
-    }
-
-    private function doAsk(Question $question): \Generator {
             $id = $this->nextId++;
             if ($this->nextId > 0xffff) {
                 $this->nextId %= 0xffff;
             }
 
+            $empty = empty($this->questions);
+
             if (isset($this->questions[$id])) {
                 $deferred = $this->questions[$id];
                 unset($this->questions[$id]);
@@ -93,12 +110,21 @@ abstract class Server {
             }
 
             $message = $this->createMessage($question, $id);
+
+            try {
+                yield $this->send($message);
+            } catch (StreamException $exception) {
+                throw new ResolutionException("Sending the request failed", 0, $exception);
+            }
+
+            if ($empty) {
+                $this->receive()->onResolve($this->onResolve);
+            }
+
             $this->questions[$id] = $deferred = new Deferred;
 
-        yield $this->send($message);
-        $this->receive()->onResolve($this->onResolve);
-
             return yield $deferred->promise();
+        });
     }
 
     protected function read(): Promise {