diff --git a/lib/honeybadger/agent.rb b/lib/honeybadger/agent.rb index 86c21251..011ac432 100644 --- a/lib/honeybadger/agent.rb +++ b/lib/honeybadger/agent.rb @@ -259,8 +259,8 @@ def track_deployment(environment: nil, revision: nil, local_username: nil, repos # (optional). # # @return [self] so that method calls can be chained. - def context(context = nil) - context_manager.set_context(context) unless context.nil? + def context(context = nil, &block) + context_manager.set_context(context, &block) unless context.nil? self end diff --git a/lib/honeybadger/context_manager.rb b/lib/honeybadger/context_manager.rb index 73379b2d..37a1d76e 100644 --- a/lib/honeybadger/context_manager.rb +++ b/lib/honeybadger/context_manager.rb @@ -21,15 +21,36 @@ def clear! # Internal helpers - def set_context(hash) + def set_context(hash, &block) + local = block_given? @mutex.synchronize do - @context ||= {} - @context.update(Context(hash)) + @global_context ||= {} + @local_context ||= [] + + new_context = Context(hash) + + if local + @local_context << new_context + else + @global_context.update(new_context) + end + end + + if local + begin + yield + ensure + @mutex.synchronize { @local_context&.pop } + end end end def get_context - @mutex.synchronize { @context } + @mutex.synchronize do + return @global_context unless @local_context + + @global_context.merge(@local_context.inject({}, :merge)) + end end def set_rack_env(env) @@ -46,10 +67,10 @@ def get_rack_env def _initialize @mutex.synchronize do - @context = nil + @global_context = nil + @local_context = nil @rack_env = nil end end - end end diff --git a/spec/unit/honeybadger_spec.rb b/spec/unit/honeybadger_spec.rb index d85e189d..1c6bedf2 100644 --- a/spec/unit/honeybadger_spec.rb +++ b/spec/unit/honeybadger_spec.rb @@ -50,6 +50,7 @@ let(:c) { {foo: :bar} } before { described_class.context(c) } + after { described_class.context.clear! } it "sets the context" do described_class.context(c) @@ -67,6 +68,48 @@ it "clears the context" do expect { described_class.context.clear! }.to change { described_class.get_context }.from(c).to(nil) end + + context 'with local context' do + it 'merges local context' do + allow(described_class).to receive(:get_context).and_call_original + + described_class.context({ bar: :baz }) do + described_class.context({ bar: :qux }) do + expect(described_class.get_context).to eq({ foo: :bar, bar: :qux }) + end + expect(described_class.get_context).to eq({ foo: :bar, bar: :baz }) + end + + expect(described_class).to have_received(:get_context).at_least(:twice) + end + + it 'clears local context' do + allow(described_class).to receive(:get_context).and_call_original + + described_class.context({ bar: :baz }) do + expect(described_class.get_context).to eq({ foo: :bar, bar: :baz }) + end + expect(described_class.get_context).to eq({ foo: :bar }) + + expect(described_class).to have_received(:get_context).at_least(:twice) + end + + it 'tracks local context correctly' do + allow(described_class).to receive(:get_context).and_call_original + + begin + described_class.context({ bar: :qux }) do + described_class.context({ baz: :qux }) + + raise 'exception' + end + rescue StandardError + # noop + end + + expect(described_class.get_context).to eq({ foo: :bar, baz: :qux }) + end + end end describe "#notify" do