From 2d4d36de552f26d8207f171c65dba6e61cf09bec Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Mon, 27 Apr 2026 12:38:48 -0700
Subject: [PATCH] perf(desktop): prune cached snapshots

---
 internal/notiondesktop/desktop.go      | 58 ++++++++++++++++++++++++++
 internal/notiondesktop/desktop_test.go | 57 +++++++++++++++++++++++++
 2 files changed, 115 insertions(+)
 create mode 100644 internal/notiondesktop/desktop_test.go

diff --git a/internal/notiondesktop/desktop.go b/internal/notiondesktop/desktop.go
index ea2ad19..ea3d384 100644
--- a/internal/notiondesktop/desktop.go
+++ b/internal/notiondesktop/desktop.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"time"
 
@@ -16,6 +17,7 @@ import (
 )
 
 const SourceName = "desktop"
+const desktopSnapshotRetention = 2
 
 type Source struct {
 	Path      string
@@ -122,9 +124,65 @@ func snapshotDB(path, cacheDir string) (string, error) {
 			return "", err
 		}
 	}
+	if err := pruneDesktopSnapshots(cacheDir, desktopSnapshotRetention, outPath); err != nil {
+		return "", err
+	}
 	return outPath, nil
 }
 
+type desktopSnapshot struct {
+	path    string
+	modTime time.Time
+}
+
+func pruneDesktopSnapshots(cacheDir string, keep int, current string) error {
+	if keep < 1 {
+		keep = 1
+	}
+	entries, err := os.ReadDir(cacheDir)
+	if err != nil {
+		return err
+	}
+	var snapshots []desktopSnapshot
+	for _, entry := range entries {
+		name := entry.Name()
+		if entry.IsDir() || !strings.HasPrefix(name, "notion-desktop-") || !strings.HasSuffix(name, ".db") {
+			continue
+		}
+		path := filepath.Join(cacheDir, name)
+		info, err := entry.Info()
+		if err != nil {
+			return err
+		}
+		snapshots = append(snapshots, desktopSnapshot{path: path, modTime: info.ModTime()})
+	}
+	sort.SliceStable(snapshots, func(i, j int) bool {
+		if snapshots[i].modTime.Equal(snapshots[j].modTime) {
+			return snapshots[i].path > snapshots[j].path
+		}
+		return snapshots[i].modTime.After(snapshots[j].modTime)
+	})
+	keepPaths := map[string]bool{}
+	if current != "" {
+		keepPaths[filepath.Clean(current)] = true
+	}
+	for i := 0; i < len(snapshots) && len(keepPaths) < keep; i++ {
+		keepPaths[filepath.Clean(snapshots[i].path)] = true
+	}
+	for _, snapshot := range snapshots {
+		path := filepath.Clean(snapshot.path)
+		if keepPaths[path] {
+			continue
+		}
+		for _, target := range []string{path, path + "-wal", path + "-shm"} {
+			if err := os.Remove(target); err != nil && !os.IsNotExist(err) {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 func copyFile(src, dst string, perm os.FileMode) error {
 	in, err := os.Open(src)
 	if err != nil {
diff --git a/internal/notiondesktop/desktop_test.go b/internal/notiondesktop/desktop_test.go
new file mode 100644
index 0000000..18b9a0d
--- /dev/null
+++ b/internal/notiondesktop/desktop_test.go
@@ -0,0 +1,57 @@
+package notiondesktop
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+func TestPruneDesktopSnapshotsKeepsNewestAndSidecars(t *testing.T) {
+	dir := t.TempDir()
+	names := []string{
+		"notion-desktop-1000.db",
+		"notion-desktop-2000.db",
+		"notion-desktop-3000.db",
+	}
+	for i, name := range names {
+		path := filepath.Join(dir, name)
+		if err := os.WriteFile(path, []byte(name), 0o600); err != nil {
+			t.Fatal(err)
+		}
+		for _, suffix := range []string{"-wal", "-shm"} {
+			if err := os.WriteFile(path+suffix, []byte(suffix), 0o600); err != nil {
+				t.Fatal(err)
+			}
+		}
+		modTime := time.Unix(int64(i+1), 0)
+		for _, target := range []string{path, path + "-wal", path + "-shm"} {
+			if err := os.Chtimes(target, modTime, modTime); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+
+	current := filepath.Join(dir, "notion-desktop-3000.db")
+	if err := pruneDesktopSnapshots(dir, 2, current); err != nil {
+		t.Fatal(err)
+	}
+
+	for _, name := range []string{"notion-desktop-2000.db", "notion-desktop-3000.db"} {
+		path := filepath.Join(dir, name)
+		for _, target := range []string{path, path + "-wal", path + "-shm"} {
+			if _, err := os.Stat(target); err != nil {
+				t.Fatalf("expected %s to remain: %v", target, err)
+			}
+		}
+	}
+	for _, target := range []string{
+		filepath.Join(dir, "notion-desktop-1000.db"),
+		filepath.Join(dir, "notion-desktop-1000.db-wal"),
+		filepath.Join(dir, "notion-desktop-1000.db-shm"),
+	} {
+		if _, err := os.Stat(target); !os.IsNotExist(err) {
+			t.Fatalf("expected %s to be pruned, got %v", target, err)
+		}
+	}
+}