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

View File

@ -93,25 +93,32 @@ SOFTWARE = {
} }
YT_POOL = Invidious::ConnectionPool::Pool.new( YT_POOL = Invidious::ConnectionPool::Pool.new(
YT_URL,
max_capacity: CONFIG.pool_size, max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_pool_size, idle_capacity: CONFIG.idle_pool_size,
timeout: CONFIG.pool_checkout_timeout timeout: CONFIG.pool_checkout_timeout
) ) do
next make_client(YT_URL, force_resolve: true)
end
# Image request pool # Image request pool
GGPHT_URL = URI.parse("https://yt3.ggpht.com")
GGPHT_POOL = Invidious::ConnectionPool::Pool.new( GGPHT_POOL = Invidious::ConnectionPool::Pool.new(
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,
timeout: CONFIG.pool_checkout_timeout 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, max_capacity: CONFIG.pool_size,
idle_capacity: CONFIG.idle_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 # CLI
Kemal.config.extra_options do |parser| Kemal.config.extra_options do |parser|

View File

@ -1,15 +1,15 @@
module Invidious::ConnectionPool module Invidious::ConnectionPool
# The base connection pool that provides the underlying logic that all connection pools are based around # A connection pool to reuse `HTTP::Client` connections
# struct Pool
# Uses `DB::Pool` for the pooling logic getter pool : DB::Pool(HTTP::Client)
abstract struct BaseConnectionPool(PoolClient)
# Creates a connection pool with the provided options, and client factory block. # Creates a connection pool with the provided options, and client factory block.
def initialize( def initialize(
*, *,
max_capacity : Int32 = 5, max_capacity : Int32 = 5,
idle_capacity : Int32? = nil, idle_capacity : Int32? = nil,
timeout : Float64 = 5.0, timeout : Float64 = 5.0,
&client_factory : -> PoolClient &client_factory : -> HTTP::Client
) )
if idle_capacity.nil? if idle_capacity.nil?
idle_capacity = max_capacity idle_capacity = max_capacity
@ -22,12 +22,9 @@ module Invidious::ConnectionPool
checkout_timeout: timeout checkout_timeout: timeout
) )
@pool = DB::Pool(PoolClient).new(pool_options, &client_factory) @pool = DB::Pool(HTTP::Client).new(pool_options, &client_factory)
end end
# Returns the underlying `DB::Pool` object
abstract def pool : DB::Pool(PoolClient)
{% for method in %w[get post put patch delete head options] %} {% for method in %w[get post put patch delete head options] %}
# Streaming API for {{method.id.upcase}} request. # Streaming API for {{method.id.upcase}} request.
# The response will have its body as an `IO` accessed via `HTTP::Client::Response#body_io`. # 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
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 class Error < Exception
end end
@ -147,12 +105,16 @@ module Invidious::ConnectionPool
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\"")
url = URI.parse("https://#{subdomain}.ytimg.com")
pool = ConnectionPool::Pool.new( pool = ConnectionPool::Pool.new(
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,
timeout: CONFIG.pool_checkout_timeout timeout: CONFIG.pool_checkout_timeout
) ) do
next make_client(url, force_resolve: true)
end
YTIMG_POOLS[subdomain] = pool YTIMG_POOLS[subdomain] = pool
return pool return pool