From 1f9689c91712dd4f6e5b2d664e5100d5c4749152 Mon Sep 17 00:00:00 2001
From: Daniel Pepper <pepper.daniel@gmail.com>
Date: Tue, 28 Nov 2023 11:21:53 -0800
Subject: [PATCH] per action modes

---
 lib/network_resiliency.rb       | 39 +++++++++++++++++++++++++++------
 spec/network_resiliency_spec.rb | 28 +++++++++++++++++++++--
 2 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/lib/network_resiliency.rb b/lib/network_resiliency.rb
index 444f0f4..3fcb12b 100644
--- a/lib/network_resiliency.rb
+++ b/lib/network_resiliency.rb
@@ -15,6 +15,7 @@ module Adapter
     autoload :Postgres, "network_resiliency/adapter/postgres"
   end
 
+  ACTIONS = [ :connect, :request ].freeze
   ADAPTERS = [ :http, :faraday, :redis, :mysql, :postgres ].freeze
   MODE = [ :observe, :resilient ].freeze
   RESILIENCY_SIZE_THRESHOLD = 1_000
@@ -92,16 +93,40 @@ def timestamp
     Process.clock_gettime(Process::CLOCK_MONOTONIC) * 1_000
   end
 
-  def mode
-    @mode || :observe
+  def mode(action)
+    unless ACTIONS.include?(action)
+      raise ArgumentError, "invalid NetworkResiliency action: #{action}"
+    end
+
+    (@mode && @mode[action]) || :observe
   end
 
   def mode=(mode)
-    unless MODE.include?(mode)
-      raise ArgumentError, "invalid NetworkResiliency mode: #{mode}"
+    @mode = {}
+
+    if mode.is_a?(Hash)
+      invalid = mode.keys - ACTIONS
+
+      unless invalid.empty?
+        raise ArgumentError, "invalid actions for mode: #{invalid}"
+      end
+
+      mode.each do |action, mode|
+        unless MODE.include?(mode)
+          raise ArgumentError, "invalid NetworkResiliency mode for #{action}: #{mode}"
+        end
+
+        @mode[action] = mode
+      end
+    else
+      unless MODE.include?(mode)
+        raise ArgumentError, "invalid NetworkResiliency mode: #{mode}"
+      end
+
+      ACTIONS.each { |action| @mode[action] = mode }
     end
 
-    @mode = mode
+    @mode.freeze
   end
 
   def normalize_request(adapter, request = nil, **context, &block)
@@ -236,7 +261,7 @@ def ignore_destination?(adapter, action, destination)
   def timeouts_for(adapter:, action:, destination:, max: nil, units: :ms)
     default = [ max ]
 
-    return default if NetworkResiliency.mode == :observe
+    return default if NetworkResiliency.mode(action.to_sym) == :observe
 
     key = [ adapter, action, destination ].join(":")
     stats = StatsEngine.get(key)
@@ -282,7 +307,7 @@ def timeouts_for(adapter:, action:, destination:, max: nil, units: :ms)
     else
       timeouts << p99
 
-      # timeouts << p99 * 10 if NetworkResiliency.mode == :resolute
+      # timeouts << p99 * 10 if NetworkResiliency.mode(action) == :resolute
 
       # unbounded second attempt
       timeouts << nil
diff --git a/spec/network_resiliency_spec.rb b/spec/network_resiliency_spec.rb
index bddb58d..4b678f1 100644
--- a/spec/network_resiliency_spec.rb
+++ b/spec/network_resiliency_spec.rb
@@ -207,7 +207,7 @@ def expect_enabled
       end
 
       expect_enabled.to be false
-      expect(NetworkResiliency.mode).to be :resilient
+      expect(NetworkResiliency.mode(:connect)).to be :resilient
     end
 
     it "will start syncing" do
@@ -226,7 +226,7 @@ def expect_enabled
   end
 
   describe ".mode" do
-    subject { NetworkResiliency.mode }
+    subject { NetworkResiliency.mode(:connect) }
 
     it "defaults to observe" do
       is_expected.to be :observe
@@ -242,6 +242,10 @@ def expect_enabled
       expect {
         NetworkResiliency.mode = :foo
       }.to raise_error(ArgumentError)
+
+      expect {
+        NetworkResiliency.mode(:foo)
+      }.to raise_error(ArgumentError)
     end
 
     it "resets" do
@@ -250,6 +254,26 @@ def expect_enabled
 
       is_expected.to be :observe
     end
+
+    context "when actions set to different modes" do
+      before do
+        NetworkResiliency.mode = { connect: :resilient }
+      end
+
+      it { is_expected.to be :resilient }
+
+      it { expect(NetworkResiliency.mode(:request)).to be :observe }
+
+      it "fails fast on invalid input" do
+        expect {
+          NetworkResiliency.mode = { connect: :foo }
+        }.to raise_error(ArgumentError)
+
+        expect {
+          NetworkResiliency.mode = { conn: :observe }
+        }.to raise_error(ArgumentError)
+      end
+    end
   end
 
   describe ".record" do