Refactor pool checkout logic

- Raises a wrapped exception on any DB:Error
- Don't use a non-pool client when client fails
- Ensure that client is always released
This commit is contained in:
syeopite 2024-11-11 15:58:35 -08:00
parent 6ff4024244
commit 71166034f0
No known key found for this signature in database
GPG Key ID: A73C186DA3955A1A
3 changed files with 72 additions and 64 deletions

View File

@ -35,6 +35,7 @@ require "protodec/utils"
require "./invidious/database/*" require "./invidious/database/*"
require "./invidious/database/migrations/*" require "./invidious/database/migrations/*"
require "./invidious/connection/*"
require "./invidious/http_server/*" require "./invidious/http_server/*"
require "./invidious/helpers/*" require "./invidious/helpers/*"
require "./invidious/yt_backend/*" require "./invidious/yt_backend/*"
@ -91,11 +92,11 @@ SOFTWARE = {
"branch" => "#{CURRENT_BRANCH}", "branch" => "#{CURRENT_BRANCH}",
} }
YT_POOL = YoutubeConnectionPool.new(YT_URL, max_capacity: CONFIG.pool_size, idle_capacity: CONFIG.idle_pool_size) YT_POOL = Invidious::ConnectionPool::Pool.new(YT_URL, max_capacity: CONFIG.pool_size, idle_capacity: CONFIG.idle_pool_size)
# Image request pool # Image request pool
GGPHT_POOL = YoutubeConnectionPool.new( GGPHT_POOL = Invidious::ConnectionPool::Pool.new(
URI.parse("https://yt3.ggpht.com"), URI.parse("https://yt3.ggpht.com"),
max_capacity: CONFIG.pool_size, max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size idle_capacity: CONFIG.idle_pool_size

View File

@ -0,0 +1,66 @@
module Invidious::ConnectionPool
struct Pool
property! url : URI
property! max_capacity : Int32
property! idle_capacity : Int32
property! timeout : Float64
property pool : DB::Pool(HTTP::Client)
def initialize(
url : URI,
*,
@max_capacity : Int32 = 5,
idle_capacity : Int32? = nil,
@timeout : Float64 = 5.0
)
if idle_capacity.nil?
@idle_capacity = @max_capacity
else
@idle_capacity = idle_capacity
end
@url = url
@pool = build_pool()
end
# Checks out a client in the pool
def client(&)
pool.checkout do |http_client|
# Proxy needs to be reinstated every time we get a client from the pool
http_client.proxy = make_configured_http_proxy_client() if CONFIG.http_proxy
response = yield http_client
return response
rescue ex : DB::Error
# Prevent broken client from being checked back into the pool
http_client.close
raise ConnectionPool::Error.new(ex.message, cause: ex)
ensure
pool.release(http_client)
end
rescue ex : DB::PoolTimeout
# Failed to checkout a client
raise ConnectionPool::Error.new(ex.message, cause: ex)
end
private def build_pool
# We call the getter for the instance variables instead of using them directly
# because the getters defined by property! ensures that the value is not a nil
options = DB::Pool::Options.new(
initial_pool_size: 0,
max_pool_size: max_capacity,
max_idle_pool_size: idle_capacity,
checkout_timeout: timeout
)
DB::Pool(HTTP::Client).new(options) do
next make_client(url, force_resolve: true)
end
end
end
class Error < Exception
end
end

View File

@ -1,65 +1,6 @@
# Mapping of subdomain => YoutubeConnectionPool # Mapping of subdomain => Invidious::ConnectionPool::Pool
# This is needed as we may need to access arbitrary subdomains of ytimg # This is needed as we may need to access arbitrary subdomains of ytimg
private YTIMG_POOLS = {} of String => YoutubeConnectionPool private YTIMG_POOLS = {} of String => Invidious::ConnectionPool::Pool
struct YoutubeConnectionPool
property! url : URI
property! max_capacity : Int32
property! idle_capacity : Int32
property! timeout : Float64
property pool : DB::Pool(HTTP::Client)
def initialize(
url : URI,
*,
@max_capacity : Int32 = 5,
idle_capacity : Int32? = nil,
@timeout : Float64 = 5.0
)
if idle_capacity.nil?
@idle_capacity = @max_capacity
else
@idle_capacity = idle_capacity
end
@url = url
@pool = build_pool()
end
def client(&)
conn = pool.checkout
# Proxy needs to be reinstated every time we get a client from the pool
conn.proxy = make_configured_http_proxy_client() if CONFIG.http_proxy
begin
response = yield conn
rescue ex
conn.close
conn = make_client(url, force_resolve: true)
response = yield conn
ensure
pool.release(conn)
end
response
end
private def build_pool
# We call the getter for the instance variables instead of using them directly
# because the getters defined by property! ensures that the value is not a nil
options = DB::Pool::Options.new(
initial_pool_size: 0,
max_pool_size: max_capacity,
max_idle_pool_size: idle_capacity,
checkout_timeout: timeout
)
DB::Pool(HTTP::Client).new(options) do
next make_client(url, force_resolve: true)
end
end
end
def add_yt_headers(request) def add_yt_headers(request)
request.headers.delete("User-Agent") if request.headers["User-Agent"] == "Crystal" request.headers.delete("User-Agent") if request.headers["User-Agent"] == "Crystal"
@ -123,7 +64,7 @@ def get_ytimg_pool(subdomain)
return pool return pool
else else
LOGGER.info("ytimg_pool: Creating a new HTTP pool for \"https://#{subdomain}.ytimg.com\"") LOGGER.info("ytimg_pool: Creating a new HTTP pool for \"https://#{subdomain}.ytimg.com\"")
pool = YoutubeConnectionPool.new( pool = Invidious::ConnectionPool::Pool.new(
URI.parse("https://#{subdomain}.ytimg.com"), URI.parse("https://#{subdomain}.ytimg.com"),
max_capacity: CONFIG.pool_size, max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size idle_capacity: CONFIG.idle_pool_size