From 0c57bd649ff7b86319151dac48abbe332ec4a64e Mon Sep 17 00:00:00 2001
From: Robert Laurin <robert.laurin@shopify.com>
Date: Thu, 10 Jun 2021 10:21:05 -0500
Subject: [PATCH] fix: use array to maintain context stack

---
 api/lib/opentelemetry/context.rb       | 33 +++++++++++++++-----------
 api/test/opentelemetry/context_test.rb | 32 +++++++++++++++++--------
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/api/lib/opentelemetry/context.rb b/api/lib/opentelemetry/context.rb
index 58e8f151e4..73b8f926a2 100644
--- a/api/lib/opentelemetry/context.rb
+++ b/api/lib/opentelemetry/context.rb
@@ -35,26 +35,27 @@ def current=(ctx)
         Thread.current[KEY] = ctx
       end
 
-      # Returns the previous context so that it can be restored
+      # Returns a token to be used with the matching call to detach
       #
       # @param [Context] context The new context
-      # @return [Context] prev The previous context
+      # @return [Integer] A token to be used when detaching
       def attach(context)
-        return current.parent if context == current
-
         prev = current
         self.current = context
-        context.parent = prev
+        stack.push(prev)
+        stack.size
       end
 
-      # Restores the current context to the context supplied or the parent context
-      # if no context is provided
+      # Restores the previous context, if a token is supplied it will
+      # be used to check if the call to detach is balanced with
+      # the original attach call. A warning is logged if the
+      # calls are unbalanced.
       #
-      # @param [Context] previous_context The previous context to restore
-      def detach(previous_context = nil)
-        OpenTelemetry.logger.warn 'Calls to detach should match corresponding calls to attach' if current.parent != previous_context
+      # @param [Integer] token The token provided by the matching call to attach
+      def detach(token = nil)
+        OpenTelemetry.logger.warn 'Calls to detach should match corresponding calls to attach' if token && token != stack.size
 
-        previous_context ||= current.parent || ROOT
+        previous_context = stack.pop || ROOT
         self.current = previous_context
       end
 
@@ -111,19 +112,23 @@ def value(key)
       end
 
       def clear
+        stack.clear
         self.current = ROOT
       end
 
       def empty
         new(EMPTY_ENTRIES)
       end
-    end
 
-    attr_accessor :parent
+      private
+
+      def stack
+        @stack ||= []
+      end
+    end
 
     def initialize(entries)
       @entries = entries.freeze
-      @parent = parent
     end
 
     # Returns the corresponding value (or nil) for key
diff --git a/api/test/opentelemetry/context_test.rb b/api/test/opentelemetry/context_test.rb
index c3dd18a783..1c6d2cace3 100644
--- a/api/test/opentelemetry/context_test.rb
+++ b/api/test/opentelemetry/context_test.rb
@@ -33,9 +33,9 @@
   end
 
   describe '.attach' do
-    it 'returns a reference to the previous context' do
-      previous_context = Context.attach(new_context)
-      _(previous_context).must_equal(Context::ROOT)
+    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
@@ -55,9 +55,19 @@
       _(Context.current[foo_key]).must_equal('c3')
     end
 
-    it 'attaching the current context does not create a circular reference' do
-      Context.attach(new_context)
+    it 'allows for attaching the same context multiple times' do
+      c1 = new_context
+      Context.attach(c1)
+      Context.attach(Context.current)
+      Context.attach(Context.current)
       Context.attach(Context.current)
+
+      Context.detach
+      _(Context.current).must_equal(c1)
+      Context.detach
+      _(Context.current).must_equal(c1)
+      Context.detach
+      _(Context.current).must_equal(c1)
       Context.detach
       _(Context.current).must_equal(Context::ROOT)
     end
@@ -75,10 +85,10 @@
     end
 
     it 'restores the context' do
-      prev = Context.attach(new_context)
+      c1_token = Context.attach(new_context)
       _(Context.current).must_equal(new_context)
 
-      Context.detach(prev)
+      Context.detach(c1_token)
       _(Context.current).must_equal(Context::ROOT)
 
       _(@log_stream.string).must_be_empty
@@ -86,10 +96,10 @@
 
     it 'warns mismatched detach calls' do
       c1 = new_context
-      Context.attach(c1)
+      c1_token = Context.attach(c1)
 
       c2 = Context.current.set_value(foo_key, 'c2')
-      c1_token = Context.attach(c2)
+      Context.attach(c2)
 
       c3 = Context.current.set_value(foo_key, 'c3')
       Context.attach(c3)
@@ -99,7 +109,7 @@
       _(@log_stream.string).must_match(/Calls to detach should match corresponding calls to attach/)
     end
 
-    it 'detaches to the parent if no context is provided' do
+    it 'detaches to the previous context' do
       c1 = new_context
       Context.attach(c1)
 
@@ -119,11 +129,13 @@
 
       Context.detach
       _(Context.current).must_equal(Context::ROOT)
+      _(@log_stream.string).must_be_empty
     end
 
     it 'detaching at the root leaves the root as the current context' do
       Context.detach
       _(Context.current).must_equal(Context::ROOT)
+      _(@log_stream.string).must_be_empty
     end
   end