diff --git a/spec/invidious/subtitle_cache_spec.cr b/spec/invidious/subtitle_cache_spec.cr index cf6b3fe5f..4a7727ab2 100644 --- a/spec/invidious/subtitle_cache_spec.cr +++ b/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 diff --git a/src/invidious/helpers/subtitle_cache.cr b/src/invidious/helpers/subtitle_cache.cr index 311766df5..2eff250b6 100644 --- a/src/invidious/helpers/subtitle_cache.cr +++ b/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 diff --git a/src/invidious/routes/companion.cr b/src/invidious/routes/companion.cr index 8b7be7db5..fe967a2ce 100644 --- a/src/invidious/routes/companion.cr +++ b/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| diff --git a/src/invidious/yt_backend/connection_pool.cr b/src/invidious/yt_backend/connection_pool.cr index 3fc450b59..cd5788841 100644 --- a/src/invidious/yt_backend/connection_pool.cr +++ b/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