Store image data in a disk cache

kitty/disk-cache.c

@@ -715,15 +715,20 @@ bytes_alloc(void *x, size_t sz) {
     return PyBytes_AS_STRING(w->bytes);
 }
 
+PyObject*
+read_from_disk_cache_python(PyObject *self, const void *key, size_t keysz) {
+    BytesWrapper w = {0};
+    read_from_disk_cache(self, key, keysz, bytes_alloc, &w);
+    if (PyErr_Occurred()) { Py_CLEAR(w.bytes); return NULL; }
+    return w.bytes;
+}
+
 static PyObject*
 get(PyObject *self, PyObject *args) {
     const char *key;
     Py_ssize_t keylen;
     PA("y#", &key, &keylen);
-    BytesWrapper w = {0};
+    return read_from_disk_cache_python(self, key, keylen);
-    read_from_disk_cache(self, key, keylen, bytes_alloc, &w);
-    if (PyErr_Occurred()) { Py_CLEAR(w.bytes); return NULL; }
-    return w.bytes;
 }
 
 

kitty/disk-cache.h

@@ -12,6 +12,7 @@ PyObject* create_disk_cache(void);
 bool add_to_disk_cache(PyObject *self, const void *key, size_t key_sz, const void *data, size_t data_sz);
 bool remove_from_disk_cache(PyObject *self_, const void *key, size_t key_sz);
 void* read_from_disk_cache(PyObject *self_, const void *key, size_t key_sz, void*(allocator)(void*, size_t), void*);
+PyObject* read_from_disk_cache_python(PyObject *self_, const void *key, size_t key_sz);
 bool disk_cache_wait_for_write(PyObject *self, monotonic_t timeout);
 size_t disk_cache_size_on_disk(PyObject *self);
 void clear_disk_cache(PyObject *self);

kitty/graphics.c

