pinot-client 1.21.0 → 1.23.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f91a5f7f465b546a3a49c78d2076097eac9b6af8fd9fbee7f8c1c45df06bea16
-  data.tar.gz: 0a04dd7945a401c01e107c95062b738ff987a23706338d7436e88ff8e50bec13
+  metadata.gz: 0752e62700790e556388c588d9a305049a5935fc4da0afdefeecab830a9dde0c
+  data.tar.gz: 4f5b1e86e84817febf245ec13498ec300ac68688c691eda58a2f4aafa79d219c
 SHA512:
-  metadata.gz: aeee0c948cca37c67a7b8bb103201ea4a4ad1dbe9e08476ca679efb10c4ea9951ff4ade65af6e8f7cc11e3c55b33ca66f969b0e4271fe97a8bd62bc7265c24aa
-  data.tar.gz: 8701f59b1adc42307e0774484a24d5c0e0cd33e8da378d9d2ccc65ad0240befcd094d58a136dc5031150fa687e88bbf87cbdde29e02d69e96462eedaf31dcbdc
+  metadata.gz: 684b1e6b39e4e1f5f84bed527de2b50f056d101b81ba7621549454d81cc08f49dd044ce2d3dfa5cbb8084075339c0b4c45a89f32c7cf185f1aa73e6c9e40ec95
+  data.tar.gz: 4ea1edf2abac766c358f0e99ed073a05e500d19fb77101475c96f6b82fa4d308a0025565bc1d8d55103da7c7af470654343f9b2b5b57366838e257853e493223

data/lib/pinot/circuit_breaker.rb

@@ -0,0 +1,115 @@
+module Pinot
+  # Per-broker circuit breaker. States: CLOSED (normal), OPEN (rejecting), HALF_OPEN (probing).
+  #
+  # failure_threshold  - consecutive failures before opening (default 5)
+  # open_timeout       - seconds to stay OPEN before moving to HALF_OPEN (default 30)
+  class CircuitBreaker
+    CLOSED    = :closed
+    OPEN      = :open
+    HALF_OPEN = :half_open
+
+    BrokerCircuitOpenError = Class.new(BrokerNotFoundError)
+
+    attr_reader :state, :failure_count
+
+    def initialize(failure_threshold: 5, open_timeout: 30)
+      @failure_threshold = failure_threshold
+      @open_timeout      = open_timeout
+      @mutex             = Mutex.new
+      @state             = CLOSED
+      @failure_count     = 0
+      @opened_at         = nil
+    end
+
+    # Call the block; record success/failure and enforce open-circuit rejection.
+    def call(broker_address)
+      @mutex.synchronize { check_state! }
+      begin
+        result = yield
+        @mutex.synchronize { on_success }
+        result
+      rescue BrokerUnavailableError, Errno::ECONNRESET, Errno::ECONNREFUSED,
+             Errno::ETIMEDOUT, Net::OpenTimeout, Net::ReadTimeout, Net::WriteTimeout => e
+        @mutex.synchronize { on_failure }
+        raise
+      end
+    end
+
+    def open?
+      @mutex.synchronize { @state == OPEN }
+    end
+
+    def reset
+      @mutex.synchronize do
+        @state         = CLOSED
+        @failure_count = 0
+        @opened_at     = nil
+      end
+    end
+
+    private
+
+    def check_state!
+      return if @state == CLOSED
+
+      if @state == OPEN
+        if elapsed_since_open >= @open_timeout
+          @state = HALF_OPEN
+        else
+          raise BrokerCircuitOpenError, "circuit open for broker (#{remaining_open_time.ceil}s remaining)"
+        end
+      end
+      # HALF_OPEN: allow the probe through
+    end
+
+    def on_success
+      @state         = CLOSED
+      @failure_count = 0
+      @opened_at     = nil
+    end
+
+    def on_failure
+      @failure_count += 1
+      if @state == HALF_OPEN || @failure_count >= @failure_threshold
+        @state     = OPEN
+        @opened_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+    end
+
+    def elapsed_since_open
+      Process.clock_gettime(Process::CLOCK_MONOTONIC) - @opened_at
+    end
+
+    def remaining_open_time
+      @open_timeout - elapsed_since_open
+    end
+  end
+
+  # Registry of per-broker CircuitBreakers, shared across transport calls.
+  class CircuitBreakerRegistry
+    def initialize(failure_threshold: 5, open_timeout: 30)
+      @failure_threshold = failure_threshold
+      @open_timeout      = open_timeout
+      @breakers          = {}
+      @mutex             = Mutex.new
+    end
+
+    def for(broker_address)
+      @mutex.synchronize do
+        @breakers[broker_address] ||= CircuitBreaker.new(
+          failure_threshold: @failure_threshold,
+          open_timeout:      @open_timeout
+        )
+      end
+    end
+
+    def open?(broker_address)
+      @mutex.synchronize { @breakers[broker_address]&.open? || false }
+    end
+
+    # Remove all state (useful for testing).
+    def reset_all
+      @mutex.synchronize { @breakers.clear }
+    end
+  end
+end

