fix(companion): bound pending caption fetches

spec/invidious/subtitle_cache_spec.cr

@@ -137,6 +137,36 @@ Spectator.describe Invidious::SubtitleCache do
       expect(results.count("hit")).to eq(4)
     end
 
+    it "bounds distinct pending fetches" do
+      cache = Invidious::SubtitleCache.new(max_entries: 2)
+      started = ::Channel(String).new
+      release = ::Channel(Nil).new
+      done = ::Channel(String).new
+
+      3.times do |i|
+        spawn do
+          key = "video:pending-#{i}"
+          result = cache.get_or_fetch(key) do
+            started.send(key)
+            release.receive
+            subtitle_response(200, "text/vtt", "WEBVTT\n")
+          end
+          done.send(result.cache_status)
+        end
+      end
+
+      3.times { started.receive }
+      expect(cache.in_flight_size).to eq(2)
+
+      3.times { release.send(nil) }
+      results = Array(String).new
+      3.times { results << done.receive }
+
+      expect(results.count("miss")).to eq(2)
+      expect(results.count("bypass")).to eq(1)
+      expect(cache.in_flight_size).to eq(0)
+    end
+
     it "evicts oldest entries when max_entries is exceeded (LRU)" do
       cache = Invidious::SubtitleCache.new(max_entries: 2)
       vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nLine"
@@ -230,6 +260,11 @@ Spectator.describe Invidious::SubtitleCache do
       oversized_result = Invidious::SubtitleCache.read_limited_body(IO::Memory.new(oversized_body))
       expect(oversized_result.oversized).to be_true
       expect(oversized_result.body).to be_nil
+
+      response = Invidious::SubtitleCache.oversized_caption_response
+      expect(response.status_code).to eq(413)
+      expect(response.headers["Cache-Control"]).to eq("no-store")
+      expect(response.body).to eq("Caption response exceeds size limit\n")
     end
   end
 end

src/invidious/helpers/subtitle_cache.cr

@@ -43,14 +43,17 @@ module Invidious
       end
     end
 
-    DEFAULT_MAX_ENTRIES = 256
-    DEFAULT_MAX_BYTES   = 64 * 1024 * 1024 # 64 MiB
-    DEFAULT_TTL         = 6.hours          # 21600 seconds
-    MAX_ENTRY_BYTES     = 2 * 1024 * 1024  # 2 MiB
-    READ_CHUNK_BYTES    = 64 * 1024
+    DEFAULT_MAX_ENTRIES  = 256
+    DEFAULT_MAX_BYTES    = 64 * 1024 * 1024 # 64 MiB
+    DEFAULT_TTL          = 6.hours          # 21600 seconds
+    MAX_ENTRY_BYTES      = 2 * 1024 * 1024  # 2 MiB
+    READ_CHUNK_BYTES     = 64 * 1024
+    MAX_IN_FLIGHT_KEYS   = 64
+    MAX_WAITING_FETCHERS = 64
 
     getter max_entries : Int32
     getter max_bytes : Int32
+    getter max_in_flight : Int32
     getter ttl : Time::Span
     getter total_bytes : Int32
 
@@ -61,6 +64,8 @@ module Invidious
       @total_bytes = 0
       @mutex = Mutex.new
       @in_flight = Hash(String, Array(::Channel(FetchResult))).new
+      @max_in_flight = {@max_entries, MAX_IN_FLIGHT_KEYS}.min
+      @max_in_flight = 1 if @max_in_flight <= 0
     end
 
     def self.valid_vtt?(body : String) : Bool
@@ -76,6 +81,14 @@ module Invidious
       [video_id, label, lang, tlang].map { |value| Base64.urlsafe_encode(value) }.join('|')
     end
 
+    def self.oversized_caption_response : Response
+      headers = HTTP::Headers.new
+      headers["Access-Control-Allow-Origin"] = "*"
+      headers["Cache-Control"] = "no-store"
+      headers["Content-Type"] = "text/plain; charset=utf-8"
+      Response.new(413, headers, "Caption response exceeds size limit\n")
+    end
+
     def self.read_limited_body(input : IO, limit : Int32 = MAX_ENTRY_BYTES) : LimitedBody
       output = IO::Memory.new
       buffer = Bytes.new(READ_CHUNK_BYTES)
@@ -96,6 +109,10 @@ module Invidious
       @mutex.synchronize { @entries.size }
     end
 
+    def in_flight_size : Int32
+      @mutex.synchronize { @in_flight.size }
+    end
+
     def clear : Nil
       @mutex.synchronize do
         @entries.clear
@@ -159,6 +176,7 @@ module Invidious
 
     def get_or_fetch(key : String, &fetch_block : -> Response?) : FetchResult
       wait_ch : ::Channel(FetchResult)? = nil
+      direct_fetch = false
 
       @mutex.synchronize do
         if entry = get_internal(key)
@@ -166,11 +184,19 @@ module Invidious
         end
 
         if waiting_list = @in_flight[key]?
-          ch = ::Channel(FetchResult).new(1)
-          waiting_list << ch
-          wait_ch = ch
+          if waiting_list.size < MAX_WAITING_FETCHERS
+            ch = ::Channel(FetchResult).new(1)
+            waiting_list << ch
+            wait_ch = ch
+          else
+            direct_fetch = true
+          end
         else
-          @in_flight[key] = Array(::Channel(FetchResult)).new
+          if @in_flight.size < @max_in_flight
+            @in_flight[key] = Array(::Channel(FetchResult)).new
+          else
+            direct_fetch = true
+          end
         end
       end
 
@@ -178,6 +204,15 @@ module Invidious
         return ch.receive
       end
 
+      if direct_fetch
+        response = begin
+          fetch_block.call
+        rescue
+          nil
+        end
+        return FetchResult.new(nil, response, "bypass")
+      end
+
       # Primary fetcher for this key
       response = begin
         fetch_block.call

src/invidious/routes/companion.cr

@@ -110,7 +110,8 @@ module Invidious::Routes::Companion
         limited_body = Invidious::SubtitleCache.read_limited_body(resp.body_io)
         if limited_body.oversized
           resp.body_io.close
-          return oversized_caption_response
+          wrapper.close
+          return Invidious::SubtitleCache.oversized_caption_response
         end
 
         body = limited_body.body
@@ -126,14 +127,6 @@ module Invidious::Routes::Companion
     nil
   end
 
-  private def self.oversized_caption_response
-    headers = HTTP::Headers.new
-    headers["Access-Control-Allow-Origin"] = "*"
-    headers["Cache-Control"] = "no-store"
-    headers["Content-Type"] = "text/plain; charset=utf-8"
-    Invidious::SubtitleCache::Response.new(413, headers, "Caption response exceeds size limit\n")
-  end
-
   private def self.proxy_companion(env, response)
     env.response.status_code = response.status_code
     response.headers.each do |key, value|

src/invidious/yt_backend/connection_pool.cr

@@ -51,9 +51,10 @@ end
 # This is used as the resource for the `CompanionPool` as to allow the ability to
 # proxy the requests to Invidious companion from Invidious directly.
 # Instead of setting up routes in a reverse proxy.
-struct CompanionWrapper
+class CompanionWrapper
   property client : HTTP::Client
   property companion : Config::CompanionConfig
+  @closed = false
 
   def initialize(companion : Config::CompanionConfig)
     @companion = companion
@@ -61,8 +62,13 @@ struct CompanionWrapper
   end
 
   def close
+    @closed = true
     @client.close
   end
+
+  def closed? : Bool
+    @closed
+  end
 end
 
 struct CompanionConnectionPool