pinot-client 1.20.0 → 1.22.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9392ad85128da7b9c3e85dfe0311316174fe36b014d573784f0de7fd60e5e9f2
-  data.tar.gz: dec15a5ce58afb3a5e9d21c374c0c089d8f7855d44df540dd94f2cf2fe7a34ac
+  metadata.gz: 8fa5d08fb08bee353252e2da56932948941faa5fdbdd4f3410e5ff8887755e04
+  data.tar.gz: c8bb9cbe06c5d209d8f9f63031776fc604a153730b4be207b7fd6d8a3b69e5f0
 SHA512:
-  metadata.gz: 7eb05753626860442cc047f379a4a0a74306a72b9d3d15bdd8fcf0d4c8f743653dc8d6213883807ac9ffa7e23263581e0a32aeefe7c5fa5ece3fcd7f1298ace8
-  data.tar.gz: f4b1d0c916e95e5c93c9792bc16cdedf85f90b582ca15c5c3c5fe6272253a24973e3bc629cb927501e7650b25eacba091218eb332bebbd22e028c08846217856
+  metadata.gz: 6d5c85b83effd2ce352b09c03d350eedd817415a18994ed23ef7be972f80dfb8d72cbbbace9c83bfddedb095b6695af0b8b4538f2f0023ee40844e6b554849a9
+  data.tar.gz: c450ffb393a8798ae3ea3e608cf2cb976bdf5aeb6f7145121764c41bed33069807fdb3087d02c3963552f79568fb6f2020088794cccf47f97218d10fe7041c99

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,8 +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
+                  :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: [],
@@ -38,7 +43,12 @@ module Pinot
       grpc_config: nil,
       zookeeper_config: nil,
       max_retries: 0,
-      retry_interval_ms: 200
+      retry_interval_ms: 200,
+      pool_size: 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
@@ -50,9 +60,13 @@ module Pinot
       @tls_config = tls_config
       @grpc_config = grpc_config
       @zookeeper_config = zookeeper_config
-      @query_timeout_ms = query_timeout_ms
       @max_retries = max_retries
       @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!
@@ -75,6 +89,14 @@ module Pinot
         raise ConfigurationError, "query_timeout_ms must be positive, got: #{query_timeout_ms}"
       end
 
+      if !pool_size.nil? && pool_size < 1
+        raise ConfigurationError, "pool_size must be at least 1, got: #{pool_size}"
+      end
+
+      if !keep_alive_timeout.nil? && keep_alive_timeout <= 0
+        raise ConfigurationError, "keep_alive_timeout must be positive, got: #{keep_alive_timeout}"
+      end
+
       self
     end
   end

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)
@@ -30,7 +32,9 @@ module Pinot
         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))
+        end
       end
     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

@@ -34,7 +34,7 @@ module Pinot
       selector = ZookeeperBrokerSelector.new(zk_path: config.zookeeper_config.zk_path)
       selector.init
 
