Merge companion and standard pool into one

This commit is contained in:
syeopite 2025-04-10 00:31:41 -07:00
parent fbccb6a221
commit 7901906092
No known key found for this signature in database
GPG Key ID: A73C186DA3955A1A
3 changed files with 37 additions and 67 deletions

View File

@ -21,6 +21,8 @@ require "../../load_config"
require "../../../src/invidious/helpers/crystal_class_overrides"
require "../../../src/invidious/connection/*"
TEST_SERVER_URL = URI.parse("http://localhost:12345")
server = HTTP::Server.new do |context|
request = context.request
response = context.response
@ -44,14 +46,14 @@ Fiber.yield
Spectator.describe Invidious::ConnectionPool do
describe "Pool" do
it "Can make a requests through standard HTTP methods" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 100)
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 100) { next make_client(TEST_SERVER_URL) }
expect(pool.get("/get").body).to eq("get")
expect(pool.post("/post").body).to eq("post")
end
it "Can make streaming requests" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 100)
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 100) { next make_client(TEST_SERVER_URL) }
expect(pool.get("/get") { |r| r.body_io.gets_to_end }).to eq("get")
expect(pool.get("/post") { |r| r.body }).to eq("")
@ -59,26 +61,25 @@ Spectator.describe Invidious::ConnectionPool do
end
it "Allows more than one clients to be checked out (if applicable)" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 100)
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 100) { next make_client(TEST_SERVER_URL) }
pool.checkout do | client |
pool.checkout do |client|
expect(pool.post("/post").body).to eq("post")
end
end
it "Can make multiple requests with the same client" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 100)
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 100) { next make_client(TEST_SERVER_URL) }
pool.checkout do | client |
pool.checkout do |client|
expect(client.get("/get").body).to eq("get")
expect(client.post("/post").body).to eq("post")
expect(client.get("/get").body).to eq("get")
end
end
it "Allows concurrent requests" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 100)
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 100) { next make_client(TEST_SERVER_URL) }
responses = [] of HTTP::Client::Response
WaitGroup.wait do |wg|
@ -91,7 +92,7 @@ Spectator.describe Invidious::ConnectionPool do
end
it "Raises on checkout timeout" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 2, timeout: 0.01)
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 2, timeout: 0.01) { next make_client(TEST_SERVER_URL) }
# Long running requests
2.times do
@ -103,8 +104,8 @@ Spectator.describe Invidious::ConnectionPool do
expect { pool.get("/get") }.to raise_error(Invidious::ConnectionPool::Error)
end
it "Raises when an error is encounter" do
pool = Invidious::ConnectionPool::Pool.new(URI.parse("http://localhost:12345"), max_capacity: 100, timeout: 0.01)
it "Raises when an error is encountered" do
pool = Invidious::ConnectionPool::Pool.new(max_capacity: 100) { next make_client(TEST_SERVER_URL) }
expect { pool.get("/get") { raise IO::Error.new } }.to raise_error(Invidious::ConnectionPool::Error)
end
end

View File

@ -93,25 +93,32 @@ SOFTWARE = {
}
YT_POOL = Invidious::ConnectionPool::Pool.new(
YT_URL,
max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size,
timeout: CONFIG.pool_checkout_timeout
)
) do
next make_client(YT_URL, force_resolve: true)
end
# Image request pool
GGPHT_URL = URI.parse("https://yt3.ggpht.com")
GGPHT_POOL = Invidious::ConnectionPool::Pool.new(
URI.parse("https://yt3.ggpht.com"),
max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size,
timeout: CONFIG.pool_checkout_timeout
)
) do
next make_client(GGPHT_URL, force_resolve: true)
end
COMPANION_POOL = Invidious::ConnectionPool::CompanionPool.new(
COMPANION_POOL = Invidious::ConnectionPool::Pool.new(
max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size
)
) do
companion = CONFIG.invidious_companion.sample
next make_client(companion.private_url, use_http_proxy: false)
end
# CLI
Kemal.config.extra_options do |parser|

View File

@ -1,15 +1,15 @@
module Invidious::ConnectionPool
# The base connection pool that provides the underlying logic that all connection pools are based around
#
# Uses `DB::Pool` for the pooling logic
abstract struct BaseConnectionPool(PoolClient)
# A connection pool to reuse `HTTP::Client` connections
struct Pool
getter pool : DB::Pool(HTTP::Client)
# Creates a connection pool with the provided options, and client factory block.
def initialize(
*,
max_capacity : Int32 = 5,
idle_capacity : Int32? = nil,
timeout : Float64 = 5.0,
&client_factory : -> PoolClient
&client_factory : -> HTTP::Client
)
if idle_capacity.nil?
idle_capacity = max_capacity
@ -22,12 +22,9 @@ module Invidious::ConnectionPool
checkout_timeout: timeout
)
@pool = DB::Pool(PoolClient).new(pool_options, &client_factory)
@pool = DB::Pool(HTTP::Client).new(pool_options, &client_factory)
end
# Returns the underlying `DB::Pool` object
abstract def pool : DB::Pool(PoolClient)
{% for method in %w[get post put patch delete head options] %}
# Streaming API for {{method.id.upcase}} request.
# The response will have its body as an `IO` accessed via `HTTP::Client::Response#body_io`.
@ -89,45 +86,6 @@ module Invidious::ConnectionPool
end
end
# A basic connection pool where each client within is set to connect to a single resource
struct Pool < BaseConnectionPool(HTTP::Client)
getter pool : DB::Pool(HTTP::Client)
# Creates a pool of clients that connects to the given url, with the provided options.
def initialize(
url : URI,
*,
max_capacity : Int32 = 5,
idle_capacity : Int32? = nil,
timeout : Float64 = 5.0,
)
super(max_capacity: max_capacity, idle_capacity: idle_capacity, timeout: timeout) do
next make_client(url, force_resolve: true)
end
end
end
# A modified connection pool for the interacting with Invidious companion.
#
# The main difference is that clients in this pool are created with different urls
# based on what is randomly selected from the configured list of companions
struct CompanionPool < BaseConnectionPool(HTTP::Client)
getter pool : DB::Pool(HTTP::Client)
# Creates a pool of clients with the provided options.
def initialize(
*,
max_capacity : Int32 = 5,
idle_capacity : Int32? = nil,
timeout : Float64 = 5.0,
)
super(max_capacity: max_capacity, idle_capacity: idle_capacity, timeout: timeout) do
companion = CONFIG.invidious_companion.sample
next make_client(companion.private_url, use_http_proxy: false)
end
end
end
class Error < Exception
end
@ -147,12 +105,16 @@ module Invidious::ConnectionPool
return pool
else
LOGGER.info("ytimg_pool: Creating a new HTTP pool for \"https://#{subdomain}.ytimg.com\"")
url = URI.parse("https://#{subdomain}.ytimg.com")
pool = ConnectionPool::Pool.new(
URI.parse("https://#{subdomain}.ytimg.com"),
max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size,
timeout: CONFIG.pool_checkout_timeout
)
) do
next make_client(url, force_resolve: true)
end
YTIMG_POOLS[subdomain] = pool
return pool