Improve fiber bound context storage (#1270)

* Improve fiber bound context storage

Now works properly in fibers once initial context is attached.

* Change `Context::storage()` return type to `ContextStorageInterface`

`ExecutionContextAwareInterface` should not be relevant for end-users / it was mainly exposed for the FFI fiber handler; calling any of its method with enabled fiber handler would have broken the storage.
Swoole context storage README creates a new storage instead of wrapping `Context::storage()`: `Context::setStorage(new SwooleContextStorage(new ContextStorage()));`.

* Add BC layer for execution context aware fiber storage

* Fix BC layer inactive execution context detection for {main}

Author: Nevay

Context.php

@@ -44,7 +44,7 @@ public static function setStorage(ContextStorageInterface&ExecutionContextAwareI
     public static function storage(): ContextStorageInterface&ExecutionContextAwareInterface
     {
         /** @psalm-suppress RedundantPropertyInitializationCheck */
-        return self::$storage ??= new ContextStorage();
+        return self::$storage ??= new FiberBoundContextStorageExecutionAwareBC();
     }
 
     /**

ContextStorage.php

@@ -7,9 +7,9 @@
 /**
  * @internal
  */
-final class ContextStorage implements ContextStorageInterface, ExecutionContextAwareInterface
+final class ContextStorage implements ContextStorageInterface, ContextStorageHeadAware, ExecutionContextAwareInterface
 {
-    public ContextStorageHead $current;
+    private ContextStorageHead $current;
     private ContextStorageHead $main;
     /** @var array<int|string, ContextStorageHead> */
     private array $forks = [];
@@ -34,6 +34,11 @@ public function destroy(int|string $id): void
         unset($this->forks[$id]);
     }
 
+    public function head(): ContextStorageHead
+    {
+        return $this->current;
+    }
+
     public function scope(): ?ContextStorageScopeInterface
     {
         return ($this->current->node->head ?? null) === $this->current

ContextStorageHead.php

@@ -11,7 +11,7 @@ final class ContextStorageHead
 {
     public ?ContextStorageNode $node = null;
 
-    public function __construct(public ContextStorage $storage)
+    public function __construct(public ContextStorageHeadAware $storage)
     {
     }
 }

ContextStorageHeadAware.php

@@ -0,0 +1,13 @@
+<?php
+
+declare(strict_types=1);
+
+namespace OpenTelemetry\Context;
+
+/**
+ * @internal
+ */
+interface ContextStorageHeadAware
+{
+    public function head(): ?ContextStorageHead;
+}

ContextStorageNode.php

@@ -20,26 +20,22 @@ public function __construct(
     ) {
     }
 
-    public function offsetExists($offset): bool
+    public function offsetExists(mixed $offset): bool
     {
         return isset($this->localStorage[$offset]);
     }
 
-    /**
-     * @phan-suppress PhanUndeclaredClassAttribute
-     */
-    #[\ReturnTypeWillChange]
-    public function offsetGet($offset)
+    public function offsetGet(mixed $offset): mixed
     {
         return $this->localStorage[$offset];
     }
 
-    public function offsetSet($offset, $value): void
+    public function offsetSet(mixed $offset, mixed $value): void
     {
         $this->localStorage[$offset] = $value;
     }
 
-    public function offsetUnset($offset): void
+    public function offsetUnset(mixed $offset): void
     {
         unset($this->localStorage[$offset]);
     }
@@ -52,7 +48,7 @@ public function context(): ContextInterface
     public function detach(): int
     {
         $flags = 0;
-        if ($this->head !== $this->head->storage->current) {
+        if ($this->head !== $this->head->storage->head()) {
             $flags |= ScopeInterface::INACTIVE;
         }
 

FiberBoundContextStorage.php

@@ -1,77 +1,79 @@
 <?php
 
-/** @noinspection PhpElementIsNotAvailableInCurrentPhpVersionInspection */
-
 declare(strict_types=1);
 
 namespace OpenTelemetry\Context;
 
 use function assert;
-use function class_exists;
 use const E_USER_WARNING;
 use Fiber;
+use function spl_object_id;
+use function sprintf;
 use function trigger_error;
+use WeakMap;
 
 /**
  * @internal
- *
- * @phan-file-suppress PhanUndeclaredClassReference
- * @phan-file-suppress PhanUndeclaredClassMethod
  */
-final class FiberBoundContextStorage implements ContextStorageInterface, ExecutionContextAwareInterface
+final class FiberBoundContextStorage implements ContextStorageInterface, ContextStorageHeadAware
 {
-    public function __construct(private readonly ContextStorageInterface&ExecutionContextAwareInterface $storage)
-    {
-    }
-
-    public function fork(int|string $id): void
-    {
-        $this->storage->fork($id);
-    }
+    /** @var WeakMap<object, ContextStorageHead> */
+    private WeakMap $heads;
 
-    public function switch(int|string $id): void
+    public function __construct()
     {
-        $this->storage->switch($id);
+        $this->heads = new WeakMap();
+        $this->heads[$this] = new ContextStorageHead($this);
     }
 
-    public function destroy(int|string $id): void
+    public function head(): ?ContextStorageHead
     {
-        $this->storage->destroy($id);
+        return $this->heads[Fiber::getCurrent() ?? $this] ?? null;
     }
 
     public function scope(): ?ContextStorageScopeInterface
     {
-        $this->checkFiberMismatch();
+        $head = $this->heads[Fiber::getCurrent() ?? $this] ?? null;
+
+        if (!$head?->node && Fiber::getCurrent()) {
+            self::triggerNotInitializedFiberContextWarning();
 
-        if (($scope = $this->storage->scope()) === null) {
             return null;
         }
 
-        return new FiberBoundContextStorageScope($scope);
+        // Starts with empty head instead of cloned parent -> no need to check for head mismatch
+        return $head->node;
     }
 
     public function current(): ContextInterface
     {
-        $this->checkFiberMismatch();
+        $head = $this->heads[Fiber::getCurrent() ?? $this] ?? null;
+
+        if (!$head?->node && Fiber::getCurrent()) {
+            self::triggerNotInitializedFiberContextWarning();
+
+            // Fallback to {main} to preserve BC
+            $head = $this->heads[$this];
+        }
 
-        return $this->storage->current();
+        return $head->node->context ?? Context::getRoot();
     }
 
     public function attach(ContextInterface $context): ContextStorageScopeInterface
     {
-        $scope = $this->storage->attach($context);
-        assert(class_exists(Fiber::class, false));
-        $scope[Fiber::class] = Fiber::getCurrent();
+        $head = $this->heads[Fiber::getCurrent() ?? $this] ??= new ContextStorageHead($this);
 
-        return new FiberBoundContextStorageScope($scope);
+        return $head->node = new ContextStorageNode($context, $head, $head->node);
     }
 
-    private function checkFiberMismatch(): void
+    private static function triggerNotInitializedFiberContextWarning(): void
     {
-        $scope = $this->storage->scope();
-        assert(class_exists(Fiber::class, false));
-        if ($scope && $scope[Fiber::class] !== Fiber::getCurrent()) {
-            trigger_error('Fiber context switching not supported', E_USER_WARNING);
-        }
+        $fiber = Fiber::getCurrent();
+        assert($fiber !== null);
+
+        trigger_error(sprintf(
+            'Access to not initialized OpenTelemetry context in fiber (id: %d), automatic forking not supported, must attach initial fiber context manually',
+            spl_object_id($fiber),
+        ), E_USER_WARNING);
     }
 }

FiberBoundContextStorageExecutionAwareBC.php

@@ -0,0 +1,74 @@
+<?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);
+            $head->storage = $this->bc;
+
+            /** @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);
+    }
+}

FiberBoundContextStorageScope.php

@@ -39,9 +39,9 @@ truthy value. Disabling is only recommended for applications using `exit` / `die
 
 ## Async applications
 
-### Fiber support
+### Fiber support - automatic context propagation to newly created fibers
 
-Requires `PHP >= 8.1`, an NTS build, `ext-ffi`, and setting the environment variable `OTEL_PHP_FIBERS_ENABLED` to a truthy value. Additionally `vendor/autoload.php` has to be preloaded for non-CLI SAPIs if [`ffi.enable`](https://www.php.net/manual/en/ffi.configuration.php#ini.ffi.enable) is set to `preload`.
+Requires an NTS build, `ext-ffi`, and setting the environment variable `OTEL_PHP_FIBERS_ENABLED` to a truthy value. Additionally `vendor/autoload.php` has to be preloaded for non-CLI SAPIs if [`ffi.enable`](https://www.php.net/manual/en/ffi.configuration.php#ini.ffi.enable) is set to `preload`.
 
 ### Event loops
 

README.md

@@ -48,7 +48,7 @@ public static function init(): bool
 
         try {
             $fibers = FFI::scope('OTEL_ZEND_OBSERVER_FIBER');
-        } catch (FFI\Exception $e) {
+        } catch (FFI\Exception) {
             try {
                 $fibers = FFI::load(__DIR__ . '/fiber/zend_observer_fiber.h');
             } catch (FFI\Exception $e) {
@@ -58,9 +58,12 @@ public static function init(): bool
             }
         }
 
-        $fibers->zend_observer_fiber_init_register(static fn (int $initializing) => Context::storage()->fork($initializing)); //@phpstan-ignore-line
-        $fibers->zend_observer_fiber_switch_register(static fn (int $from, int $to) => Context::storage()->switch($to)); //@phpstan-ignore-line
-        $fibers->zend_observer_fiber_destroy_register(static fn (int $destroying) => Context::storage()->destroy($destroying)); //@phpstan-ignore-line
+        $storage = new ContextStorage();
+        $fibers->zend_observer_fiber_init_register(static fn (int $initializing) => $storage->fork($initializing)); //@phpstan-ignore-line
+        $fibers->zend_observer_fiber_switch_register(static fn (int $from, int $to) => $storage->switch($to)); //@phpstan-ignore-line
+        $fibers->zend_observer_fiber_destroy_register(static fn (int $destroying) => $storage->destroy($destroying)); //@phpstan-ignore-line
+
+        Context::setStorage($storage);
 
         return true;
     }