-      inner = http_client || HttpClient.new(timeout: config.http_timeout, tls_config: config.tls_config)
+      inner = http_client || build_http_client(config)
       transport = JsonHttpTransport.new(
         http_client: inner,
         extra_headers: config.extra_http_header || {},
@@ -49,11 +49,12 @@ 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
 
-    inner = http_client || HttpClient.new(timeout: config.http_timeout, tls_config: config.tls_config)
+    inner = http_client || build_http_client(config)
 
     transport = JsonHttpTransport.new(
       http_client: inner,
@@ -71,13 +72,24 @@ 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
     conn
   end
 
+  def self.build_http_client(config)
+    HttpClient.new(
+      timeout: config.http_timeout,
+      tls_config: config.tls_config,
+      pool_size: config.pool_size,
+      keep_alive_timeout: config.keep_alive_timeout
+    )
+  end
+  private_class_method :build_http_client
+
   def self.build_selector(config, http_client)
     if config.broker_list && !config.broker_list.empty?
       SimpleBrokerSelector.new(config.broker_list)
@@ -86,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/errors.rb

@@ -3,6 +3,9 @@ module Pinot
   class BrokerNotFoundError < Error; end
   class TableNotFoundError < Error; end
   class TransportError < Error; end
+  class BrokerUnavailableError < TransportError; end
+  class QueryTimeoutError < TransportError; end
+  class RateLimitError < TransportError; end
   class PreparedStatementClosedError < Error; end
   class ConfigurationError < Error; end
 end

data/lib/pinot/transport.rb

@@ -6,14 +6,16 @@ require "openssl"
 
 module Pinot
   class HttpClient
-    MAX_POOL_SIZE = 5
-    KEEP_ALIVE_TIMEOUT = 30
+    DEFAULT_POOL_SIZE = 5
+    DEFAULT_KEEP_ALIVE_TIMEOUT = 30
 
     PoolEntry = Struct.new(:http, :checked_in_at)
 
-    def initialize(timeout: nil, tls_config: nil)
+    def initialize(timeout: nil, tls_config: nil, pool_size: nil, keep_alive_timeout: nil)
       @timeout = timeout
       @tls_config = tls_config
+      @max_pool_size = pool_size || DEFAULT_POOL_SIZE
+      @keep_alive_timeout = keep_alive_timeout || DEFAULT_KEEP_ALIVE_TIMEOUT
       @pool = {}
       @pool_mutex = Mutex.new
       @reaper = start_reaper
@@ -70,7 +72,7 @@ module Pinot
         entries = @pool[key] ||= []
         fresh = nil
         while (entry = entries.pop)
-          if now - entry.checked_in_at < KEEP_ALIVE_TIMEOUT
+          if now - entry.checked_in_at < @keep_alive_timeout
             fresh = entry.http
             break
           else
@@ -85,7 +87,7 @@ module Pinot
     def checkin(key, http)
       @pool_mutex.synchronize do
         pool_for_key = @pool[key] ||= []
-        if pool_for_key.size < MAX_POOL_SIZE
+        if pool_for_key.size < @max_pool_size
           pool_for_key.push(PoolEntry.new(http, Process.clock_gettime(Process::CLOCK_MONOTONIC)))
         else
           http.finish rescue nil
@@ -96,7 +98,7 @@ module Pinot
     def start_reaper
       t = Thread.new do
         loop do
-          sleep KEEP_ALIVE_TIMEOUT / 2.0
+          sleep @keep_alive_timeout / 2.0
           reap_stale_connections
         end
       end
@@ -109,7 +111,7 @@ module Pinot
       @pool_mutex.synchronize do
         @pool.each_value do |entries|
           entries.reject! do |entry|
-            if now - entry.checked_in_at >= KEEP_ALIVE_TIMEOUT
+            if now - entry.checked_in_at >= @keep_alive_timeout
               entry.http.finish rescue nil
               true
             else
@@ -167,6 +169,15 @@ module Pinot
       Net::OpenTimeout, Net::ReadTimeout, Net::WriteTimeout
     ].freeze
 
+    # HTTP status codes that map to specific error classes and are safe to retry
+    HTTP_ERROR_MAP = {
+      "429" => RateLimitError,
+      "503" => BrokerUnavailableError,
+      "504" => BrokerUnavailableError
+    }.freeze
+
+    RETRYABLE_HTTP_ERRORS = [RateLimitError, BrokerUnavailableError].freeze
+
     def initialize(http_client:, extra_headers: {}, timeout_ms: nil, logger: nil,
                    max_retries: 0, retry_interval_ms: 200)
       @http_client = http_client
@@ -194,9 +205,9 @@ module Pinot
 
         resp = @http_client.post(url, body: body, headers: headers)
 
-        if resp.code == "503"
+        if (error_class = HTTP_ERROR_MAP[resp.code])
           logger.error "Pinot broker returned HTTP #{resp.code}"
-          raise TransportError, "http exception with HTTP status code #{resp.code}"
+          raise error_class, "http exception with HTTP status code #{resp.code}"
         end
 
         unless resp.code.to_i == 200
@@ -209,7 +220,7 @@ module Pinot
         rescue JSON::ParserError => e
           raise e.message
         end
-      rescue TransportError, *RETRYABLE_ERRORS => e
+      rescue *RETRYABLE_HTTP_ERRORS, *RETRYABLE_ERRORS => e
         if attempts < max_attempts
           sleep_ms = (@retry_interval_ms || 200) * (2 ** (attempts - 1))
           sleep(sleep_ms / 1000.0)

data/lib/pinot/version.rb

@@ -1,3 +1,3 @@
 module Pinot
-  VERSION = "1.20.0"
+  VERSION = "1.22.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.20.0
+  version: 1.22.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