From 32971271e81464364b383ecec861c4c566e850ea Mon Sep 17 00:00:00 2001
From: Tobias Bachert <git@b-privat.de>
Date: Fri, 5 Apr 2024 18:56:01 +0200
Subject: [PATCH] Add BC layer for execution context aware fiber storage

---
 src/Context/Context.php                       |  8 +--
 ...berBoundContextStorageExecutionAwareBC.php | 72 +++++++++++++++++++
 ...oundContextStorageExecutionAwareBCTest.php | 58 +++++++++++++++
 3 files changed, 134 insertions(+), 4 deletions(-)
 create mode 100644 src/Context/FiberBoundContextStorageExecutionAwareBC.php
 create mode 100644 tests/Unit/Context/FiberBoundContextStorageExecutionAwareBCTest.php

diff --git a/src/Context/Context.php b/src/Context/Context.php
index 87529543a..e64e533a4 100644
--- a/src/Context/Context.php
+++ b/src/Context/Context.php
@@ -16,7 +16,7 @@ final class Context implements ContextInterface
 {
     private const OTEL_PHP_DEBUG_SCOPES_DISABLED = 'OTEL_PHP_DEBUG_SCOPES_DISABLED';
 
-    private static ContextStorageInterface $storage;
+    private static ContextStorageInterface&ExecutionContextAwareInterface $storage;
 
     // Optimization for spans to avoid copying the context array.
     private static ContextKeyInterface $spanContextKey;
@@ -36,15 +36,15 @@ public static function createKey(string $key): ContextKeyInterface
         return new ContextKey($key);
     }
 
-    public static function setStorage(ContextStorageInterface $storage): void
+    public static function setStorage(ContextStorageInterface&ExecutionContextAwareInterface $storage): void
     {
         self::$storage = $storage;
     }
 
-    public static function storage(): ContextStorageInterface
+    public static function storage(): ContextStorageInterface&ExecutionContextAwareInterface
     {
         /** @psalm-suppress RedundantPropertyInitializationCheck */
-        return self::$storage ??= new FiberBoundContextStorage();
+        return self::$storage ??= new FiberBoundContextStorageExecutionAwareBC();
     }
 
     /**
diff --git a/src/Context/FiberBoundContextStorageExecutionAwareBC.php b/src/Context/FiberBoundContextStorageExecutionAwareBC.php
new file mode 100644
index 000000000..adab2f9eb
--- /dev/null
+++ b/src/Context/FiberBoundContextStorageExecutionAwareBC.php
@@ -0,0 +1,72 @@
+<?php
+
+declare(strict_types=1);
+
+namespace OpenTelemetry\Context;
+
+/**
+ * @internal
+ */
+final class FiberBoundContextStorageExecutionAwareBC implements ContextStorageInterface, ExecutionContextAwareInterface
+{
+    private readonly FiberBoundContextStorage $storage;
+    private ?ContextStorage $bc = null;
+
+    public function __construct()
+    {
+        $this->storage = new FiberBoundContextStorage();
+    }
+
+    public function fork(int|string $id): void
+    {
+        $this->bcStorage()->fork($id);
+    }
+
+    public function switch(int|string $id): void
+    {
+        $this->bcStorage()->switch($id);
+    }
+
+    public function destroy(int|string $id): void
+    {
+        $this->bcStorage()->destroy($id);
+    }
+
+    private function bcStorage(): ContextStorage
+    {
+        if ($this->bc === null) {
+            $this->bc = new ContextStorage();
+
+            // Copy head into $this->bc storage to preserve already attached scopes
+            /** @psalm-suppress PossiblyNullFunctionCall */
+            $head = (static fn ($storage) => $storage->heads[$storage])
+                ->bindTo(null, FiberBoundContextStorage::class)($this->storage);
+            /** @psalm-suppress PossiblyNullFunctionCall */
+            (static fn ($storage) => $storage->current = $storage->main = $head)
+                ->bindTo(null, ContextStorage::class)($this->bc);
+        }
+
+        return $this->bc;
+    }
+
+    public function scope(): ?ContextStorageScopeInterface
+    {
+        return $this->bc
+            ? $this->bc->scope()
+            : $this->storage->scope();
+    }
+
+    public function current(): ContextInterface
+    {
+        return $this->bc
+            ? $this->bc->current()
+            : $this->storage->current();
+    }
+
+    public function attach(ContextInterface $context): ContextStorageScopeInterface
+    {
+        return $this->bc
+            ? $this->bc->attach($context)
+            : $this->storage->attach($context);
+    }
+}
diff --git a/tests/Unit/Context/FiberBoundContextStorageExecutionAwareBCTest.php b/tests/Unit/Context/FiberBoundContextStorageExecutionAwareBCTest.php
new file mode 100644
index 000000000..bb5cdf137
--- /dev/null
+++ b/tests/Unit/Context/FiberBoundContextStorageExecutionAwareBCTest.php
@@ -0,0 +1,58 @@
+<?php
+
+declare(strict_types=1);
+
+namespace OpenTelemetry\Context;
+
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\TestCase;
+
+#[CoversClass(FiberBoundContextStorageExecutionAwareBC::class)]
+final class FiberBoundContextStorageExecutionAwareBCTest extends TestCase
+{
+    public function test_retains_scope_after_bc_switch(): void
+    {
+        $storage = new FiberBoundContextStorageExecutionAwareBC();
+
+        $storage->attach($storage->current());
+
+        $scope = $storage->scope();
+
+        $storage->fork(1);
+
+        $this->assertSame($scope, $storage->scope());
+    }
+
+    public function test_storage_switch_switches_context(): void
+    {
+        $storage = new FiberBoundContextStorageExecutionAwareBC();
+        $main = $storage->current();
+        $fork = $storage->current()->with(Context::createKey('-'), 42);
+
+        $scopeMain = $storage->attach($main);
+
+        // Coroutine start
+        $storage->fork(1);
+        $storage->switch(1);
+        $this->assertSame($main, $storage->current());
+
+        $scopeFork = $storage->attach($fork);
+        $this->assertSame($fork, $storage->current());
+
+        // Coroutine suspend
+        $storage->switch(0);
+        $this->assertSame($main, $storage->current());
+
+        // Coroutine resume
+        $storage->switch(1);
+        $this->assertSame($fork, $storage->current());
+
+        $scopeFork->detach();
+
+        // Coroutine return
+        $storage->switch(0);
+        $storage->destroy(1);
+
+        $scopeMain->detach();
+    }
+}