data/lib/pinot/config.rb

@@ -23,10 +23,13 @@ module Pinot
     attr_accessor :broker_list, :http_timeout, :query_timeout_ms, :extra_http_header,
                   :use_multistage_engine, :controller_config, :logger, :tls_config,
                   :grpc_config, :zookeeper_config,
-                  :max_retries,        # Integer, default 0 (no retry)
-                  :retry_interval_ms,  # Integer ms base interval, default 200
-                  :pool_size,          # max idle connections per broker, default 5
-                  :keep_alive_timeout  # seconds before idle connection is reaped, default 30
+                  :max_retries,                  # Integer, default 0 (no retry)
+                  :retry_interval_ms,            # Integer ms base interval, default 200
+                  :pool_size,                    # max idle connections per broker, default 5
+                  :keep_alive_timeout,           # seconds before idle connection is reaped, default 30
+                  :circuit_breaker_enabled,      # bool, default false
+                  :circuit_breaker_threshold,    # consecutive failures before opening, default 5
+                  :circuit_breaker_timeout       # seconds circuit stays open, default 30
 
     def initialize(
       broker_list: [],
@@ -42,7 +45,10 @@ module Pinot
       max_retries: 0,
       retry_interval_ms: 200,
       pool_size: nil,
-      keep_alive_timeout: nil
+      keep_alive_timeout: nil,
+      circuit_breaker_enabled: false,
+      circuit_breaker_threshold: 5,
+      circuit_breaker_timeout: 30
     )
       @broker_list = broker_list
       @http_timeout = http_timeout
@@ -58,6 +64,9 @@ module Pinot
       @retry_interval_ms = retry_interval_ms
       @pool_size = pool_size
       @keep_alive_timeout = keep_alive_timeout
+      @circuit_breaker_enabled = circuit_breaker_enabled
+      @circuit_breaker_threshold = circuit_breaker_threshold
+      @circuit_breaker_timeout = circuit_breaker_timeout
     end
 
     def validate!

data/lib/pinot/connection.rb

@@ -4,13 +4,15 @@ module Pinot
   class Connection
     attr_accessor :query_timeout_ms
 
-    def initialize(transport:, broker_selector:, use_multistage_engine: false, logger: nil, query_timeout_ms: nil)
+    def initialize(transport:, broker_selector:, use_multistage_engine: false, logger: nil,
+                   query_timeout_ms: nil, circuit_breaker_registry: nil)
       @transport = transport
       @broker_selector = broker_selector
       @use_multistage_engine = use_multistage_engine
       @trace = false
       @logger = logger
       @query_timeout_ms = query_timeout_ms
+      @circuit_breaker_registry = circuit_breaker_registry
     end
 
     def use_multistage_engine=(val)
@@ -25,12 +27,15 @@ module Pinot
       @trace = false
     end
 
-    def execute_sql(table, query, query_timeout_ms: nil)
+    def execute_sql(table, query, query_timeout_ms: nil, headers: {})
       Pinot::Instrumentation.instrument(table: table, query: query) do
         logger.debug "Executing SQL on table=#{table}: #{query}"
         broker = @broker_selector.select_broker(table)
         effective_timeout = query_timeout_ms || @query_timeout_ms
-        @transport.execute(broker, build_request(query, timeout_ms: effective_timeout))
+        run_with_circuit_breaker(broker) do
+          @transport.execute(broker, build_request(query, timeout_ms: effective_timeout),
+                             extra_request_headers: headers)
+        end
       end
     end
 
@@ -38,11 +43,10 @@ module Pinot
       execute_sql(table, query, query_timeout_ms: timeout_ms)
     end
 
-    def execute_sql_with_params(table, query_pattern, params, query_timeout_ms: nil)
+    def execute_sql_with_params(table, query_pattern, params, query_timeout_ms: nil, headers: {})
       query = format_query(query_pattern, params)
-      execute_sql(table, query, query_timeout_ms: query_timeout_ms)
+      execute_sql(table, query, query_timeout_ms: query_timeout_ms, headers: headers)
     rescue => e
