From 8201f0dd0e77aa5034b1807c85c98a7ec755ceaf Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Tue, 23 Jul 2024 13:37:53 +0530 Subject: [PATCH] Move caching implementation to Python Less code, more performant since the cache is used from Python. And we can make the Go code a pure image format conversion filter. --- kitty/render_cache.py | 111 +++++++++++++++++++ kitty/utils.py | 39 ++----- kitty_tests/graphics.py | 34 ++++-- tools/cmd/tool/main.go | 4 +- tools/utils/images/convert.go | 97 +++++++++++++++++ tools/utils/images/render.go | 171 ------------------------------ tools/utils/images/render_test.go | 99 ----------------- 7 files changed, 246 insertions(+), 309 deletions(-) create mode 100644 kitty/render_cache.py create mode 100644 tools/utils/images/convert.go delete mode 100644 tools/utils/images/render.go delete mode 100644 tools/utils/images/render_test.go diff --git a/kitty/render_cache.py b/kitty/render_cache.py new file mode 100644 index 000000000..f89ef6cfa --- /dev/null +++ b/kitty/render_cache.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +# License: GPLv3 Copyright: 2024, Kovid Goyal + +import os +import time +from contextlib import closing, suppress +from typing import Iterator, Tuple + +from .constants import cache_dir, kitten_exe +from .utils import lock_file, unlock_file + + +class ImageRenderCache: + + lock_file_name = '.lock' + + def __init__(self, subdirname: str = 'rgba', max_entries: int = 32, cache_path: str = ''): + self.subdirname = subdirname + self.cache_path = cache_path + self.cache_dir = '' + self.max_entries = max_entries + + def ensure_subdir(self) -> None: + if not self.cache_dir: + x = os.path.abspath(os.path.join(self.cache_path or cache_dir(), self.subdirname)) + os.makedirs(x, exist_ok=True) + self.cache_dir = x + + def __enter__(self) -> None: + self.ensure_subdir() + self.lock_file = open(os.path.join(self.cache_dir, self.lock_file_name), 'wb') + lock_file(self.lock_file) + + def __exit__(self, *a: object) -> None: + with closing(self.lock_file): + unlock_file(self.lock_file) + + def entries(self) -> 'Iterator[os.DirEntry[str]]': + for x in os.scandir(self.cache_dir): + if x.name != self.lock_file_name: + yield x + + def prune_entries(self) -> None: + entries = list(self.entries()) + if len(entries) <= self.max_entries: + return + def sort_key(e: 'os.DirEntry[str]') -> float: + with suppress(OSError): + st = e.stat() + return st.st_mtime + return 0. + + entries.sort(key=sort_key, reverse=True) + for e in entries[self.max_entries:]: + with suppress(FileNotFoundError): + os.remove(e.path) + + def touch(self, path: str) -> None: + os.utime(path, follow_symlinks=False) + + def render_image(self, src_path: str, output_path: str) -> None: + import subprocess + with open(src_path, 'rb') as src, open(output_path, 'wb') as output: + cp = subprocess.run([kitten_exe(), '__convert_image__', 'RGBA'], stdin=src, stdout=output, stderr=subprocess.PIPE) + if cp.returncode != 0: + raise ValueError(f'Failed to convert path to RGBA data with error: {cp.stderr.decode("utf-8", "replace")}') + + def read_metadata(self, output_path: str) -> Tuple[int, int, int]: + with open(output_path, 'rb') as f: + header = f.read(8) + import struct + width, height = struct.unpack(' str: + from hashlib import sha256 + + with self: + output_name = sha256(src_path.encode()).hexdigest() + output_path = os.path.join(self.cache_dir, output_name) + src_info = os.stat(src_path) + with suppress(OSError), open(output_path, 'rb') as f: + dest_info = os.stat(f.fileno()) + if dest_info.st_size == src_info.st_size and dest_info.st_mtime >= src_info.st_mtime: + self.touch(output_path) + return output_path + + self.render_image(src_path, output_path) + self.prune_entries() + return output_path + + def __call__(self, src: str) -> Tuple[int, int, int]: + return self.read_metadata(self.render(src)) + + +class ImageRenderCacheForTesting(ImageRenderCache): + + def __init__(self, cache_path: str): + super().__init__(max_entries=2, cache_path=cache_path) + self.current_time = time.time_ns() + + def render_image(self, src_path: str, output_path: str) -> None: + super().render_image(src_path, output_path) + self.touch(output_path) + + def touch(self, path:str) -> None: + self.current_time += 3 * int(1e9) + os.utime(path, ns=(self.current_time, self.current_time)) + + +default_image_render_cache = ImageRenderCache() diff --git a/kitty/utils.py b/kitty/utils.py index dd1694811..c20b38533 100644 --- a/kitty/utils.py +++ b/kitty/utils.py @@ -34,7 +34,6 @@ from typing import ( from .constants import ( appname, - cache_dir, clear_handled_signals, config_dir, is_macos, @@ -1227,33 +1226,15 @@ def timed_debug_print(*a: Any, sep: str = ' ', end: str = '\n') -> None: _timed_debug_print(sep.join(map(str, a)) + end) -def cached_rgba_file_descriptor_for_image_path(path: str) -> Tuple[int, int, int]: - from hashlib import sha256 - try: - path = os.path.realpath(path, strict=True) - except TypeError: - path = os.path.realpath(path) - src_info = os.stat(path) - output_name = sha256(path.encode()).hexdigest() + '.rgba' - output_path = os.path.join(cache_dir(), 'rgba', output_name) +def lock_file(f: BinaryIO) -> None: + if not f.writable(): + raise ValueError('Cannot lock files not opened in writable mode') + import fcntl + fcntl.lockf(f, fcntl.LOCK_EX) - def read_data(f: BinaryIO) -> Tuple[int, int, int]: - header = f.read(8) - import struct - width, height = struct.unpack('= src_info.st_mtime: - return read_data(f) - - import subprocess - cp = subprocess.run([kitten_exe(), '__render_image__', path], capture_output=True) - if cp.returncode != 0: - raise ValueError(f'Failed to convert path to RGBA data with error: {cp.stderr.decode("utf-8", "replace")}') - ans = cp.stdout.decode().strip() - if ans != output_path: - raise ValueError(f'The two cache name algorithms dont agree for path: {path}\n{output_path} != {ans}') - with open(ans, 'rb') as f: - return read_data(f) +def unlock_file(f: BinaryIO) -> None: + if not f.writable(): + raise ValueError('Cannot unlock files not opened in writable mode') + import fcntl + fcntl.lockf(f, fcntl.LOCK_UN) diff --git a/kitty_tests/graphics.py b/kitty_tests/graphics.py index 2e21ebd7b..458286667 100644 --- a/kitty_tests/graphics.py +++ b/kitty_tests/graphics.py @@ -12,7 +12,6 @@ from dataclasses import dataclass from io import BytesIO from kitty.fast_data_types import base64_decode, base64_encode, has_avx2, has_sse4_2, load_png_data, shm_unlink, shm_write, test_xor64 -from kitty.utils import cached_rgba_file_descriptor_for_image_path from . import BaseTest, parse_bytes @@ -1244,16 +1243,35 @@ class TestGraphics(BaseTest): @unittest.skipIf(Image is None, 'PIL not available, skipping PNG tests') def test_cached_rgba_conversion(self): + from kitty.render_cache import ImageRenderCacheForTesting w, h = 5, 3 rgba_data = byte_block(w * h * 4) img = Image.frombytes('RGBA', (w, h), rgba_data) buf = BytesIO() img.save(buf, 'PNG') png_data = buf.getvalue() - with tempfile.NamedTemporaryFile(suffix='.png') as png: - png.write(png_data) - png.flush() - os.fsync(png.fileno()) - qw, qh, fd = cached_rgba_file_descriptor_for_image_path(png.name) - os.close(fd) - self.ae((qw, qh), (w, h)) + with tempfile.TemporaryDirectory() as cache_path: + irc = ImageRenderCacheForTesting(cache_path) + srcs, outputs = [], [] + for i in range(2 * irc.max_entries): + with open(os.path.join(cache_path, f'{i}.png'), 'wb') as f: + f.write(png_data) + srcs.append(f.name) + outputs.append(irc.render(f.name)) + entries = list(irc.entries()) + self.assertLessEqual(len(entries), irc.max_entries) + remaining_outputs = outputs[-irc.max_entries:] + for x in remaining_outputs: + self.assertTrue(os.path.exists(x)) + for x in outputs[:-irc.max_entries]: + self.assertFalse(os.path.exists(x)) + self.assertLess(os.path.getmtime(remaining_outputs[0]), os.path.getmtime(remaining_outputs[1])) + remaining_srcs = srcs[-irc.max_entries:] + self.ae(irc.render(remaining_srcs[0]), remaining_outputs[0]) + self.assertGreater(os.path.getmtime(remaining_outputs[0]), os.path.getmtime(remaining_outputs[1])) + + width, height, fd = irc(remaining_srcs[-1]) + with open(fd, 'rb') as f: + self.ae((width, height), (w, h)) + f.seek(8) + self.ae(rgba_data, f.read()) diff --git a/tools/cmd/tool/main.go b/tools/cmd/tool/main.go index ac925c584..e5f44ef5c 100644 --- a/tools/cmd/tool/main.go +++ b/tools/cmd/tool/main.go @@ -104,8 +104,8 @@ func KittyToolEntryPoints(root *cli.Command) { return confirm_and_run_shebang(args) }, }) - // __render_image__ - images.RenderEntryPoint(root) + // __convert_image__ + images.ConvertEntryPoint(root) // __generate_man_pages__ root.AddSubCommand(&cli.Command{ Name: "__generate_man_pages__", diff --git a/tools/utils/images/convert.go b/tools/utils/images/convert.go new file mode 100644 index 000000000..56a79a585 --- /dev/null +++ b/tools/utils/images/convert.go @@ -0,0 +1,97 @@ +package images + +import ( + "bytes" + "encoding/binary" + "fmt" + "image" + "io" + "os" + "strings" + + "kitty/tools/cli" + "kitty/tools/utils" +) + +var _ = fmt.Print + +func encode_rgba(output io.Writer, img image.Image) (err error) { + var final_img *image.NRGBA + switch img.(type) { + case *image.NRGBA: + final_img = img.(*image.NRGBA) + default: + b := img.Bounds() + final_img = image.NewNRGBA(image.Rect(0, 0, b.Dx(), b.Dy())) + ctx := Context{} + ctx.PasteCenter(final_img, img, nil) + } + b := final_img.Bounds() + header := make([]byte, 8) + var width = utils.Abs(b.Dx()) + var height = utils.Abs(b.Dy()) + binary.LittleEndian.PutUint32(header, uint32(width)) + binary.LittleEndian.PutUint32(header[4:], uint32(height)) + readers := []io.Reader{bytes.NewReader(header)} + stride := 4 * width + + if final_img.Stride == stride { + readers = append(readers, bytes.NewReader(final_img.Pix)) + } else { + p := final_img.Pix + for y := 0; y < b.Dy(); y++ { + readers = append(readers, bytes.NewReader(p[:min(stride, len(p))])) + p = p[final_img.Stride:] + } + } + _, err = io.Copy(output, io.MultiReader(readers...)) + return +} + +func convert_image(input io.ReadSeeker, output io.Writer, format string) (err error) { + image_data, err := OpenNativeImageFromReader(input) + if err != nil { + return err + } + if len(image_data.Frames) == 0 { + return fmt.Errorf("Image has no frames") + } + img := image_data.Frames[0].Img + switch strings.ToUpper(format) { + case "RGBA": + return encode_rgba(output, img) + case "JPEG", "JPG": + return Encode(output, img, "image/jpeg") + case "PNG": + return Encode(output, img, "image/png") + case "GIF": + return Encode(output, img, "image/gif") + case "BMP": + return Encode(output, img, "image/bmp") + case "TIFF": + return Encode(output, img, "image/tiff") + } + return +} + +func ConvertEntryPoint(root *cli.Command) { + root.AddSubCommand(&cli.Command{ + Name: "__convert_image__", + Hidden: true, + OnlyArgsAllowed: true, + Run: func(cmd *cli.Command, args []string) (rc int, err error) { + format := "RGBA" + if len(args) > 0 { + format = args[0] + } + buf := bytes.NewBuffer(make([]byte, 0, 1024*1024)) + if _, err = io.Copy(buf, os.Stdin); err != nil { + return 1, err + } + if err = convert_image(bytes.NewReader(buf.Bytes()), os.Stdout, format); err != nil { + rc = 1 + } + return + }, + }) +} diff --git a/tools/utils/images/render.go b/tools/utils/images/render.go deleted file mode 100644 index 20cc6bac2..000000000 --- a/tools/utils/images/render.go +++ /dev/null @@ -1,171 +0,0 @@ -package images - -import ( - "bytes" - "crypto/sha256" - "encoding/binary" - "errors" - "fmt" - "image" - "io" - "io/fs" - "os" - "path/filepath" - "time" - - "kitty/tools/cli" - "kitty/tools/utils" -) - -var _ = fmt.Print - -func convert_and_save_as_rgba_data(input_path, output_path string, perm os.FileMode) (err error) { - f, err := os.Open(input_path) - if err != nil { - return err - } - defer f.Close() - image_data, err := OpenNativeImageFromReader(f) - if err != nil { - return err - } - if len(image_data.Frames) == 0 { - return fmt.Errorf("Image at %s has no frames", input_path) - } - img := image_data.Frames[0].Img - var final_img *image.NRGBA - switch img.(type) { - case *image.NRGBA: - final_img = img.(*image.NRGBA) - default: - b := img.Bounds() - final_img = image.NewNRGBA(image.Rect(0, 0, b.Dx(), b.Dy())) - ctx := Context{} - ctx.PasteCenter(final_img, img, nil) - } - b := final_img.Bounds() - header := make([]byte, 8) - var width = utils.Abs(b.Dx()) - var height = utils.Abs(b.Dy()) - binary.LittleEndian.PutUint32(header, uint32(width)) - binary.LittleEndian.PutUint32(header[4:], uint32(height)) - readers := []io.Reader{bytes.NewReader(header)} - stride := 4 * width - - if final_img.Stride == stride { - readers = append(readers, bytes.NewReader(final_img.Pix)) - } else { - p := final_img.Pix - for y := 0; y < b.Dy(); y++ { - readers = append(readers, bytes.NewReader(p[:min(stride, len(p))])) - p = p[final_img.Stride:] - } - } - return utils.AtomicUpdateFile(output_path, io.MultiReader(readers...), perm) -} - -var now_implementation = time.Now -var chmtime_after_creation = false - -func prune_cache(cdir string, max_entries int) error { - entries, err := os.ReadDir(cdir) - if err != nil { - return err - } - if len(entries) <= max_entries { - return nil - } - now := now_implementation() - epoch := time.Unix(0, 0) - entries = utils.SortWithKey(entries, func(a fs.DirEntry) time.Duration { - if info, err := a.Info(); err == nil { - return now.Sub(info.ModTime()) - } - return now.Sub(epoch) - }) - for _, x := range entries[max_entries:] { - if err = os.Remove(filepath.Join(cdir, x.Name())); err != nil { - return err - } - } - return nil -} - -func render_image(src_path, cdir string, max_cache_entries int) (output_path string, err error) { - if src_path, err = filepath.EvalSymlinks(src_path); err != nil { - return - } - if src_path, err = filepath.Abs(src_path); err != nil { - return - } - lock_file := filepath.Join(cdir, "lock") - lockf, err := os.OpenFile(lock_file, os.O_CREATE|os.O_RDWR, 0600) - defer lockf.Close() - if err != nil { - return - } - if err = utils.LockFileExclusive(lockf); err != nil { - return "", fmt.Errorf("Failed to lock cache file %s with error: %w", lock_file, err) - } - defer func() { - utils.UnlockFile(lockf) - }() - output_path = filepath.Join(cdir, fmt.Sprintf("%x", sha256.Sum256([]byte(src_path)))) + ".rgba" - needs_update := true - input_info, err := os.Stat(src_path) - if err != nil { - return - } - output_info, err := os.Stat(output_path) - if err == nil { - needs_update = input_info.Size() != output_info.Size() || input_info.ModTime().After(output_info.ModTime()) - } else { - if !errors.Is(err, fs.ErrNotExist) { - return - } - } - if needs_update { - if err = convert_and_save_as_rgba_data(src_path, output_path, 0600); err != nil { - return - } - if chmtime_after_creation { - n := now_implementation() - if err = os.Chtimes(output_path, n, n); err != nil { - return - } - } - if err = prune_cache(cdir, max_cache_entries); err != nil { - return - } - } else { - n := now_implementation() - if err = os.Chtimes(output_path, n, n); err != nil { - return - } - } - return -} - -func RenderEntryPoint(root *cli.Command) { - root.AddSubCommand(&cli.Command{ - Name: "__render_image__", - Hidden: true, - OnlyArgsAllowed: true, - Run: func(cmd *cli.Command, args []string) (rc int, err error) { - if len(args) != 1 { - return 1, fmt.Errorf("Usage: render input_image_path") - } - cdir := utils.CacheDir() - cdir = filepath.Join(cdir, "rgba") - if err = os.MkdirAll(cdir, 0755); err != nil { - return 1, err - } - if output_path, err := render_image(args[0], cdir, 32); err != nil { - return 1, err - } else { - fmt.Println(output_path) - } - return - }, - }) -} diff --git a/tools/utils/images/render_test.go b/tools/utils/images/render_test.go deleted file mode 100644 index f4eec6894..000000000 --- a/tools/utils/images/render_test.go +++ /dev/null @@ -1,99 +0,0 @@ -package images - -import ( - "encoding/binary" - "fmt" - "os" - "path/filepath" - "testing" - "time" - - "github.com/google/go-cmp/cmp" -) - -var _ = fmt.Print - -var one_pixel_gray_png = []byte{137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 0, 0, 0, 0, 58, 126, 155, 85, 0, 0, 0, 10, 73, 68, 65, 84, 120, 1, 99, 176, 5, 0, 0, 63, 0, 62, 18, 174, 200, 16, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130} - -func TestRenderCache(t *testing.T) { - chmtime_after_creation = true - epoch := time.Now() - now_implementation = func() time.Time { - epoch = epoch.Add(3 * time.Second) - return epoch - } - defer func() { - chmtime_after_creation = false - now_implementation = time.Now - }() - tmp := t.TempDir() - cdir := filepath.Join(tmp, "cache") - if err := os.Mkdir(cdir, 0777); err != nil { - t.Fatal(err) - } - srcs := make([]string, 0, 5) - outputs := make([]string, 0, 5) - const max_cache_entries = 2 - for i := range max_cache_entries * 2 { - name := fmt.Sprintf(`%d.png`, i) - src_path := filepath.Join(tmp, name) - srcs = append(srcs, src_path) - if err := os.WriteFile(src_path, one_pixel_gray_png, 0644); err != nil { - t.Fatal(err) - } - output_path, err := render_image(src_path, cdir, max_cache_entries) - if err != nil { - t.Fatal(err) - } - outputs = append(outputs, output_path) - if entries, err := os.ReadDir(cdir); err != nil { - t.Fatal(err) - } else if len(entries) > max_cache_entries { - t.Fatalf("Too many files in cache dir %d > %d", len(entries), max_cache_entries) - } - } - exists := func(path string) bool { - _, err := os.Stat(path) - return err == nil - } - mtime := func(path string) time.Time { - ans, err := os.Stat(path) - if err != nil { - t.Fatal(err) - } - return ans.ModTime() - } - for i, x := range outputs[len(outputs)-max_cache_entries:] { - if !exists(x) { - t.Fatalf("The %d cache entry does not exist", max_cache_entries+i) - } - } - o := outputs[len(outputs)-max_cache_entries:] - if mtime(o[0]).After(mtime(o[1])) { - t.Fatalf("The mtimes are not monotonic") - } - s := srcs[len(srcs)-max_cache_entries:] - output_path, err := render_image(s[0], cdir, max_cache_entries) - if err != nil { - t.Fatal(err) - } - if output_path != o[0] { - t.Fatalf("Output path change on rerun") - } - if mtime(o[1]).After(mtime(o[0])) || mtime(o[1]).Equal(mtime(o[0])) { - t.Fatalf("The mtime was not updated") - } - data, err := os.ReadFile(output_path) - if err != nil { - t.Fatal(err) - } - if len(data) != 12 { - t.Fatalf("unexpected data length: %d != %d", len(data), 12) - } - if width, height := binary.LittleEndian.Uint32(data), binary.LittleEndian.Uint32(data[4:]); width != 1 || height != 1 { - t.Fatalf("unexpected dimensions: %dx%d", width, height) - } - if diff := cmp.Diff(data[8:], []byte{61, 61, 61, 255}); diff != "" { - t.Fatalf("unexpected pixel: %s", diff) - } -}