From c012206f802074e5a27a3454db4f263305397b62 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 4b50171a..76dc7be7 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 94add5a8..e086ac3b 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