From 06b2390193cca3222ab2c0845353cf3dcd934a5d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matyas?= <github@higher.lv>
Date: Sat, 4 May 2024 12:02:48 +0200
Subject: [PATCH 1/2] Implement local contexts

---
 lib/honeybadger/agent.rb           |  4 +--
 lib/honeybadger/context_manager.rb | 33 ++++++++++++++++++-----
 spec/unit/honeybadger_spec.rb      | 43 ++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/lib/honeybadger/agent.rb b/lib/honeybadger/agent.rb
index 86c21251d..011ac432f 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 73379b2d0..46cc393e5 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 nil unless @global_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 d85e189d0..1c6bedf26 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

From 4dd4459338cc5637b08318d1fc162077b9d0b031 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matyas?= <github@higher.lv>
Date: Wed, 8 May 2024 21:15:59 +0200
Subject: [PATCH 2/2] Add extra guards in case @local_context is nil

---
 lib/honeybadger/context_manager.rb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/honeybadger/context_manager.rb b/lib/honeybadger/context_manager.rb
index 46cc393e5..37a1d76e3 100644
--- a/lib/honeybadger/context_manager.rb
+++ b/lib/honeybadger/context_manager.rb
@@ -40,14 +40,14 @@ def set_context(hash, &block)
         begin
           yield
         ensure
-          @mutex.synchronize { @local_context.pop }
+          @mutex.synchronize { @local_context&.pop }
         end
       end
     end
 
     def get_context
       @mutex.synchronize do
-        return nil unless @global_context
+        return @global_context unless @local_context
 
         @global_context.merge(@local_context.inject({}, :merge))
       end