From dbf318a20e282bd473d0ba9106f4ee2d98e7b151 Mon Sep 17 00:00:00 2001 From: Yuzu Date: Tue, 14 Jul 2026 05:49:10 +0700 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=80=20perf:=20Performance=20and=20Memo?= =?UTF-8?q?ry=20optimization?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 15 ++-- README.md | 17 ++-- src/app_test.go | 44 +++++++++- src/config.go | 4 +- src/frameformat.go | 2 +- src/frames.go | 213 ++++++++++++++++++++++++++++++++------------- src/main.go | 46 ++++++++-- src/mmap_other.go | 9 ++ src/mmap_unix.go | 30 +++++++ src/server.go | 5 +- 10 files changed, 300 insertions(+), 85 deletions(-) create mode 100644 src/mmap_other.go create mode 100644 src/mmap_unix.go diff --git a/.env.example b/.env.example index 9798e07..1b2ea06 100644 --- a/.env.example +++ b/.env.example @@ -1,12 +1,7 @@ HOST=0.0.0.0 PORT=22 -# Generation settings -# Stored frame resolution in pixels. Higher = sharper but bigger .tsf files. -# regenerate your frames after changing it. -FRAME_RESOLUTION=512 - -# Playback settings. +# Playback settings # Playthroughs before the session is closed (0, unlimited). MAX_LOOP=5 # Whether to keep looping the same frame set or pick a random one after each @@ -30,7 +25,13 @@ INVERT=false FORCE_GRAYSCALE=false # Max rendered width/height in characters. -MAX_DIMENSION=1080 +MAX_DIMENSION=512 + +# Memory budget in MB for the rendered-frame cache (0 disables caching). +RENDER_CACHE_MB=256 + +# Go soft memory limit; set below your container limit to avoid OOM kills. +GOMEMLIMIT=1GiB # Connection limits. New connections over a limit are dropped immediately. # Max simultaneous connections from a single client IP. diff --git a/README.md b/README.md index 0e3ce84..59f2d98 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Generate a frame set from a video through container image ```sh docker run --rm -v ./video.mp4:/home/app/video.mp4 -v ./frames:/home/app/frames \ - ghcr.io/yuzuzensai/trollssh:v1.0.1 trollssh --generate --video video.mp4 + ghcr.io/yuzuzensai/trollssh:v1.0.1 trollssh --generate --video video.mp4 --resolution 512 ``` This writes `frames/.tsf`, a simple container of color JPEG frames plus @@ -51,13 +51,20 @@ ssh anyone@localhost ## Configuration -Configuration is via environment variables, loaded from a `.env` file if one -exists (see [`.env.example`](.env.example) for the full annotated list). -Durations are in milliseconds. +Server configuration is via environment variables, loaded from a `.env` file +if one exists (see [`.env.example`](.env.example) for the full annotated +list). Durations are in milliseconds. Host keys (`data/id_rsa`, `data/id_ed25519`) are generated on first run and reused afterwards. +Frame generation is configured with flags: + +| Flag | Default | Description | +| -------------------- | ------- | -------------------------------------------------- | +| `--generate`, `-g` | | Generate a `.tsf` frame set instead of serving | +| `--video`, `-v` | | Source video path | +| `--resolution`, `-r` | `512` | Stored frame max dimension in pixels. Higher = sharper but bigger `.tsf` files and slower rendering | ## Customization @@ -75,7 +82,7 @@ Requirements: Go 1.25+ and `ffmpeg` / `ffprobe` on `PATH` (only for `--generate`). ```sh -go run ./src --generate --video video.mp4 +go run ./src --generate --video video.mp4 --resolution 512 go run ./src ``` diff --git a/src/app_test.go b/src/app_test.go index cc5a015..db4d5f5 100644 --- a/src/app_test.go +++ b/src/app_test.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "os" "path/filepath" "strings" @@ -168,7 +169,7 @@ func TestFrameToAscii(t *testing.T) { // Below threshold -> first ramp char; full brightness -> last. opts := asciiOptions{brightnessThreshold: 40, charset: "standard"} ramp := []rune(resolveCharset("standard")) - out := []rune(frameToAscii([]byte{0, 255}, opts)) + out := []rune(frameToAscii([]byte{0, 255}, ramp, opts)) if out[0] != ramp[0] { t.Errorf("dark px = %q, want %q", out[0], ramp[0]) } @@ -180,12 +181,51 @@ func TestFrameToAscii(t *testing.T) { func TestFrameToAsciiInvert(t *testing.T) { opts := asciiOptions{brightnessThreshold: 40, charset: "standard", invert: true} ramp := []rune(resolveCharset("standard")) - out := []rune(frameToAscii([]byte{255}, opts)) + out := []rune(frameToAscii([]byte{255}, ramp, opts)) if out[0] != ramp[0] { t.Errorf("inverted bright = %q, want %q", out[0], ramp[0]) } } +func TestRenderCacheEvictsByBytes(t *testing.T) { + budget := 3 * entryCost("k", strings.Repeat("x", 1000)) + c := newRenderCache(budget) + for i := range 5 { + c.put(fmt.Sprintf("%d", i), strings.Repeat("x", 1000)) + } + if c.size > budget { + t.Errorf("size %d exceeds budget %d", c.size, budget) + } + if _, ok := c.get("0"); ok { + t.Error("oldest entry should have been evicted") + } + if _, ok := c.get("4"); !ok { + t.Error("newest entry should be cached") + } +} + +func TestRenderCacheDisabled(t *testing.T) { + c := newRenderCache(0) + if c != nil { + t.Fatal("zero budget should disable the cache") + } + c.put("k", "v") + if _, ok := c.get("k"); ok { + t.Error("nil cache should never hit") + } +} + +func TestRenderCacheRejectsOversizedEntry(t *testing.T) { + c := newRenderCache(256) + c.put("big", strings.Repeat("x", 10_000)) + if _, ok := c.get("big"); ok { + t.Error("entry larger than budget should not be cached") + } + if c.size != 0 { + t.Errorf("size = %d, want 0", c.size) + } +} + func TestConnectionTracker(t *testing.T) { tr := newConnectionTracker() tr.increment("1.2.3.4") diff --git a/src/config.go b/src/config.go index 0706099..072c6b4 100644 --- a/src/config.go +++ b/src/config.go @@ -28,7 +28,7 @@ type Config struct { MaxAuthAttempts int HandshakeTimeout time.Duration MaxDimension int - FrameResolution int + RenderCacheMB int BrightnessThreshold int Charset string Invert bool @@ -126,7 +126,7 @@ func loadConfig() Config { MaxAuthAttempts: envInt("MAX_AUTH_ATTEMPTS", 6, 1, maxInt), HandshakeTimeout: envDurationMs("HANDSHAKE_TIMEOUT", 10*time.Second), MaxDimension: envInt("MAX_DIMENSION", 512, 1, 4096), - FrameResolution: envInt("FRAME_RESOLUTION", 360, 16, 1080), + RenderCacheMB: envInt("RENDER_CACHE_MB", 256, 0, maxInt), BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100), Charset: envString("CHARSET", "detailed"), Invert: envBool("INVERT", false), diff --git a/src/frameformat.go b/src/frameformat.go index 7dbf319..09ee8a3 100644 --- a/src/frameformat.go +++ b/src/frameformat.go @@ -48,7 +48,7 @@ func writeTSF(output string, data *FramesContainer) error { } func loadTSF(filename string) (*FramesContainer, error) { - raw, err := os.ReadFile(filename) + raw, err := readFrameFile(filename) if err != nil { return nil, err } diff --git a/src/frames.go b/src/frames.go index 1637aa4..66b1ffc 100644 --- a/src/frames.go +++ b/src/frames.go @@ -7,8 +7,10 @@ import ( "image" "image/color" "image/jpeg" + "strconv" "strings" "sync" + "unicode/utf8" "golang.org/x/image/draw" ) @@ -62,18 +64,49 @@ func resolveCharset(charset string) string { return charset } -func resizeFrame(frame []byte, width, height int, keepAspectRatio bool, tier colorTier) (draw.Image, error) { +var pixPool sync.Pool + +func getPixBuf(n int) []byte { + if v := pixPool.Get(); v != nil { + if b := *v.(*[]byte); cap(b) >= n { + return b[:n] + } + } + return make([]byte, n) +} + +func putPixBuf(b []byte) { + pixPool.Put(&b) +} + +var outPool sync.Pool + +func getOutBuf(capacity int) []byte { + if v := outPool.Get(); v != nil { + if b := *v.(*[]byte); cap(b) >= capacity { + return b[:0] + } + } + return make([]byte, 0, capacity) +} + +func putOutBuf(b []byte) { + outPool.Put(&b) +} + +func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool, tier colorTier) (draw.Image, error) { src, err := jpeg.Decode(bytes.NewReader(frame)) if err != nil { return nil, err } + rect := image.Rect(0, 0, width, height) var dst draw.Image var bg color.Color if tier == colorTierNone { - dst = image.NewGray(image.Rect(0, 0, width, height)) + dst = &image.Gray{Pix: pix[:width*height], Stride: width, Rect: rect} bg = color.Gray{0} } else { - dst = image.NewNRGBA(image.Rect(0, 0, width, height)) + dst = &image.NRGBA{Pix: pix[:4*width*height], Stride: 4 * width, Rect: rect} bg = color.Black } if keepAspectRatio { @@ -114,16 +147,17 @@ func rampIndex(brightness, threshold, total int, invert bool) int { return index } -func frameToAscii(pixels []byte, options asciiOptions) string { - ramp := []rune(resolveCharset(options.charset)) +func frameToAscii(pixels []byte, ramp []rune, options asciiOptions) string { total := len(ramp) - var b strings.Builder + buf := getOutBuf(len(pixels) * 4) for _, p := range pixels { brightness := int(p) * 100 / 255 index := rampIndex(brightness, options.brightnessThreshold, total, options.invert) - b.WriteRune(ramp[index]) + buf = utf8.AppendRune(buf, ramp[index]) } - return b.String() + ascii := string(buf) + putOutBuf(buf) + return ascii } const ansiReset = "\x1b[0m" @@ -147,11 +181,25 @@ func quantize256(r, g, b uint8) int { return 16 + 36*toLevel(r) + 6*toLevel(g) + toLevel(b) } -func frameToAnsi(img *image.NRGBA, options asciiOptions, tier colorTier) string { - ramp := []rune(resolveCharset(options.charset)) +func appendColor(buf []byte, r, g, b uint8, tier colorTier) []byte { + if tier == colorTierTrueColor { + buf = append(buf, "\x1b[38;2;"...) + buf = strconv.AppendUint(buf, uint64(r), 10) + buf = append(buf, ';') + buf = strconv.AppendUint(buf, uint64(g), 10) + buf = append(buf, ';') + buf = strconv.AppendUint(buf, uint64(b), 10) + } else { + buf = append(buf, "\x1b[38;5;"...) + buf = strconv.AppendUint(buf, uint64(quantize256(r, g, b)), 10) + } + return append(buf, 'm') +} + +func frameToAnsi(img *image.NRGBA, ramp []rune, options asciiOptions, tier colorTier) string { total := len(ramp) bounds := img.Bounds() - var b strings.Builder + buf := getOutBuf(bounds.Dx() * bounds.Dy() * 16) var lastR, lastG, lastB uint8 first := true for y := bounds.Min.Y; y < bounds.Max.Y; y++ { @@ -161,33 +209,29 @@ func frameToAnsi(img *image.NRGBA, options asciiOptions, tier colorTier) string brightness := (int(r)*299 + int(g)*587 + int(bl)*114) / 255 / 10 index := rampIndex(brightness, options.brightnessThreshold, total, options.invert) if first || r != lastR || g != lastG || bl != lastB { - if tier == colorTierTrueColor { - fmt.Fprintf(&b, "\x1b[38;2;%d;%d;%dm", r, g, bl) - } else { - fmt.Fprintf(&b, "\x1b[38;5;%dm", quantize256(r, g, bl)) - } + buf = appendColor(buf, r, g, bl, tier) lastR, lastG, lastB = r, g, bl first = false } - b.WriteRune(ramp[index]) + buf = utf8.AppendRune(buf, ramp[index]) } if y < bounds.Max.Y-1 { - b.WriteString(ansiReset + "\r\n") + buf = append(buf, ansiReset+"\r\n"...) first = true } } - b.WriteString(ansiReset) - return b.String() + buf = append(buf, ansiReset...) + ascii := string(buf) + putOutBuf(buf) + return ascii } -type FrameRenderer struct { - colorFrames [][]byte - options asciiOptions - maxEntries int - - mu sync.Mutex - cache map[string]*list.Element - order *list.List +type renderCache struct { + mu sync.Mutex + maxBytes int64 + size int64 + entries map[string]*list.Element + order *list.List } type cacheEntry struct { @@ -195,52 +239,101 @@ type cacheEntry struct { ascii string } -func newFrameRenderer(colorFrames [][]byte, options asciiOptions) *FrameRenderer { +func entryCost(key, ascii string) int64 { + return int64(len(key)+len(ascii)) + 128 +} + +func newRenderCache(maxBytes int64) *renderCache { + if maxBytes <= 0 { + return nil + } + return &renderCache{ + maxBytes: maxBytes, + entries: make(map[string]*list.Element), + order: list.New(), + } +} + +func (c *renderCache) get(key string) (string, bool) { + if c == nil { + return "", false + } + c.mu.Lock() + defer c.mu.Unlock() + el, ok := c.entries[key] + if !ok { + return "", false + } + c.order.MoveToBack(el) + return el.Value.(*cacheEntry).ascii, true +} + +func (c *renderCache) put(key, ascii string) { + if c == nil { + return + } + cost := entryCost(key, ascii) + if cost > c.maxBytes { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if _, ok := c.entries[key]; ok { + return + } + c.entries[key] = c.order.PushBack(&cacheEntry{key, ascii}) + c.size += cost + for c.size > c.maxBytes { + oldest := c.order.Front() + c.order.Remove(oldest) + evicted := oldest.Value.(*cacheEntry) + delete(c.entries, evicted.key) + c.size -= entryCost(evicted.key, evicted.ascii) + } +} + +type FrameRenderer struct { + setID int + colorFrames [][]byte + options asciiOptions + ramp []rune + cache *renderCache +} + +func newFrameRenderer(setID int, colorFrames [][]byte, options asciiOptions, cache *renderCache) *FrameRenderer { return &FrameRenderer{ + setID: setID, colorFrames: colorFrames, options: options, - maxEntries: 4096, - cache: make(map[string]*list.Element), - order: list.New(), + ramp: []rune(resolveCharset(options.charset)), + cache: cache, } } func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, tier colorTier) (string, error) { - key := fmt.Sprintf("%d:%dx%d:%t:%d", index, width, height, keepAspectRatio, tier) - - r.mu.Lock() - if el, ok := r.cache[key]; ok { - r.order.MoveToBack(el) - ascii := el.Value.(*cacheEntry).ascii - r.mu.Unlock() + key := fmt.Sprintf("%d:%d:%dx%d:%t:%d", r.setID, index, width, height, keepAspectRatio, tier) + if ascii, ok := r.cache.get(key); ok { return ascii, nil } - r.mu.Unlock() + n := width * height + if tier != colorTierNone { + n *= 4 + } + pix := getPixBuf(n) + img, err := resizeFrame(r.colorFrames[index], pix, width, height, keepAspectRatio, tier) + if err != nil { + putPixBuf(pix) + return "", err + } var ascii string if tier == colorTierNone { - img, err := resizeFrame(r.colorFrames[index], width, height, keepAspectRatio, tier) - if err != nil { - return "", err - } - ascii = frameToAscii(img.(*image.Gray).Pix, r.options) + ascii = frameToAscii(img.(*image.Gray).Pix, r.ramp, r.options) } else { - img, err := resizeFrame(r.colorFrames[index], width, height, keepAspectRatio, tier) - if err != nil { - return "", err - } - ascii = frameToAnsi(img.(*image.NRGBA), r.options, tier) + ascii = frameToAnsi(img.(*image.NRGBA), r.ramp, r.options, tier) } + putPixBuf(pix) - r.mu.Lock() - if _, ok := r.cache[key]; !ok { - r.cache[key] = r.order.PushBack(&cacheEntry{key, ascii}) - if r.order.Len() > r.maxEntries { - oldest := r.order.Front() - r.order.Remove(oldest) - delete(r.cache, oldest.Value.(*cacheEntry).key) - } - } - r.mu.Unlock() + r.cache.put(key, ascii) return ascii, nil } diff --git a/src/main.go b/src/main.go index 89c5177..5c2c321 100644 --- a/src/main.go +++ b/src/main.go @@ -6,7 +6,9 @@ import ( "os/signal" "path/filepath" "runtime" + "runtime/debug" "sort" + "strconv" "strings" "sync" "syscall" @@ -16,12 +18,13 @@ import ( ) type cliArgs struct { - generate bool - video string + generate bool + video string + resolution int } func parseArgs(argv []string) cliArgs { - var args cliArgs + args := cliArgs{resolution: 512} for i := 0; i < len(argv); i++ { switch argv[i] { case "--generate", "-g": @@ -31,6 +34,13 @@ func parseArgs(argv []string) cliArgs { i++ args.video = argv[i] } + case "--resolution", "-r": + if i+1 < len(argv) { + i++ + if n, err := strconv.Atoi(argv[i]); err == nil { + args.resolution = max(n, 16) + } + } } } return args @@ -52,7 +62,7 @@ func resolveVideoPath(explicitPath string) string { return "" } -func generateFrames(config Config, framesDir, videoArg string) { +func generateFrames(framesDir, videoArg string, resolution int) { if videoArg == "" { fail("No source video given. Pass --video .") } @@ -68,7 +78,7 @@ func generateFrames(config Config, framesDir, videoArg string) { output := filepath.Join(framesDir, base+".tsf") logInfo(fmt.Sprintf("Generating frames from %q -> %s", videoPath, output)) - if err := processVideo(videoPath, output, config.FrameResolution); err != nil { + if err := processVideo(videoPath, output, resolution); err != nil { fail(fmt.Sprintf("Failed to generate frames from %q: %s", videoPath, err.Error())) } } @@ -143,9 +153,33 @@ func loadAllFrames(framesDir string) []*FramesContainer { return results } +func applyMemoryLimit() { + if os.Getenv("GOMEMLIMIT") != "" { + return + } + for _, path := range []string{ + "/sys/fs/cgroup/memory.max", + "/sys/fs/cgroup/memory/memory.limit_in_bytes", + } { + raw, err := os.ReadFile(path) + if err != nil { + continue + } + n, err := strconv.ParseInt(strings.TrimSpace(string(raw)), 10, 64) + if err != nil || n <= 0 || n > 1<<48 { + return + } + limit := n * 9 / 10 + debug.SetMemoryLimit(limit) + logInfo(fmt.Sprintf("Memory limit set to %d MB (90%% of cgroup limit)", limit>>20)) + return + } +} + func main() { _ = godotenv.Load() logThreshold = resolveThreshold() + applyMemoryLimit() config := loadConfig() args := parseArgs(os.Args[1:]) @@ -158,7 +192,7 @@ func main() { framesDir := filepath.Join(cwd, "frames") if args.generate { - generateFrames(config, framesDir, args.video) + generateFrames(framesDir, args.video, args.resolution) return } diff --git a/src/mmap_other.go b/src/mmap_other.go new file mode 100644 index 0000000..c792989 --- /dev/null +++ b/src/mmap_other.go @@ -0,0 +1,9 @@ +//go:build !unix + +package main + +import "os" + +func readFrameFile(filename string) ([]byte, error) { + return os.ReadFile(filename) +} diff --git a/src/mmap_unix.go b/src/mmap_unix.go new file mode 100644 index 0000000..ede9ec9 --- /dev/null +++ b/src/mmap_unix.go @@ -0,0 +1,30 @@ +//go:build unix + +package main + +import ( + "os" + "syscall" +) + +func readFrameFile(filename string) ([]byte, error) { + f, err := os.Open(filename) + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + + info, err := f.Stat() + if err != nil { + return nil, err + } + size := info.Size() + if size <= 0 || size != int64(int(size)) { + return os.ReadFile(filename) + } + data, err := syscall.Mmap(int(f.Fd()), 0, int(size), syscall.PROT_READ, syscall.MAP_SHARED) + if err != nil { + return os.ReadFile(filename) + } + return data, nil +} diff --git a/src/server.go b/src/server.go index 7bd74d2..d4d1413 100644 --- a/src/server.go +++ b/src/server.go @@ -105,15 +105,16 @@ func clampDimension(value, max int) int { func createServer(deps ServerDeps) *Server { config := deps.Config + cache := newRenderCache(int64(config.RenderCacheMB) << 20) sets := make([]frameSet, len(deps.VideoSets)) for i, data := range deps.VideoSets { sets[i] = frameSet{ data: data, - renderer: newFrameRenderer(data.ColorFrames, asciiOptions{ + renderer: newFrameRenderer(i, data.ColorFrames, asciiOptions{ brightnessThreshold: config.BrightnessThreshold, charset: config.Charset, invert: config.Invert, - }), + }, cache), } }