terminal/kitty-gfx: add per-screen storage limit

src/terminal/kitty/graphics_image.zig

@@ -223,6 +223,13 @@ pub const LoadingImage = struct {
             log.warn("failed to calculate size for base64 data: {}", .{err});
             return error.InvalidData;
         };
+
+        // If our data would get too big, return an error
+        if (self.data.items.len + size > max_size) {
+            log.warn("image data too large max_size={}", .{max_size});
+            return error.InvalidData;
+        }
+
         try self.data.ensureUnusedCapacity(alloc, size);
 
         // We decode directly into the arraylist
@@ -355,6 +362,10 @@ pub const LoadingImage = struct {
         ) orelse return error.InvalidData;
         defer stb.stbi_image_free(data);
         const len: usize = @intCast(width * height * bpp);
+        if (len > max_size) {
+            log.warn("png image too large size={} max_size={}", .{ len, max_size });
+            return error.InvalidData;
+        }
 
         // Validate our bpp
         if (bpp != 3 and bpp != 4) return error.UnsupportedDepth;

src/terminal/kitty/graphics_storage.zig

@@ -41,6 +41,12 @@ pub const ImageStorage = struct {
     /// Non-null if there is an in-progress loading image.
     loading: ?*LoadingImage = null,
 
+    /// The total bytes of image data that have been loaded and the limit.
+    /// If the limit is reached, the oldest images will be evicted to make
+    /// space. Unused images take priority.
+    total_bytes: usize = 0,
+    total_limit: usize = 320 * 1000 * 1000, // 320MB
+
     pub fn deinit(self: *ImageStorage, alloc: Allocator) void {
         if (self.loading) |loading| loading.destroy(alloc);
 
@@ -54,6 +60,20 @@ pub const ImageStorage = struct {
     /// Add an already-loaded image to the storage. This will automatically
     /// free any existing image with the same ID.
     pub fn addImage(self: *ImageStorage, alloc: Allocator, img: Image) Allocator.Error!void {
+        // If the image itself is over the limit, then error immediately
+        if (img.data.len > self.total_limit) return error.OutOfMemory;
+
+        // If this would put us over the limit, then evict.
+        const total_bytes = self.total_bytes + img.data.len;
+        if (total_bytes > self.total_limit) {
+            const req_bytes = total_bytes - self.total_limit;
+            log.info("evicting images to make space for {} bytes", .{req_bytes});
+            if (!try self.evictImage(alloc, req_bytes)) {
+                log.warn("failed to evict enough images for required bytes", .{});
+                return error.OutOfMemory;
+            }
+        }
+
         // Do the gop op first so if it fails we don't get a partial state
         const gop = try self.images.getOrPut(alloc, img.id);
 
@@ -64,8 +84,13 @@ pub const ImageStorage = struct {
         }});
 
         // Write our new image
-        if (gop.found_existing) gop.value_ptr.deinit(alloc);
+        if (gop.found_existing) {
+            self.total_bytes -= gop.value_ptr.data.len;
+            gop.value_ptr.deinit(alloc);
+        }
+
         gop.value_ptr.* = img;
+        self.total_bytes += img.data.len;
 
         self.dirty = true;
     }
@@ -251,6 +276,7 @@ pub const ImageStorage = struct {
 
         // If we get here, we can delete the image.
         if (self.images.getEntry(image_id)) |entry| {
+            self.total_bytes -= entry.value_ptr.data.len;
             entry.value_ptr.deinit(alloc);
             self.images.removeByPtr(entry.key_ptr);
         }
@@ -278,6 +304,106 @@ pub const ImageStorage = struct {
         }
     }
 
+    /// Evict image to make space. This will evict the oldest image,
+    /// prioritizing unused images first, as recommended by the published
+    /// Kitty spec.
+    ///
+    /// This will evict as many images as necessary to make space for
+    /// req bytes.
+    fn evictImage(self: *ImageStorage, alloc: Allocator, req: usize) !bool {
+        assert(req <= self.total_limit);
+
+        // Ironically we allocate to evict. We should probably redesign the
+        // data structures to avoid this but for now allocating a little
+        // bit is fine compared to the megabytes we're looking to save.
+        const Candidate = struct {
+            id: u32,
+            time: std.time.Instant,
+            used: bool,
+        };
+
+        var candidates = std.ArrayList(Candidate).init(alloc);
+        defer candidates.deinit();
+
+        var it = self.images.iterator();
+        while (it.next()) |kv| {
+            const img = kv.value_ptr;
+
+            // This is a huge waste. See comment above about redesigning
+            // our data structures to avoid this. Eviction should be very
+            // rare though and we never have that many images/placements
+            // so hopefully this will last a long time.
+            const used = used: {
+                var p_it = self.placements.iterator();
+                while (p_it.next()) |p_kv| {
+                    if (p_kv.key_ptr.image_id == img.id) {
+                        break :used true;
+                    }
+                }
+
+                break :used false;
+            };
+
+            try candidates.append(.{
+                .id = img.id,
+                .time = img.transmit_time,
+                .used = used,
+            });
+        }
+
+        // Sort
+        std.mem.sortUnstable(
+            Candidate,
+            candidates.items,
+            {},
+            struct {
+                fn lessThan(
+                    ctx: void,
+                    lhs: Candidate,
+                    rhs: Candidate,
+                ) bool {
+                    _ = ctx;
+
+                    // If they're usage matches, then its based on time.
+                    if (lhs.used == rhs.used) return switch (lhs.time.order(rhs.time)) {
+                        .lt => true,
+                        .gt => false,
+                        .eq => lhs.id < rhs.id,
+                    };
+
+                    // If not used, then its a better candidate
+                    return !lhs.used;
+                }
+            }.lessThan,
+        );
+
+        // They're in order of best to evict.
+        var evicted: usize = 0;
+        for (candidates.items) |c| {
+            // Delete all the placements for this image and the image.
+            var p_it = self.placements.iterator();
+            while (p_it.next()) |entry| {
+                if (entry.key_ptr.image_id == c.id) {
+                    self.placements.removeByPtr(entry.key_ptr);
+                }
+            }
+
+            if (self.images.getEntry(c.id)) |entry| {
+                log.info("evicting image id={} bytes={}", .{ c.id, entry.value_ptr.data.len });
+
+                evicted += entry.value_ptr.data.len;
+                self.total_bytes -= entry.value_ptr.data.len;
+
+                entry.value_ptr.deinit(alloc);
+                self.images.removeByPtr(entry.key_ptr);
+
+                if (evicted > req) return true;
+            }
+        }
+
+        return false;
+    }
+
     /// Every placement is uniquely identified by the image ID and the
     /// placement ID. If an image ID isn't specified it is assumed to be 0.
     /// Likewise, if a placement ID isn't specified it is assumed to be 0.