From 15086f4b6287bdc034aedae62101c3ae4a3b7f0b Mon Sep 17 00:00:00 2001
From: syeopite <syeopite@syeopite.dev>
Date: Tue, 3 Jun 2025 17:10:10 -0700
Subject: [PATCH] Improve cache size check to be more performant

Summing the sizes of each cached file every time is very inefficient.
Instead we can simply store the cache size in an constant and increase
it everytime a file is added into the cache.
---
 .../handlers/static_assets_handler_spec.cr    | 31 +++++++++++++++++++
 .../http_server/static_assets_handler.cr      |  7 +++--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/spec/http_server/handlers/static_assets_handler_spec.cr b/spec/http_server/handlers/static_assets_handler_spec.cr
index 4b50171ae..76dc7be77 100644
--- a/spec/http_server/handlers/static_assets_handler_spec.cr
+++ b/spec/http_server/handlers/static_assets_handler_spec.cr
@@ -198,5 +198,36 @@ Spectator.describe StaticAssetsHandler do
     end
   end
 
+  it "Will not cache additional files if the cache limit is reached" do
+    5.times do |times|
+      data = "a" * 1_000_000
+
+      make_temporary_file("test cache size limit #{times}") do |temporary_file, file_link|
+        cycle_temporary_file_contents(temporary_file, data) do
+          response = handle HTTP::Request.new("GET", file_link)
+          expect(response.status_code).to eq(200)
+          expect(response.body).to eq(data)
+        end
+
+        response = handle HTTP::Request.new("GET", file_link)
+        expect(response.status_code).to eq(200)
+        expect(response.body).to eq(data)
+      end
+    end
+
+    # Cache should be 5 mb so no more files will be cached.
+    make_temporary_file("test cache size limit uncached") do |temporary_file, file_link|
+      cycle_temporary_file_contents(temporary_file, "a") do
+        response = handle HTTP::Request.new("GET", file_link)
+        expect(response.status_code).to eq(200)
+        expect(response.body).to eq("a")
+      end
+
+      response = handle HTTP::Request.new("GET", file_link)
+      expect(response.status_code).to eq(200)
+      expect(response.body).to_not eq("a")
+    end
+  end
+
   after_each { Invidious::HttpServer::StaticAssetsHandler.clear_cache }
 end
diff --git a/src/invidious/http_server/static_assets_handler.cr b/src/invidious/http_server/static_assets_handler.cr
index 94add5a80..e086ac3b0 100644
--- a/src/invidious/http_server/static_assets_handler.cr
+++ b/src/invidious/http_server/static_assets_handler.cr
@@ -20,6 +20,7 @@ module Invidious::HttpServer
     end
 
     CACHE_LIMIT = 5_000_000 # 5MB
+    @@current_cache_size = 0
     @@cached_files = {} of Path => CachedFile
 
     # Returns metadata for the requested file
@@ -71,9 +72,8 @@ module Invidious::HttpServer
 
     # Writes file data to the cache
     private def flush_io_to_cache(io, file_path, file_info)
-      if @@cached_files.sum(&.[1].size) + file_info.size < CACHE_LIMIT
-        data_slice = io.to_slice
-        @@cached_files[file_path] = CachedFile.new(data_slice, file_info.size, file_info.modification_time)
+      if (@@current_cache_size += file_info.size) <= CACHE_LIMIT
+        @@cached_files[file_path] = CachedFile.new(io.to_slice, file_info.size, file_info.modification_time)
       end
     end
 
@@ -112,6 +112,7 @@ module Invidious::HttpServer
     #
     # This is only used in the specs to clear the cache before each handler test
     def self.clear_cache
+      @@current_cache_size = 0
       return @@cached_files.clear
     end
   end