From 480e073fa9be184b6839619c38795af582247c19 Mon Sep 17 00:00:00 2001
From: syeopite <syeopite@syeopite.dev>
Date: Fri, 8 Dec 2023 18:20:17 -0800
Subject: [PATCH] Use HTTP pools for image requests to YouTube

---
 src/invidious.cr                            |  8 ++++++++
 src/invidious/routes/images.cr              | 12 +++++-------
 src/invidious/yt_backend/connection_pool.cr | 15 +++++++++++++++
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/invidious.cr b/src/invidious.cr
index 3804197e..81db2c6c 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -92,6 +92,14 @@ SOFTWARE = {
 
 YT_POOL = YoutubeConnectionPool.new(YT_URL, capacity: CONFIG.pool_size)
 
+# Image request pool
+
+GGPHT_POOL = YoutubeConnectionPool.new(URI.parse("https://yt3.ggpht.com"), capacity: CONFIG.pool_size)
+
+# Mapping of subdomain => YoutubeConnectionPool
+# This is needed as we may need to access arbitrary subdomains of ytimg
+YTIMG_POOLS = {} of String => YoutubeConnectionPool
+
 # CLI
 Kemal.config.extra_options do |parser|
   parser.banner = "Usage: invidious [arguments]"
diff --git a/src/invidious/routes/images.cr b/src/invidious/routes/images.cr
index b6a2e110..1964d597 100644
--- a/src/invidious/routes/images.cr
+++ b/src/invidious/routes/images.cr
@@ -32,7 +32,7 @@ module Invidious::Routes::Images
     }
 
     begin
-      HTTP::Client.get("https://yt3.ggpht.com#{url}") do |resp|
+      GGPHT_POOL.client &.get(url) do |resp|
         return request_proc.call(resp)
       end
     rescue ex
@@ -80,7 +80,7 @@ module Invidious::Routes::Images
     }
 
     begin
-      HTTP::Client.get("https://#{authority}.ytimg.com#{url}") do |resp|
+      get_ytimg_pool(authority).client &.get(url) do |resp|
         return request_proc.call(resp)
       end
     rescue ex
@@ -119,7 +119,7 @@ module Invidious::Routes::Images
     }
 
     begin
-      HTTP::Client.get("https://i9.ytimg.com#{url}") do |resp|
+      get_ytimg_pool("i9").client &.get(url) do |resp|
         return request_proc.call(resp)
       end
     rescue ex
@@ -165,8 +165,7 @@ module Invidious::Routes::Images
     if name == "maxres.jpg"
       build_thumbnails(id).each do |thumb|
         thumbnail_resource_path = "/vi/#{id}/#{thumb[:url]}.jpg"
-        # This can likely be optimized into a (small) pool sometime in the future.
-        if HTTP::Client.head("https://i.ytimg.com#{thumbnail_resource_path}").status_code == 200
+        if get_ytimg_pool("i9").client &.head(thumbnail_resource_path).status_code == 200
           name = thumb[:url] + ".jpg"
           break
         end
@@ -199,8 +198,7 @@ module Invidious::Routes::Images
     }
 
     begin
-      # This can likely be optimized into a (small) pool sometime in the future.
-      HTTP::Client.get("https://i.ytimg.com#{url}") do |resp|
+      get_ytimg_pool("i").client &.get(url) do |resp|
         return request_proc.call(resp)
       end
     rescue ex
diff --git a/src/invidious/yt_backend/connection_pool.cr b/src/invidious/yt_backend/connection_pool.cr
index ca612083..26bf2773 100644
--- a/src/invidious/yt_backend/connection_pool.cr
+++ b/src/invidious/yt_backend/connection_pool.cr
@@ -77,3 +77,18 @@ def make_client(url : URI, region = nil, force_resolve : Bool = false, &)
     client.close
   end
 end
+
+# Fetches a HTTP pool for the specified subdomain of ytimg.com
+#
+# Creates a new one when the specified pool for the subdomain does not exist
+def get_ytimg_pool(subdomain)
+  if pool = YTIMG_POOLS[subdomain]?
+    return pool
+  else
+    LOGGER.info("ytimg_pool: Creating a new HTTP pool for \"https://#{subdomain}.ytimg.com\"")
+    pool = YoutubeConnectionPool.new(URI.parse("https://#{subdomain}.ytimg.com"), capacity: CONFIG.pool_size)
+    YTIMG_POOLS[subdomain] = pool
+
+    return pool
+  end
+end