diff --git a/common/lib/opentelemetry/common/utilities.rb b/common/lib/opentelemetry/common/utilities.rb index 8667a6fd94..df58722392 100644 --- a/common/lib/opentelemetry/common/utilities.rb +++ b/common/lib/opentelemetry/common/utilities.rb @@ -87,9 +87,12 @@ def truncate_attribute_value(value, limit) end # Disables tracing within the provided block. - def untraced - Context.with_value(UNTRACED_KEY, true) do |ctx, _| - yield ctx + def untraced(context = Context.current) + context = context.set_value(UNTRACED_KEY, true) + if block_given? + Context.with_current(context) { |ctx| yield ctx } + else + context end end diff --git a/common/test/opentelemetry/common/utilities_test.rb b/common/test/opentelemetry/common/utilities_test.rb index f931242fd8..555b87121d 100644 --- a/common/test/opentelemetry/common/utilities_test.rb +++ b/common/test/opentelemetry/common/utilities_test.rb @@ -26,6 +26,13 @@ def shutdown(timeout: nil); end common_utils.untraced {} assert_equal(false, common_utils.untraced?) end + + it 'supports non block format' do + token = OpenTelemetry::Context.attach(common_utils.untraced) + assert_equal(true, common_utils.untraced?) + OpenTelemetry::Context.detach(token) + assert_equal(false, common_utils.untraced?) + end end describe '#utf8_encode' do