Support cancellation of FallbackExecutor

Author: clue

src/Query/FallbackExecutor.php

@@ -2,6 +2,8 @@
 
 namespace React\Dns\Query;
 
+use React\Promise\Promise;
+
 final class FallbackExecutor implements ExecutorInterface
 {
     private $executor;
@@ -15,16 +17,33 @@ public function __construct(ExecutorInterface $executor, ExecutorInterface $fall
 
     public function query(Query $query)
     {
+        $cancelled = false;
         $fallback = $this->fallback;
-        return $this->executor->query($query)->then(null, function (\Exception $e1) use ($query, $fallback) {
-            return $fallback->query($query)->then(null, function (\Exception $e2) use ($e1) {
-                $append = $e2->getMessage();
-                if (($pos = strpos($append, ':')) !== false) {
-                    $append = substr($append, $pos + 2);
+        $promise = $this->executor->query($query);
+
+        return new Promise(function ($resolve, $reject) use (&$promise, $fallback, $query, &$cancelled) {
+            $promise->then($resolve, function (\Exception $e1) use ($fallback, $query, $resolve, $reject, &$cancelled, &$promise) {
+                // reject if primary resolution rejected due to cancellation
+                if ($cancelled) {
+                    $reject($e1);
+                    return;
                 }
 
-                throw new \RuntimeException($e1->getMessage() . '. ' . $append);
+                // start fallback query if primary query rejected
+                $promise = $fallback->query($query)->then($resolve, function (\Exception $e2) use ($e1, $reject) {
+                    $append = $e2->getMessage();
+                    if (($pos = strpos($append, ':')) !== false) {
+                        $append = substr($append, $pos + 2);
+                    }
+
+                    // reject with combined error message if both queries fail
+                    $reject(new \RuntimeException($e1->getMessage() . '. ' . $append));
+                });
             });
+        }, function () use (&$promise, &$cancelled) {
+            // cancel pending query (primary or fallback)
+            $cancelled = true;
+            $promise->cancel();
         });
     }
 }

tests/Query/FallbackExecutorTest.php

@@ -134,8 +134,6 @@ public function testQueryWillRejectWithExceptionMessagesConcatenatedInFullWhenPr
 
     public function testCancelQueryWillReturnRejectedPromiseWithoutCallingSecondaryExecutorWhenPrimaryExecutorIsStillPending()
     {
-        $this->markTestIncomplete();
-
         $query = new Query('reactphp.org', Message::TYPE_A, Message::CLASS_IN);
 
         $primary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
@@ -152,4 +150,75 @@ public function testCancelQueryWillReturnRejectedPromiseWithoutCallingSecondaryE
         $this->assertInstanceOf('React\Promise\PromiseInterface', $promise);
         $promise->then($this->expectCallableNever(), $this->expectCallableOnce());
     }
+
+    public function testCancelQueryWillReturnRejectedPromiseWhenPrimaryExecutorRejectsAndSecondaryExecutorIsStillPending()
+    {
+        $query = new Query('reactphp.org', Message::TYPE_A, Message::CLASS_IN);
+
+        $primary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
+        $primary->expects($this->once())->method('query')->with($query)->willReturn(\React\Promise\reject(new \RuntimeException()));
+
+        $secondary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
+        $secondary->expects($this->once())->method('query')->with($query)->willReturn(new Promise(function () { }, function () { throw new \RuntimeException(); }));
+
+        $executor = new FallbackExecutor($primary, $secondary);
+
+        $promise = $executor->query($query);
+        $promise->cancel();
+
+        $this->assertInstanceOf('React\Promise\PromiseInterface', $promise);
+        $promise->then($this->expectCallableNever(), $this->expectCallableOnce());
+    }
+
+    public function testCancelQueryShouldNotCauseGarbageReferencesWhenCancellingPrimaryExecutor()
+    {
+        if (class_exists('React\Promise\When')) {
+            $this->markTestSkipped('Not supported on legacy Promise v1 API');
+        }
+
+        $primary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
+        $primary->expects($this->once())->method('query')->willReturn(new Promise(function () { }, function () { throw new \RuntimeException(); }));
+
+        $secondary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
+        $secondary->expects($this->never())->method('query');
+
+        $executor = new FallbackExecutor($primary, $secondary);
+
+        gc_collect_cycles();
+        gc_collect_cycles(); // clear twice to avoid leftovers in PHP 7.4 with ext-xdebug and code coverage turned on
+
+        $query = new Query('reactphp.org', Message::TYPE_A, Message::CLASS_IN);
+
+        $promise = $executor->query($query);
+        $promise->cancel();
+        $promise = null;
+
+        $this->assertEquals(0, gc_collect_cycles());
+    }
+
+    public function testCancelQueryShouldNotCauseGarbageReferencesWhenCancellingSecondaryExecutor()
+    {
+        if (class_exists('React\Promise\When')) {
+            $this->markTestSkipped('Not supported on legacy Promise v1 API');
+        }
+
+        $primary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
+        $primary->expects($this->once())->method('query')->willReturn(\React\Promise\reject(new \RuntimeException()));
+
+        $secondary = $this->getMockBuilder('React\Dns\Query\ExecutorInterface')->getMock();
+        $secondary->expects($this->once())->method('query')->willReturn(new Promise(function () { }, function () { throw new \RuntimeException(); }));
+
+        $executor = new FallbackExecutor($primary, $secondary);
+
+        gc_collect_cycles();
+        gc_collect_cycles(); // clear twice to avoid leftovers in PHP 7.4 with ext-xdebug and code coverage turned on
+
+        $query = new Query('reactphp.org', Message::TYPE_A, Message::CLASS_IN);
+
+        $promise = $executor->query($query);
+        $promise->cancel();
+        $promise = null;
+
+        $this->assertEquals(0, gc_collect_cycles());
+    }
 }