Use array lookup instead of recursion in Context (#821)

* Use array instead of recursion in `Context`

* Move span and baggage context key to context package

* Add additional context tests

Author: Nevay

Context.php

@@ -4,94 +4,68 @@
 
 namespace OpenTelemetry\Context;
 
-use function assert;
+use function spl_object_id;
 
 /**
- * @see https://github.com/open-telemetry/opentelemetry-specification/blob/v1.6.1/specification/context/context.md#overview
+ * @see https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/context/README.md#context
  */
 final class Context
 {
-    /**
-     * @var ContextStorageInterface&ExecutionContextAwareInterface
-     */
+    /** @var ContextStorageInterface&ExecutionContextAwareInterface */
     private static ContextStorageInterface $storage;
-    private static ?self $root = null;
-    private ?ContextKey $key;
-    /**
-     * @var mixed
-     */
-    private $value;
-    private ?self $parent;
 
-    /**
-     * This is a general purpose read-only key-value store. Read-only in the sense that adding a new value does not
-     * mutate the existing context, but returns a new Context which has the new value added.
-     *
-     * In practical terms, this is implemented as a linked list of Context instances, with each one holding a reference
-     * to the key object, the value that corresponds to the key, and an optional reference to the parent Context
-     * (i.e. the next link in the linked list chain)
-     *
-     * @param ContextKey|null $key The key object. Should only be null when creating an "empty" context
-     * @param mixed $value
-     * @param Context|null $parent Reference to the parent object
-     */
-    private function __construct(?ContextKey $key=null, $value=null, ?self $parent=null)
+    // Optimization for spans to avoid copying the context array.
+    private static ContextKey $spanContextKey;
+    private ?object $span = null;
+    /** @var array<int, mixed> */
+    private array $context = [];
+    /** @var array<int, ContextKey> */
+    private array $contextKeys = [];
+
+    private function __construct()
     {
-        $this->key = $key;
-        $this->value = $value;
-        $this->parent = $parent;
+        self::$spanContextKey = ContextKeys::span();
     }
 
-    /**
-     * Static version of get()
-     * This is primarily useful when the caller doesn't already have a reference to the Context that they want to mutate.
-     * This will operate on the "current" global context in that scenario.
-     *
-     * There are two ways to call this function:
-     * 1) With a $context value:
-     *    Context::getValue($key, $context) is functionally equivalent to $context->get($key)
-     * 2) Without a $context value:
-     *    This will fetch the "current" Context if one exists or create one if not, then attempt to get the value from it.
-     *
-     * @param ContextKey $key
-     * @param Context|null $context
-     *
-     * @return mixed
-     */
-    public static function getValue(ContextKey $key, ?self $context=null)
+    public static function createKey(string $key): ContextKey
     {
-        $context ??= self::getCurrent();
-
-        return $context->get($key);
+        return new ContextKey($key);
     }
 
     /**
-     * This adds a key/value pair to this Context. We do this by instantiating a new Context instance with the key/value and pass
-     * a reference to $this as the "parent" creating the linked list chain.
-     *
-     * @param ContextKey $key
-     * @param mixed $value
+     * @internal
      *
-     * @return Context a new Context containing the key/value
+     * @param ContextStorageInterface&ExecutionContextAwareInterface $storage
      */
-    public function with(ContextKey $key, $value): self
+    public static function setStorage(ContextStorageInterface $storage): void
     {
-        return new self($key, $value, $this);
+        self::$storage = $storage;
     }
 
     /**
-     * @todo: Implement this on the API side
+     * @return ContextStorageInterface&ExecutionContextAwareInterface
      */
-    public function withContextValue(ImplicitContextKeyedInterface $value): self
+    public static function storage(): ContextStorageInterface
     {
-        return $value->storeInContext($this);
+        /** @psalm-suppress RedundantPropertyInitializationCheck */
+        return self::$storage ??= new ContextStorage();
     }
 
     /**
-     * Makes `$this` the currently active {@see Context}.
-     *
-     * @todo: Implement this on the API side
+     * @internal
      */
+    public static function getRoot(): Context
+    {
+        static $empty;
+
+        return $empty ??= new self();
+    }
+
+    public static function getCurrent(): Context
+    {
+        return self::storage()->current();
+    }
+
     public function activate(): ScopeInterface
     {
         $scope = self::storage()->attach($this);
@@ -101,58 +75,45 @@ public function activate(): ScopeInterface
         return $scope;
     }
 
-    /**
-     * Fetch a value from the Context given a key value.
-     *
-     * @return mixed
-     */
-    public function get(ContextKey $key)
+    public function withContextValue(ImplicitContextKeyedInterface $value): Context
     {
-        for ($context = $this; $context; $context = $context->parent) {
-            if ($context->key === $key) {
-                return $context->value;
-            }
-        }
-
-        return null;
+        return $value->storeInContext($this);
     }
 
-    /**
-     * @internal
-     *
-     * @param ContextStorageInterface&ExecutionContextAwareInterface $storage
-     */
-    public static function setStorage(ContextStorageInterface $storage): void
+    public function with(ContextKey $key, $value): self
     {
-        self::$storage = $storage;
-    }
+        if ($this->get($key) === $value) {
+            return $this;
+        }
 
-    /**
-     * @return ContextStorageInterface&ExecutionContextAwareInterface
-     */
-    public static function storage(): ContextStorageInterface
-    {
-        /** @psalm-suppress RedundantPropertyInitializationCheck */
-        return self::$storage ??= new ContextStorage();
-    }
+        $self = clone $this;
 
-    /**
-     * @param non-empty-string $key
-     *
-     * @see https://github.com/open-telemetry/opentelemetry-specification/blob/v1.6.1/specification/context/context.md#create-a-key
-     */
-    public static function createKey(string $key): ContextKey
-    {
-        return new ContextKey($key);
-    }
+        if ($key === self::$spanContextKey) {
+            $self->span = $value; // @phan-suppress-current-line PhanTypeMismatchPropertyReal
 
-    public static function getCurrent(): self
-    {
-        return self::storage()->current();
+            return $self;
+        }
+
+        $id = spl_object_id($key);
+        if ($value !== null) {
+            $self->context[$id] = $value;
+            $self->contextKeys[$id] ??= $key;
+        } else {
+            unset(
+                $self->context[$id],
+                $self->contextKeys[$id],
+            );
+        }
+
+        return $self;
     }
 
-    public static function getRoot(): self
+    public function get(ContextKey $key)
     {
-        return self::$root ??= new self();
+        if ($key === self::$spanContextKey) {
+            return $this->span;
+        }
+
+        return $this->context[spl_object_id($key)] ?? null;
     }
 }

ContextKeys.php

@@ -0,0 +1,25 @@
+<?php
+
+declare(strict_types=1);
+
+namespace OpenTelemetry\Context;
+
+/**
+ * @psalm-internal OpenTelemetry
+ */
+final class ContextKeys
+{
+    public static function span(): ContextKey
+    {
+        static $instance;
+
+        return $instance ??= Context::createKey('opentelemetry-trace-span-key');
+    }
+
+    public static function baggage(): ContextKey
+    {
+        static $instance;
+
+        return $instance ??= Context::createKey('opentelemetry-trace-baggage-key');
+    }
+}