From 122aff0c74600f8f839ffeb0186efe2aa2c750a2 Mon Sep 17 00:00:00 2001
From: Tobias Bachert <git@b-privat.de>
Date: Tue, 21 Jun 2022 02:32:31 +0200
Subject: [PATCH] Prevent leaking `TracerProvider` references in
 `register_shutdown_function` (#716)

---
 src/SDK/Trace/TracerProvider.php            | 41 ++++++++++++++++++++-
 tests/Unit/SDK/Trace/TracerProviderTest.php | 13 +++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/src/SDK/Trace/TracerProvider.php b/src/SDK/Trace/TracerProvider.php
index ddf1c0a5e..37ec398d1 100644
--- a/src/SDK/Trace/TracerProvider.php
+++ b/src/SDK/Trace/TracerProvider.php
@@ -14,11 +14,15 @@
 use OpenTelemetry\SDK\Trace\Sampler\AlwaysOnSampler;
 use OpenTelemetry\SDK\Trace\Sampler\ParentBased;
 use function register_shutdown_function;
+use function spl_object_id;
+use WeakReference;
 
 final class TracerProvider implements API\TracerProviderInterface
 {
     public const DEFAULT_TRACER_NAME = 'io.opentelemetry.contrib.php';
 
+    /** @var array<int, WeakReference<self>>|null */
+    private static ?array $tracerProviders = null;
     private static ?API\TracerInterface $defaultTracer = null;
 
     /** @var array<string, API\TracerInterface> */
@@ -53,7 +57,7 @@ public function __construct(
             $spanProcessors
         );
 
-        register_shutdown_function([$this, 'shutdown']);
+        self::registerShutdownFunction($this);
     }
 
     public function forceFlush(): ?bool
@@ -116,6 +120,41 @@ public function shutdown(): bool
             return true;
         }
 
+        self::unregisterShutdownFunction($this);
+
         return $this->tracerSharedState->shutdown();
     }
+
+    public function __destruct()
+    {
+        $this->shutdown();
+    }
+
+    private static function registerShutdownFunction(TracerProvider $tracerProvider): void
+    {
+        if (self::$tracerProviders === null) {
+            register_shutdown_function(static function (): void {
+                $tracerProviders = self::$tracerProviders;
+                self::$tracerProviders = null;
+
+                // Push tracer provider shutdown to end of queue
+                // @phan-suppress-next-line PhanTypeMismatchArgumentInternal
+                register_shutdown_function(static function (array $tracerProviders): void {
+                    foreach ($tracerProviders as $reference) {
+                        if ($tracerProvider = $reference->get()) {
+                            $tracerProvider->shutdown();
+                        }
+                    }
+                }, $tracerProviders);
+            });
+        }
+
+        self::$tracerProviders[spl_object_id($tracerProvider)] = WeakReference::create($tracerProvider);
+    }
+
+    private static function unregisterShutdownFunction(TracerProvider $tracerProvider): void
+    {
+        /** @psalm-suppress PossiblyNullArrayAccess */
+        unset(self::$tracerProviders[spl_object_id($tracerProvider)]);
+    }
 }
diff --git a/tests/Unit/SDK/Trace/TracerProviderTest.php b/tests/Unit/SDK/Trace/TracerProviderTest.php
index caec2789d..7636ee4ab 100644
--- a/tests/Unit/SDK/Trace/TracerProviderTest.php
+++ b/tests/Unit/SDK/Trace/TracerProviderTest.php
@@ -10,6 +10,7 @@
 use OpenTelemetry\SDK\Trace\Tracer;
 use OpenTelemetry\SDK\Trace\TracerProvider;
 use PHPUnit\Framework\TestCase;
+use WeakReference;
 
 /**
  * @coversDefaultClass \OpenTelemetry\SDK\Trace\TracerProvider
@@ -174,4 +175,16 @@ public function test_get_tracer_returns_noop_tracer_after_shutdown(): void
             $provider->getTracer('foo')
         );
     }
+
+    /**
+     * @coversNothing
+     */
+    public function test_tracer_register_shutdown_function_does_not_leak_reference(): void
+    {
+        $provider = new TracerProvider();
+        $reference = WeakReference::create($provider);
+
+        $provider = null;
+        $this->assertTrue($reference->get() === null);
+    }
 }