-      # Re-raise format errors directly (they already have the right message)
       raise e
     end
 
@@ -97,6 +101,11 @@ module Pinot
 
     private
 
+    def run_with_circuit_breaker(broker, &block)
+      return yield unless @circuit_breaker_registry
+      @circuit_breaker_registry.for(broker).call(broker, &block)
+    end
+
     def logger
       @logger || Pinot::Logging.logger
     end

data/lib/pinot/connection_factory.rb

@@ -49,7 +49,8 @@ module Pinot
         broker_selector: selector,
         use_multistage_engine: config.use_multistage_engine || false,
         logger: config.logger,
-        query_timeout_ms: config.query_timeout_ms
+        query_timeout_ms: config.query_timeout_ms,
+        circuit_breaker_registry: build_circuit_breaker_registry(config)
       )
     end
 
@@ -71,7 +72,8 @@ module Pinot
       broker_selector: selector,
       use_multistage_engine: config.use_multistage_engine || false,
       logger: config.logger,
-      query_timeout_ms: config.query_timeout_ms
+      query_timeout_ms: config.query_timeout_ms,
+      circuit_breaker_registry: build_circuit_breaker_registry(config)
     )
 
     selector.init
@@ -96,4 +98,14 @@ module Pinot
     end
   end
   private_class_method :build_selector
+
+  def self.build_circuit_breaker_registry(config)
+    return nil unless config.circuit_breaker_enabled
+
+    CircuitBreakerRegistry.new(
+      failure_threshold: config.circuit_breaker_threshold || 5,
+      open_timeout:      config.circuit_breaker_timeout   || 30
+    )
+  end
+  private_class_method :build_circuit_breaker_registry
 end

data/lib/pinot/prepared_statement.rb

@@ -60,7 +60,7 @@ module Pinot
       nil
     end
 
-    def execute
+    def execute(headers: {})
       @mutex.synchronize do
         raise PreparedStatementClosedError, "prepared statement is closed" if @closed
         @parameters.each_with_index do |p, i|
@@ -72,10 +72,10 @@ module Pinot
       rescue => e
         raise "failed to build query: #{e.message}"
       end
-      @connection.execute_sql(@table, query)
+      @connection.execute_sql(@table, query, headers: headers)
     end
 
-    def execute_with_params(*params)
+    def execute_with_params(*params, headers: {})
       @mutex.synchronize { raise PreparedStatementClosedError, "prepared statement is closed" if @closed }
       if params.length != @param_count
         raise "expected #{@param_count} parameters, got #{params.length}"
@@ -85,7 +85,7 @@ module Pinot
       rescue => e
         raise "failed to build query: #{e.message}"
       end
-      @connection.execute_sql(@table, query)
+      @connection.execute_sql(@table, query, headers: headers)
     end
 
     def clear_parameters

data/lib/pinot/transport.rb

@@ -188,7 +188,7 @@ module Pinot
       @retry_interval_ms = retry_interval_ms
     end
 
-    def execute(broker_address, request)
+    def execute(broker_address, request, extra_request_headers: {})
       logger.debug "Pinot query to #{broker_address}: #{request.query}"
 
       attempts = 0
@@ -202,6 +202,7 @@ module Pinot
         headers = DEFAULT_HEADERS
           .merge(@extra_headers)
           .merge("X-Correlation-Id" => SecureRandom.uuid)
+          .merge(extra_request_headers)
 
         resp = @http_client.post(url, body: body, headers: headers)
 

data/lib/pinot/version.rb

@@ -1,3 +1,3 @@
 module Pinot
-  VERSION = "1.21.0"
+  VERSION = "1.23.0"
 end

data/lib/pinot.rb

@@ -18,6 +18,7 @@ require_relative "pinot/table_aware_broker_selector"
 require_relative "pinot/controller_response"
 require_relative "pinot/controller_based_broker_selector"
 require_relative "pinot/transport"
+require_relative "pinot/circuit_breaker"
 require_relative "pinot/connection"
 require_relative "pinot/prepared_statement"
 require_relative "pinot/connection_factory"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pinot-client
 version: !ruby/object:Gem::Version
-  version: 1.21.0
+  version: 1.23.0
 platform: ruby
 authors:
 - Xiang Fu
@@ -104,6 +104,7 @@ files:
 - README.md
 - lib/pinot.rb
 - lib/pinot/broker_selector.rb
+- lib/pinot/circuit_breaker.rb
 - lib/pinot/config.rb
 - lib/pinot/connection.rb
 - lib/pinot/connection_factory.rb