From f38e756a3deb0ece2b5f74b79ae4c6b2f3d78b4b Mon Sep 17 00:00:00 2001
From: Aaron Piotrowski <aaron@trowski.com>
Date: Sun, 20 Feb 2022 12:33:37 -0600
Subject: [PATCH] Fixed memory leak in CombinedCancellationToken (#384)

Co-authored-by: Artur Khasanov <ah@self.team>
---
 src/CompositeCancellation.php                 | 18 +++++---
 .../CompositeCancellationTest.php             | 41 +++++++++++++++++++
 2 files changed, 53 insertions(+), 6 deletions(-)
 create mode 100644 test/Cancellation/CompositeCancellationTest.php

diff --git a/src/CompositeCancellation.php b/src/CompositeCancellation.php
index 1a8c8c7..678a65e 100644
--- a/src/CompositeCancellation.php
+++ b/src/CompositeCancellation.php
@@ -11,22 +11,28 @@ final class CompositeCancellation implements Cancellation
 
     private string $nextId = "a";
 
-    /** @var \Closure(CancelledException)[] */
+    /** @var array<string, \Closure(CancelledException): void> */
     private array $callbacks = [];
 
     private ?CancelledException $exception = null;
 
     public function __construct(Cancellation ...$cancellations)
     {
-        foreach ($cancellations as $cancellation) {
-            $id = $cancellation->subscribe(function (CancelledException $exception): void {
-                $this->exception = $exception;
+        $thatException = &$this->exception;
+        $thatCallbacks = &$this->callbacks;
 
-                foreach ($this->callbacks as $callback) {
+        foreach ($cancellations as $cancellation) {
+            $id = $cancellation->subscribe(static function (CancelledException $exception) use (
+                &$thatException,
+                &$thatCallbacks
+            ): void {
+                $thatException = $exception;
+
+                foreach ($thatCallbacks as $callback) {
                     EventLoop::queue($callback, $exception);
                 }
 
-                $this->callbacks = [];
+                $thatCallbacks = [];
             });
 
             $this->cancellations[] = [$cancellation, $id];
diff --git a/test/Cancellation/CompositeCancellationTest.php b/test/Cancellation/CompositeCancellationTest.php
new file mode 100644
index 0000000..8688e19
--- /dev/null
+++ b/test/Cancellation/CompositeCancellationTest.php
@@ -0,0 +1,41 @@
+<?php
+
+namespace Cancellation;
+
+use Amp\CompositeCancellation;
+use Amp\DeferredCancellation;
+use Amp\PHPUnit\AsyncTestCase;
+use function Amp\delay;
+
+class CompositeCancellationTest extends AsyncTestCase
+{
+    private const LOOP_COUNT = 20;
+    private const TOKENS_COUNT = 1000;
+
+    public function testBenchmark(): void
+    {
+        $firstMemoryMeasure = 0;
+
+        for ($i = 0; $i < self::LOOP_COUNT; $i++) {
+            $tokens = [];
+            for ($j = 0; $j < self::TOKENS_COUNT; $j++) {
+                $tokens[] = (new DeferredCancellation())->getCancellation();
+            }
+            $combinedToken = new CompositeCancellation(...$tokens);
+
+            if (!$firstMemoryMeasure && $i > self::LOOP_COUNT / 2) {
+                // Warmup and store first memory usage after 50% of iterations
+                $firstMemoryMeasure = \memory_get_usage(true);
+            }
+            // Remove tokens from memory
+            unset($combinedToken);
+
+            delay(0.001); // Tick loop to allow resources to be freed.
+
+            // Asserts
+            if ($firstMemoryMeasure > 0) {
+                self::assertEquals($firstMemoryMeasure, \memory_get_usage(true));
+            }
+        }
+    }
+}