From 8610008c02d5b9c12f1334af7015a66d211ef2fe Mon Sep 17 00:00:00 2001
From: Kir Shatrov <kirill.shatrov@shopify.com>
Date: Tue, 29 Oct 2024 15:47:50 -0700
Subject: [PATCH] Allow to disable bulkheads in a given thread

---
 CHANGELOG.md        |  1 +
 lib/semian.rb       | 17 ++++++++++++++++-
 test/semian_test.rb | 15 +++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3e251bcc..ab2b663e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 * Change: semians will trigger on ER_PROXYSQL_MAX_CONN_TIMEOUT (#520)
 * Change: added support for dynamic semian_configurations to the Semian Redis adapter
+* Change: `Semian.disable_bulkheads_for_thread` disables bulkheads for the given thread
 
 # v0.21.3
 
diff --git a/lib/semian.rb b/lib/semian.rb
index b117e96f..a53033a8 100644
--- a/lib/semian.rb
+++ b/lib/semian.rb
@@ -275,6 +275,21 @@ def thread_safe=(thread_safe)
     @thread_safe = thread_safe
   end
 
+  THREAD_BULKHEAD_DISABLED_VAR = :semian_bulkheads_disabled
+  private_constant(:THREAD_BULKHEAD_DISABLED_VAR)
+
+  def bulkheads_disabled_in_thread?(thread)
+    thread.thread_variable_get(THREAD_BULKHEAD_DISABLED_VAR)
+  end
+
+  def disable_bulkheads_for_thread(thread)
+    old_value = thread.thread_variable_get(THREAD_BULKHEAD_DISABLED_VAR)
+    thread.thread_variable_set(THREAD_BULKHEAD_DISABLED_VAR, true)
+    yield
+  ensure
+    thread.thread_variable_set(THREAD_BULKHEAD_DISABLED_VAR, old_value)
+  end
+
   private
 
   def create_circuit_breaker(name, **options)
@@ -318,7 +333,7 @@ def implementation(**options)
   end
 
   def create_bulkhead(name, **options)
-    return if ENV.key?("SEMIAN_BULKHEAD_DISABLED")
+    return if ENV.key?("SEMIAN_BULKHEAD_DISABLED") || bulkheads_disabled_in_thread?(Thread.current)
     return unless options.fetch(:bulkhead, true)
 
     permissions = options[:permissions] || default_permissions
diff --git a/test/semian_test.rb b/test/semian_test.rb
index 0e5a519f..85af7a7a 100644
--- a/test/semian_test.rb
+++ b/test/semian_test.rb
@@ -171,6 +171,21 @@ def test_disabled_bulkhead_via_env_with_option_enabled
     ENV.delete("SEMIAN_BULKHEAD_DISABLED")
   end
 
+  def test_disabled_bulkhead_via_thread
+    Semian.disable_bulkheads_for_thread(Thread.current) do
+      resource = Semian.register(
+        :disabled_bulkhead_via_env,
+        bulkhead: true,
+        tickets: 1,
+        success_threshold: 1,
+        error_threshold: 1,
+        error_timeout: 1,
+      )
+
+      assert_nil(resource.bulkhead)
+    end
+  end
+
   def test_disabled_circuit_breaker
     resource = Semian.register(
       :disabled_circuit_breaker,