fix: update context to match spec (#807)

* fix: update context to match spec

* fix: restore context value class method

* fix: use array to maintain context stack

* fix: misc Francis feedback

Co-authored-by: Francis Bogsanyi <francis.bogsanyi@shopify.com>

* fix: use error handler for logging out of sync detach calls

* fix: add detach specific error

* fix: relax description of return type for attach detach

Co-authored-by: Francis Bogsanyi <francis.bogsanyi@shopify.com>

* fix: Efficiency concerns

Co-authored-by: Francis Bogsanyi <francis.bogsanyi@shopify.com>

* fix: failing test from efficiency fixup

* fix: remove dead code

Co-authored-by: Francis Bogsanyi <francis.bogsanyi@shopify.com>

* fix: remove current setting in context clear

Co-authored-by: Francis Bogsanyi <francis.bogsanyi@shopify.com>

api/lib/opentelemetry/context.rb

@@ -10,8 +10,11 @@
 module OpenTelemetry
   # Manages context on a per-fiber basis
   class Context
-    KEY = :__opentelemetry_context__
     EMPTY_ENTRIES = {}.freeze
+    STACK_KEY = :__opentelemetry_context_storage__
+    private_constant :EMPTY_ENTRIES, :STACK_KEY
+
+    DetachError = Class.new(OpenTelemetry::Error)
 
     class << self
       # Returns a key used to index a value in a Context
@@ -26,14 +29,36 @@ def create_key(name)
       #
       # @return [Context]
       def current
-        Thread.current[KEY] ||= ROOT
+        stack.last || ROOT
       end
 
-      # Sets the current context
+      # Associates a Context with the caller's current Fiber. Every call to
+      # this operation should be paired with a corresponding call to detach.
       #
-      # @param [Context] ctx The context to be made active
-      def current=(ctx)
-        Thread.current[KEY] = ctx
+      # Returns a token to be used with the matching call to detach
+      #
+      # @param [Context] context The new context
+      # @return [Object] A token to be used when detaching
+      def attach(context)
+        s = stack
+        s.push(context)
+        s.size
+      end
+
+      # Restores the previous Context associated with the current Fiber.
+      # The supplied token is used to check if the call to detach is balanced
+      # with a corresponding attach call. A warning is logged if the
+      # calls are unbalanced.
+      #
+      # @param [Object] token The token provided by the matching call to attach
+      # @return [Boolean] True if the calls matched, false otherwise
+      def detach(token)
+        s = stack
+        calls_matched = (token == s.size)
+        OpenTelemetry.handle_error(exception: DetachError.new('calls to detach should match corresponding calls to attach.')) unless calls_matched
+
+        s.pop
+        calls_matched
       end
 
       # Executes a block with ctx as the current context. It restores
@@ -42,10 +67,10 @@ def current=(ctx)
       # @param [Context] ctx The context to be made active
       # @yield [context] Yields context to the block
       def with_current(ctx)
-        prev = ctx.attach
+        token = attach(ctx)
         yield ctx
       ensure
-        ctx.detach(prev)
+        detach(token)
       end
 
       # Execute a block in a new context with key set to value. Restores the
@@ -58,10 +83,10 @@ def with_current(ctx)
       #   the block
       def with_value(key, value)
         ctx = current.set_value(key, value)
-        prev = ctx.attach
+        token = attach(ctx)
         yield ctx, value
       ensure
-        ctx.detach(prev)
+        detach(token)
       end
 
       # Execute a block in a new context where its values are merged with the
@@ -75,10 +100,10 @@ def with_value(key, value)
       #   to the block
       def with_values(values)
         ctx = current.set_values(values)
-        prev = ctx.attach
+        token = attach(ctx)
         yield ctx, values
       ensure
-        ctx.detach(prev)
+        detach(token)
       end
 
       # Returns the value associated with key in the current context
@@ -89,16 +114,21 @@ def value(key)
       end
 
       def clear
-        self.current = ROOT
+        stack.clear
       end
 
       def empty
-        new(nil, EMPTY_ENTRIES)
+        new(EMPTY_ENTRIES)
+      end
+
+      private
+
+      def stack
+        Thread.current[STACK_KEY] ||= []
       end
     end
 
-    def initialize(parent, entries)
-      @parent = parent
+    def initialize(entries)
       @entries = entries.freeze
     end
 
@@ -120,7 +150,7 @@ def value(key)
     def set_value(key, value)
       new_entries = @entries.dup
       new_entries[key] = value
-      Context.new(self, new_entries)
+      Context.new(new_entries)
     end
 
     # Returns a new Context with the current context's entries merged with the
@@ -131,22 +161,7 @@ def set_value(key, value)
     # @param [Object] value Object to be stored under key
     # @return [Context]
     def set_values(values) # rubocop:disable Naming/AccessorMethodName:
-      Context.new(self, @entries.merge(values))
-    end
-
-    # @api private
-    def attach
-      prev = self.class.current
-      self.class.current = self
-      prev
-    end
-
-    # @api private
-    def detach(ctx_to_attach = nil)
-      OpenTelemetry.logger.warn 'Calls to detach should match corresponding calls to attach' if self.class.current != self
-
-      ctx_to_attach ||= @parent || ROOT
-      ctx_to_attach.attach
+      Context.new(@entries.merge(values))
     end
 
     ROOT = empty.freeze

api/test/opentelemetry/context_test.rb

@@ -11,21 +11,145 @@
 describe OpenTelemetry::Context do
   Context = OpenTelemetry::Context
 
-  after do
-    Context.clear
-  end
+  after { Context.clear }
 
   let(:foo_key) { Context.create_key('foo') }
   let(:bar_key) { Context.create_key('bar') }
   let(:baz_key) { Context.create_key('baz') }
   let(:new_context) { Context.empty.set_value(foo_key, 'bar') }
 
+  describe '.create_key' do
+    it 'returns a Context::Key' do
+      key = Context.create_key('testing')
+      _(key).must_be_instance_of(Context::Key)
+      _(key.name).must_equal('testing')
+    end
+  end
+
   describe '.current' do
     it 'defaults to the root context' do
       _(Context.current).must_equal(Context::ROOT)
     end
   end
 
+  describe '.attach' do
+    it 'returns a token to be used when detaching' do
+      c1_token = Context.attach(new_context)
+      _(c1_token).wont_be_nil
+    end
+
+    it 'sets the current context' do
+      c1 = new_context
+      Context.attach(c1)
+      _(Context.current).must_equal(c1)
+      _(Context.current[foo_key]).must_equal('bar')
+
+      c2 = Context.current.set_value(foo_key, 'c2')
+      Context.attach(c2)
+      _(Context.current).must_equal(c2)
+      _(Context.current[foo_key]).must_equal('c2')
+
+      c3 = Context.current.set_value(foo_key, 'c3')
+      Context.attach(c3)
+      _(Context.current).must_equal(c3)
+      _(Context.current[foo_key]).must_equal('c3')
+    end
+
+    it 'allows for attaching the same context multiple times' do
+      c1 = new_context
+      token0 = Context.attach(c1)
+      token1 = Context.attach(Context.current)
+      token2 = Context.attach(Context.current)
+      token3 = Context.attach(Context.current)
+
+      Context.detach(token0)
+      _(Context.current).must_equal(c1)
+      Context.detach(token1)
+      _(Context.current).must_equal(c1)
+      Context.detach(token2)
+      _(Context.current).must_equal(c1)
+      Context.detach(token3)
+      _(Context.current).must_equal(Context::ROOT)
+    end
+  end
+
+  describe '.detach' do
+    before do
+      @log_stream = StringIO.new
+      @_logger = OpenTelemetry.logger
+      OpenTelemetry.logger = ::Logger.new(@log_stream)
+    end
+
+    after do
+      OpenTelemetry.logger = @_logger
+    end
+
+    it 'restores the context' do
+      c1_token = Context.attach(new_context)
+      _(Context.current).must_equal(new_context)
+
+      Context.detach(c1_token)
+      _(Context.current).must_equal(Context::ROOT)
+
+      _(@log_stream.string).must_be_empty
+    end
+
+    it 'warns mismatched detach calls' do
+      c1 = new_context
+      c1_token = Context.attach(c1)
+
+      c2 = Context.current.set_value(foo_key, 'c2')
+      Context.attach(c2)
+
+      c3 = Context.current.set_value(foo_key, 'c3')
+      Context.attach(c3)
+
+      Context.detach(c1_token)
+
+      _(@log_stream.string).must_match(/OpenTelemetry error: calls to detach should match corresponding calls to attach/)
+    end
+
+    it 'detaches to the previous context' do
+      c1 = new_context
+      c1_token = Context.attach(c1)
+
+      c2 = Context.current.set_value(foo_key, 'c2')
+      c2_token = Context.attach(c2)
+
+      c3 = Context.current.set_value(foo_key, 'c3')
+      c3_token = Context.attach(c3)
+
+      _(Context.current).must_equal(c3)
+
+      Context.detach(c3_token)
+      _(Context.current).must_equal(c2)
+
+      Context.detach(c2_token)
+      _(Context.current).must_equal(c1)
+
+      Context.detach(c1_token)
+      _(Context.current).must_equal(Context::ROOT)
+      _(@log_stream.string).must_be_empty
+    end
+
+    it 'detaching with a junk token leaves the current context as root' do
+      Context.detach('junk')
+      _(Context.current).must_equal(Context::ROOT)
+      _(@log_stream.string).must_match(/OpenTelemetry error: calls to detach should match corresponding calls to attach/)
+    end
+
+    it 'with a raising error handler' do
+      OpenTelemetry.error_handler = lambda { |exception: nil, message: nil|
+        raise exception, "OpenTelemetry error: #{[message, exception&.message].compact.join(' - ')}"
+      }
+
+      _(-> { Context.detach('junk') }).must_raise(OpenTelemetry::Context::DetachError)
+
+    ensure
+      OpenTelemetry.error_handler = ->(exception: nil, message: nil) { OpenTelemetry.logger.error("OpenTelemetry error: #{[message, exception&.message].compact.join(' - ')}") }
+    end
+  end
+
   describe '.with_current' do
     it 'handles nested contexts' do
       c1 = new_context
@@ -41,7 +165,7 @@
 
     it 'resets context when an exception is raised' do
       c1 = new_context
-      Context.current = c1
+      Context.attach(c1)
 
       _(proc do
         c2 = Context.current.set_value(bar_key, 'baz')
@@ -85,13 +209,6 @@
     end
   end
 
-  describe '#value' do
-    it 'returns corresponding value for key' do
-      ctx = new_context
-      _(ctx.value(foo_key)).must_equal('bar')
-    end
-  end
-
   describe '.with_values' do
     it 'executes block within new context' do
       orig_ctx = Context.current
@@ -119,6 +236,44 @@
     end
   end
 
+  describe '.value' do
+    it 'returns the value from the current context' do
+      Context.attach(new_context)
+      _(Context.value(foo_key)).must_equal('bar')
+
+      c2 = Context.current.set_value(bar_key, 'baz')
+      Context.attach(c2)
+      _(Context.value(bar_key)).must_equal('baz')
+    end
+  end
+
+  describe '.clear' do
+    it 'clears the context' do
+      Context.attach(new_context)
+      _(Context.current).must_equal(new_context)
+
+      Context.clear
+
+      _(Context.current).must_equal(Context::ROOT)
+    end
+  end
+
+  describe '#value' do
+    it 'returns corresponding value for key' do
+      ctx = new_context
+      _(ctx.value(foo_key)).must_equal('bar')
+    end
+  end
+
+  describe '#set_value' do
+    it 'returns new context with entry' do
+      c1 = Context.current
+      c2 = c1.set_value(foo_key, 'bar')
+      _(c1.value(foo_key)).must_be_nil
+      _(c2.value(foo_key)).must_equal('bar')
+    end
+  end
+
   describe '#set_values' do
     it 'assigns multiple values' do
       ctx = new_context
@@ -136,15 +291,6 @@
     end
   end
 
-  describe '#update' do
-    it 'returns new context with entry' do
-      c1 = Context.current
-      c2 = c1.set_value(foo_key, 'bar')
-      _(c1.value(foo_key)).must_be_nil
-      _(c2.value(foo_key)).must_equal('bar')
-    end
-  end
-
   describe 'threading' do
     it 'unwinds the stack on each thread' do
       ctx = new_context