Merge companion and standard pool into one

Author: syeopite

spec/helpers/networking/connection_pool_spec.cr

@@ -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,41 +46,40 @@ 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("")
       expect(pool.post("/post") { |r| r.body_io.gets_to_end }).to eq("post")
     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

src/invidious.cr

@@ -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|

src/invidious/connection/pool.cr

@@ -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