circuit_store.rb in httpx-1.0.0

- old
+ new

@@ -1,7 +1,9 @@
 # frozen_string_literal: true
 
+require "mutex_m"
+
 module HTTPX::Plugins::CircuitBreaker
   using HTTPX::URIExtensions
 
   class CircuitStore
     def initialize(options)
@@ -11,31 +13,40 @@
           options.circuit_breaker_reset_attempts_in,
           options.circuit_breaker_break_in,
           options.circuit_breaker_half_open_drip_rate
         )
       end
+      @circuits.extend(Mutex_m)
     end
 
     def try_open(uri, response)
-      circuit = get_circuit_for_uri(uri)
+      circuit = @circuits.synchronize { get_circuit_for_uri(uri) }
 
       circuit.try_open(response)
     end
 
+    def try_close(uri)
+      circuit = @circuits.synchronize do
+        return unless @circuits.key?(uri.origin) || @circuits.key?(uri.to_s)
+
+        get_circuit_for_uri(uri)
+      end
+
+      circuit.try_close
+    end
+
     # if circuit is open, it'll respond with the stored response.
     # if not, nil.
     def try_respond(request)
-      circuit = get_circuit_for_uri(request.uri)
+      circuit = @circuits.synchronize { get_circuit_for_uri(request.uri) }
 
       circuit.respond
     end
 
     private
 
     def get_circuit_for_uri(uri)
-      uri = URI(uri)
-
-      if @circuits.key?(uri.origin)
+      if uri.respond_to?(:origin) && @circuits.key?(uri.origin)
         @circuits[uri.origin]
       else
         @circuits[uri.to_s]
       end
     end