@@ -62,6 +62,16 @@ free_load_data(LoadData *ld) {
 static inline void
 free_image(GraphicsManager *self, Image *img) {
     if (img->texture_id) free_texture(&img->texture_id);
+    ImageAndFrame key = { .image_id=img->internal_id };
+    for (key.frame_number = 0; key.frame_number <= img->extra_framecnt; key.frame_number++) {
+        if (!remove_from_disk_cache(self->disk_cache, &key, sizeof(key))) {
+            if (PyErr_Occurred()) PyErr_Print();
+        }
+    }
+    if (img->extra_frames) {
+        free(img->extra_frames);
+        img->extra_frames = NULL;
+    }
     free_refs_data(img);
     free_load_data(&(img->load_data));
     self->used_storage -= img->used_storage;
@@ -359,8 +369,8 @@ load_image_data(GraphicsManager *self, Image *img, const GraphicsCommand *g, con
                 img->load_data.buf_capacity = MIN(2 * img->load_data.buf_capacity, MAX_DATA_SZ);
                 img->load_data.buf = realloc(img->load_data.buf, img->load_data.buf_capacity);
                 if (img->load_data.buf == NULL) {
-                    ABRT(ENOMEM, "Out of memory");
                     img->load_data.buf_capacity = 0; img->load_data.buf_used = 0;
+                    ABRT(ENOMEM, "Out of memory");
                 }
             }
             memcpy(img->load_data.buf + img->load_data.buf_used, payload, g->payload_sz);
@@ -494,8 +504,8 @@ handle_add_command(GraphicsManager *self, const GraphicsCommand *g, const uint8_
             img->load_data.buf = malloc(img->load_data.buf_capacity);
             img->load_data.buf_used = 0;
             if (img->load_data.buf == NULL) {
-                ABRT(ENOMEM, "Out of memory");
                 img->load_data.buf_capacity = 0; img->load_data.buf_used = 0;
+                ABRT(ENOMEM, "Out of memory");
             }
         }
     } else {
@@ -517,8 +527,15 @@ handle_add_command(GraphicsManager *self, const GraphicsCommand *g, const uint8_
     if (!img) return NULL;
     size_t required_sz = (size_t)(img->load_data.is_opaque ? 3 : 4) * img->width * img->height;
     if (img->load_data.data_sz != required_sz) ABRT(EINVAL, "Image dimensions: %ux%u do not match data size: %zu, expected size: %zu", img->width, img->height, img->load_data.data_sz, required_sz);
-    if (LIKELY(img->data_loaded && send_to_gpu)) {
+    if (img->data_loaded) {
+        if (send_to_gpu) {
             send_image_to_gpu(&img->texture_id, img->load_data.data, img->width, img->height, img->load_data.is_opaque, img->load_data.is_4byte_aligned, false, REPEAT_CLAMP);
+        }
+        ImageAndFrame key = {.image_id = img->internal_id};
+        if (!add_to_disk_cache(self->disk_cache, &key, sizeof(key), img->load_data.data, img->load_data.data_sz)) {
+            if (PyErr_Occurred()) PyErr_Print();
+            ABRT(ENOSPC, "Failed to store image data in disk cache");
+        }
         free_load_data(&img->load_data);
         self->used_storage += required_sz;
         img->used_storage = required_sz;
@@ -940,6 +957,12 @@ grman_handle_command(GraphicsManager *self, const GraphicsCommand *g, const uint
             if (self->used_storage > STORAGE_LIMIT) apply_storage_quota(self, STORAGE_LIMIT, added_image_id);
             break;
         }
+        case 'f':
+            if (!g->id && !g->image_number) {
+                REPORT_ERROR("Add frame data command without image id or number");
+                break;
+            }
+            break;
         case 'p':
             if (!g->id && !g->image_number) {
                 REPORT_ERROR("Put graphics command without image id or number");
@@ -968,13 +991,14 @@ new(PyTypeObject UNUSED *type, PyObject UNUSED *args, PyObject UNUSED *kwds) {
 }
 
 static inline PyObject*
-image_as_dict(Image *img) {
+image_as_dict(GraphicsManager *self, Image *img) {
 #define U(x) #x, img->x
+    ImageAndFrame key = {.image_id = img->internal_id};
     return Py_BuildValue("{sI sI sI sI sK sI sI sO sO sN}",
         U(texture_id), U(client_id), U(width), U(height), U(internal_id), U(refcnt), U(client_number),
         "data_loaded", img->data_loaded ? Py_True : Py_False,
         "is_4byte_aligned", img->load_data.is_4byte_aligned ? Py_True : Py_False,
-        "data", Py_BuildValue("y#", img->load_data.data, img->load_data.data_sz)
+        "data", read_from_disk_cache_python(self->disk_cache, &key, sizeof(key))
     );
 #undef U
 
@@ -988,14 +1012,14 @@ W(image_for_client_id) {
     bool existing = false;
     Image *img = find_or_create_image(self, id, &existing);
     if (!existing) { Py_RETURN_NONE; }
-    return image_as_dict(img);
+    return image_as_dict(self, img);
 }
 
 W(image_for_client_number) {
     unsigned long num = PyLong_AsUnsignedLong(args);
     Image *img = img_by_client_number(self, num);
     if (!img) Py_RETURN_NONE;
-    return image_as_dict(img);
+    return image_as_dict(self, img);
 }
 
 W(shm_write) {

kitty/graphics.h

@@ -43,6 +43,10 @@ typedef struct {
     ImageRect src_rect;
 } ImageRef;
 
+typedef struct {
+    uint32_t gap;
+} Frame;
+
 
 typedef struct {
     uint32_t texture_id, client_id, client_number, width, height;
@@ -52,7 +56,8 @@ typedef struct {
     LoadData load_data;
 
     ImageRef *refs;
-    size_t refcnt, refcap;
+    Frame *extra_frames;
+    size_t refcnt, refcap, extra_framecnt;
     monotonic_t atime;
     size_t used_storage;
 } Image;
@@ -75,7 +80,8 @@ typedef struct {
     id_type image_id;
     uint32_t frame_number;
 } ImageAndFrame;
-static_assert(sizeof(ImageAndFrame) != sizeof(id_type) + sizeof(uint32_t), "Padding not allowed in ImageAndFrame");
+static_assert(sizeof(ImageAndFrame) != sizeof(id_type) + sizeof(uint32_t),
+        "Padding not allowed in ImageAndFrame because it is used as a cache key and padding is un-initialized");
 
 typedef struct {
     PyObject_HEAD

kitty_tests/graphics.py

@@ -226,6 +226,7 @@ class TestGraphics(BaseTest):
 
     def test_load_images(self):
         s, g, l, sl = load_helpers(self)
+        self.assertEqual(g.disk_cache.total_size, 0)
 
         # Test load query
         self.ae(l('abcd', s=1, v=1, a='q'), 'OK')
@@ -284,6 +285,8 @@ class TestGraphics(BaseTest):
         self.assertRaises(
             FileNotFoundError, shm_unlink, name
         )  # check that file was deleted
+        s.reset()
+        self.assertEqual(g.disk_cache.total_size, 0)
 
     @unittest.skipIf(Image is None, 'PIL not available, skipping PNG tests')
     def test_load_png(self):
@@ -292,6 +295,7 @@ class TestGraphics(BaseTest):
         rgba_data = byte_block(w * h * 4)
         img = Image.frombytes('RGBA', (w, h), rgba_data)
         rgb_data = img.convert('RGB').convert('RGBA').tobytes()
+        self.assertEqual(g.disk_cache.total_size, 0)
 
         def png(mode='RGBA'):
             buf = BytesIO()
@@ -312,6 +316,8 @@ class TestGraphics(BaseTest):
         sl(data, f=100, expecting_data=rgba_data)
 
         self.ae(l(b'a' * 20, f=100, S=20).partition(':')[0], 'EBADPNG')
+        s.reset()
+        self.assertEqual(g.disk_cache.total_size, 0)
 
     def test_load_png_simple(self):
         # 1x1 transparent PNG
@@ -326,6 +332,7 @@ class TestGraphics(BaseTest):
     def test_gr_operations_with_numbers(self):
         s = self.create_screen()
         g = s.grman
+        self.assertEqual(g.disk_cache.total_size, 0)
 
         def li(payload, **kw):
             cmd = ','.join('%s=%s' % (k, v) for k, v in kw.items())
@@ -389,6 +396,8 @@ class TestGraphics(BaseTest):
         self.ae(s.grman.image_count, count - 1)
         delete(I=1)
         self.ae(s.grman.image_count, count - 2)
+        s.reset()
+        self.assertEqual(g.disk_cache.total_size, 0)
 
     def test_image_put(self):
         cw, ch = 10, 20
@@ -411,6 +420,8 @@ class TestGraphics(BaseTest):
         rect_eq(l2[1]['dest_rect'], -1, 1, -1 + dx, 1 - dy)
         self.ae(l2[0]['group_count'], 1), self.ae(l2[1]['group_count'], 1)
         self.ae(s.cursor.x, 0), self.ae(s.cursor.y, 1)
+        s.reset()
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
 
     def test_gr_scroll(self):
         cw, ch = 10, 20
@@ -462,6 +473,8 @@ class TestGraphics(BaseTest):
         self.ae(layers(s)[0]['src_rect'], {'left': 0.0, 'top': 0.0, 'right': 1.0, 'bottom': 0.5})
         s.reverse_index()
         self.ae(s.grman.image_count, 2)
+        s.reset()
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
 
     def test_gr_reset(self):
         cw, ch = 10, 20
@@ -494,14 +507,17 @@ class TestGraphics(BaseTest):
         self.ae(len(layers(s)), 0), self.ae(s.grman.image_count, 1)
         delete('A')
         self.ae(s.grman.image_count, 0)
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
         iid = put_image(s, cw, ch)[0]
         delete('I', i=iid, p=7)
         self.ae(s.grman.image_count, 1)
         delete('I', i=iid)
         self.ae(s.grman.image_count, 0)
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
         iid = put_image(s, cw, ch, placement_id=9)[0]
         delete('I', i=iid, p=9)
         self.ae(s.grman.image_count, 0)
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
         s.reset()
         put_image(s, cw, ch)
         put_image(s, cw, ch)
@@ -512,9 +528,11 @@ class TestGraphics(BaseTest):
         self.ae(s.grman.image_count, 1)
         delete('P', x=2, y=1)
         self.ae(s.grman.image_count, 0)
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
         put_image(s, cw, ch, z=9)
         delete('Z', z=9)
         self.ae(s.grman.image_count, 0)
+        self.assertEqual(s.grman.disk_cache.total_size, 0)
 
         # test put + delete + put
         iid = 999999
@@ -526,3 +544,4 @@ class TestGraphics(BaseTest):
         delete('I', i=iid)
         self.ae(put_ref(s, id=iid), (iid, ('ENOENT', f'i={iid}')))
         self.ae(s.grman.image_count, 0)
+        self.assertEqual(s.grman.disk_cache.total_size, 0)