From 4dcd6dd2f45f3f644581e3f4dad45bd9c8a670ac Mon Sep 17 00:00:00 2001
From: Kevinf100 <12779728+kevinf100@users.noreply.github.com>
Date: Tue, 3 Jun 2025 10:34:56 -0400
Subject: [PATCH] Update API to only try working backends with companion

---
 src/invidious/helpers/backend_info.cr   | 13 +++++++++++++
 src/invidious/routes/before_all.cr      |  2 +-
 src/invidious/yt_backend/youtube_api.cr |  3 ++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/invidious/helpers/backend_info.cr b/src/invidious/helpers/backend_info.cr
index 107258a5..b41768b4 100644
--- a/src/invidious/helpers/backend_info.cr
+++ b/src/invidious/helpers/backend_info.cr
@@ -3,10 +3,19 @@ module BackendInfo
   @@exvpp_url : Array(String) = Array.new(CONFIG.invidious_companion.size, "")
   @@status : Array(Int32) = Array.new(CONFIG.invidious_companion.size, 0)
   @@csp : Array(String) = Array.new(CONFIG.invidious_companion.size, "")
+  @@working_ends : Array(Int32) = Array(Int32).new(0)
   @@mutex : Mutex = Mutex.new
 
   def check_backends
     check_companion()
+    LOGGER.debug("Invidious companion: Updating working_ends")
+    @@working_ends.clear
+    @@status.each_with_index do |_, index|
+      if @@status[index] == 2
+          @@working_ends.push(index)
+      end
+    end
+    LOGGER.debug("Invidious companion: New working_ends \"#{@@working_ends}\"")
   end
 
   private def check_companion
@@ -70,6 +79,10 @@ module BackendInfo
     return @@status
   end
 
+  def get_working_ends
+    return @@working_ends
+  end
+
   def get_exvpp
     return @@exvpp_url
   end
diff --git a/src/invidious/routes/before_all.cr b/src/invidious/routes/before_all.cr
index 20113a6c..e9228b4e 100644
--- a/src/invidious/routes/before_all.cr
+++ b/src/invidious/routes/before_all.cr
@@ -46,7 +46,7 @@ module Invidious::Routes::BeforeAll
         begin
           current_companion = env.request.cookies[CONFIG.server_id_cookie_name].value.try &.to_i
         rescue
-          current_companion = rand(CONFIG.invidious_companion.size)
+          current_companion = BackendInfo.get_working_ends.sample
         end
 
         if current_companion > CONFIG.invidious_companion.size
diff --git a/src/invidious/yt_backend/youtube_api.cr b/src/invidious/yt_backend/youtube_api.cr
index e811e2ba..d4a43e0e 100644
--- a/src/invidious/yt_backend/youtube_api.cr
+++ b/src/invidious/yt_backend/youtube_api.cr
@@ -689,10 +689,11 @@ module YoutubeAPI
 
     begin
       if env.nil?
-        current_companion = rand(CONFIG.invidious_companion.size)
+        current_companion = BackendInfo.get_working_ends.sample
       else
         current_companion = env.get("current_companion").as(Int32)
       end
+      LOGGER.trace("Invidious companion: current_companion: #{current_companion}")
       response = COMPANION_POOL[current_companion].client &.post(endpoint, headers: headers, body: data.to_json)
       body = response.body
       if (response.status_code != 200)