pinot-client 1.20.0 → 1.21.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9392ad85128da7b9c3e85dfe0311316174fe36b014d573784f0de7fd60e5e9f2
-  data.tar.gz: dec15a5ce58afb3a5e9d21c374c0c089d8f7855d44df540dd94f2cf2fe7a34ac
+  metadata.gz: f91a5f7f465b546a3a49c78d2076097eac9b6af8fd9fbee7f8c1c45df06bea16
+  data.tar.gz: 0a04dd7945a401c01e107c95062b738ff987a23706338d7436e88ff8e50bec13
 SHA512:
-  metadata.gz: 7eb05753626860442cc047f379a4a0a74306a72b9d3d15bdd8fcf0d4c8f743653dc8d6213883807ac9ffa7e23263581e0a32aeefe7c5fa5ece3fcd7f1298ace8
-  data.tar.gz: f4b1d0c916e95e5c93c9792bc16cdedf85f90b582ca15c5c3c5fe6272253a24973e3bc629cb927501e7650b25eacba091218eb332bebbd22e028c08846217856
+  metadata.gz: aeee0c948cca37c67a7b8bb103201ea4a4ad1dbe9e08476ca679efb10c4ea9951ff4ade65af6e8f7cc11e3c55b33ca66f969b0e4271fe97a8bd62bc7265c24aa
+  data.tar.gz: 8701f59b1adc42307e0774484a24d5c0e0cd33e8da378d9d2ccc65ad0240befcd094d58a136dc5031150fa687e88bbf87cbdde29e02d69e96462eedaf31dcbdc

data/lib/pinot/config.rb

@@ -23,8 +23,10 @@ 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
 
     def initialize(
       broker_list: [],
@@ -38,7 +40,9 @@ 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
     )
       @broker_list = broker_list
       @http_timeout = http_timeout
@@ -50,9 +54,10 @@ 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
     end
 
     def validate!
@@ -75,6 +80,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_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 || {},
@@ -53,7 +53,7 @@ module Pinot
       )
     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,
@@ -78,6 +78,16 @@ module Pinot
     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)

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.21.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pinot-client
 version: !ruby/object:Gem::Version
-  version: 1.20.0
+  version: 1.21.0
 platform: ruby
 authors:
 - Xiang Fu