From 71166034f0d71d5b3ff53af0bedb1ab127782bbf Mon Sep 17 00:00:00 2001
From: syeopite <syeopite@syeopite.dev>
Date: Mon, 11 Nov 2024 15:58:35 -0800
Subject: [PATCH] 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
---
 src/invidious.cr                            |  5 +-
 src/invidious/connection/pool.cr            | 66 +++++++++++++++++++++
 src/invidious/yt_backend/connection_pool.cr | 65 +-------------------
 3 files changed, 72 insertions(+), 64 deletions(-)
 create mode 100644 src/invidious/connection/pool.cr

diff --git a/src/invidious.cr b/src/invidious.cr
index e2a91d8d..189cde40 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -35,6 +35,7 @@ require "protodec/utils"
 
 require "./invidious/database/*"
 require "./invidious/database/migrations/*"
+require "./invidious/connection/*"
 require "./invidious/http_server/*"
 require "./invidious/helpers/*"
 require "./invidious/yt_backend/*"
@@ -91,11 +92,11 @@ SOFTWARE = {
   "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
 
-GGPHT_POOL = YoutubeConnectionPool.new(
+GGPHT_POOL = Invidious::ConnectionPool::Pool.new(
   URI.parse("https://yt3.ggpht.com"),
   max_capacity: CONFIG.pool_size,
   idle_capacity: CONFIG.idle_pool_size
diff --git a/src/invidious/connection/pool.cr b/src/invidious/connection/pool.cr
new file mode 100644
index 00000000..a6b09f02
--- /dev/null
+++ b/src/invidious/connection/pool.cr
@@ -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
diff --git a/src/invidious/yt_backend/connection_pool.cr b/src/invidious/yt_backend/connection_pool.cr
index 5eca9184..1066cade 100644
--- a/src/invidious/yt_backend/connection_pool.cr
+++ b/src/invidious/yt_backend/connection_pool.cr
@@ -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
-private YTIMG_POOLS = {} of String => YoutubeConnectionPool
-
-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
+private YTIMG_POOLS = {} of String => Invidious::ConnectionPool::Pool
 
 def add_yt_headers(request)
   request.headers.delete("User-Agent") if request.headers["User-Agent"] == "Crystal"
@@ -123,7 +64,7 @@ def get_ytimg_pool(subdomain)
     return pool
   else
     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"),
       max_capacity: CONFIG.pool_size,
       idle_capacity: CONFIG.idle_pool_size