Context API improvements

api/lib/opentelemetry/context.rb

@@ -40,9 +40,10 @@ def current=(ctx)
       # the previous context upon exiting.
       #
       # @param [Context] ctx The context to be made active
+      # @yield [context] Yields context to the block
       def with_current(ctx)
         prev = ctx.attach
-        yield
+        yield ctx
       ensure
         ctx.detach(prev)
       end
@@ -53,10 +54,12 @@ def with_current(ctx)
       # @param [String] key The lookup key
       # @param [Object] value The object stored under key
       # @param [Callable] Block to execute in a new context
+      # @yield [context, value] Yields the newly created context and value to
+      #   the block
       def with_value(key, value)
         ctx = current.set_value(key, value)
         prev = ctx.attach
-        yield value
+        yield ctx, value
       ensure
         ctx.detach(prev)
       end
@@ -68,10 +71,12 @@ def with_value(key, value)
       # @param [Hash] values Will be merged with values of the current context
       #  and returned in a new context
       # @param [Callable] Block to execute in a new context
+      # @yield [context, values] Yields the newly created context and values
+      #   to the block
       def with_values(values)
         ctx = current.set_values(values)
         prev = ctx.attach
-        yield values
+        yield ctx, values
       ensure
         ctx.detach(prev)
       end

api/lib/opentelemetry/trace/tracer.rb

@@ -13,8 +13,12 @@ class Tracer
 
       private_constant :EXTRACTED_SPAN_CONTEXT_KEY, :CURRENT_SPAN_KEY
 
-      def current_span
-        Context.value(CURRENT_SPAN_KEY) || Span::INVALID
+      # Returns the current span from the current or provided context
+      #
+      # @param [optional Context] context The context to lookup the current
+      #   {Span} from. Defaults to Context.current
+      def current_span(context = Context.current)
+        context.value(CURRENT_SPAN_KEY) || Span::INVALID
       end
 
       # Returns the the active span context from the given {Context}, or current
@@ -46,9 +50,11 @@ def active_span_context(context = nil)
       # On exit, the Span that was active before calling this method will be reactivated. If an
       # exception occurs during the execution of the provided block, it will be recorded on the
       # span and reraised.
+      # @yield [span, context] yields the newly created span and a context containing the
+      #   span to the block.
       def in_span(name, attributes: nil, links: nil, start_timestamp: nil, kind: nil, with_parent: nil, with_parent_context: nil)
         span = start_span(name, attributes: attributes, links: links, start_timestamp: start_timestamp, kind: kind, with_parent: with_parent, with_parent_context: with_parent_context)
-        with_span(span) { |s| yield s }
+        with_span(span) { |s, c| yield s, c }
       rescue Exception => e # rubocop:disable Lint/RescueException
         span.record_error(e)
         span.status = Status.new(Status::UNKNOWN_ERROR,
@@ -62,8 +68,11 @@ def in_span(name, attributes: nil, links: nil, start_timestamp: nil, kind: nil,
       # available implicitly.
       #
       # On exit, the Span that was active before calling this method will be reactivated.
+      #
+      # @param [Span] span the span to activate
+      # @yield [span, context] yields span and a context containing the span to the block.
       def with_span(span)
-        Context.with_value(CURRENT_SPAN_KEY, span) { |s| yield s }
+        Context.with_value(CURRENT_SPAN_KEY, span) { |c, s| yield s, c }
       end
 
       def start_root_span(name, attributes: nil, links: nil, start_timestamp: nil, kind: nil)

api/test/opentelemetry/context_test.rb

@@ -52,6 +52,13 @@
 
       _(Context.current).must_equal(c1)
     end
+
+    it 'yields the current context to the block' do
+      ctx = new_context
+      Context.with_current(ctx) do |c|
+        _(c).must_equal(ctx)
+      end
+    end
   end
 
   describe '.with_value' do
@@ -60,7 +67,7 @@
 
       block_called = false
 
-      Context.with_value(foo_key, 'bar') do |value|
+      Context.with_value(foo_key, 'bar') do |_, value|
         _(Context.current.value(foo_key)).must_equal('bar')
         _(value).must_equal('bar')
         block_called = true
@@ -69,6 +76,13 @@
       _(Context.current).must_equal(orig_ctx)
       _(block_called).must_equal(true)
     end
+
+    it 'yields the current context and value to the block' do
+      Context.with_value(foo_key, 'bar') do |c, v|
+        _(v).must_equal('bar')
+        _(c.value(foo_key)).must_equal('bar')
+      end
+    end
   end
 
   describe '#value' do
@@ -84,7 +98,7 @@
 
       block_called = false
 
-      Context.with_values(foo_key => 'bar', bar_key => 'baz') do |values|
+      Context.with_values(foo_key => 'bar', bar_key => 'baz') do |_, values|
         _(Context.current.value(foo_key)).must_equal('bar')
         _(Context.current.value(bar_key)).must_equal('baz')
         _(values).must_equal(foo_key => 'bar', bar_key => 'baz')
@@ -94,6 +108,15 @@
       _(Context.current).must_equal(orig_ctx)
       _(block_called).must_equal(true)
     end
+
+    it 'yields the current context and values to the block' do
+      values = { foo_key => 'bar', bar_key => 'baz' }
+      Context.with_values(values) do |c, v|
+        _(v).must_equal(values)
+        _(c.value(foo_key)).must_equal('bar')
+        _(c.value(bar_key)).must_equal('baz')
+      end
+    end
   end
 
   describe '#set_values' do

api/test/opentelemetry/trace/tracer_test.rb

@@ -35,6 +35,9 @@ def start_span(*)
       parent_span_context
     )
   end
+  let(:current_span_key) do
+    OpenTelemetry::Trace::Propagation::ContextKeys.current_span_key
+  end
 
   describe '#current_span' do
     let(:current_span) { tracer.start_span('current') }
@@ -50,6 +53,13 @@ def start_span(*)
         _(tracer.current_span).must_equal(wrapper_span)
       end
     end
+
+    it 'returns the current span from the provided context' do
+      span = tracer.start_span('a-span')
+      context = Context.empty.set_value(current_span_key, span)
+      _(tracer.current_span).wont_equal(span)
+      _(tracer.current_span(context)).must_equal(span)
+    end
   end
 
   describe '#active_span_context' do
@@ -94,6 +104,13 @@ def start_span(*)
       end
     end
 
+    it 'yields context containing span' do
+      tracer.in_span('wrapper') do |span, context|
+        _(context).must_equal(OpenTelemetry::Context.current)
+        _(context[current_span_key]).must_equal(span)
+      end
+    end
+
     it 'returns the result of the block' do
       result = tracer.in_span('wrapper') { 'my-result' }
       _(result).must_equal('my-result')
@@ -130,6 +147,15 @@ def start_span(*)
       end
     end
 
+    it 'yields context containing span' do
+      wrapper_span = tracer.start_span('wrapper')
+
+      tracer.with_span(wrapper_span) do |span, context|
+        _(context).must_equal(OpenTelemetry::Context.current)
+        _(context[current_span_key]).must_equal(span)
+      end
+    end
+
     it 'should reactive the span after the block' do
       outer = tracer.start_span('outer')
       inner = tracer.start_span('inner')