perf: bound indexing memory and flush partial progress on interrupt

cmd/gocate/main.go

@@ -4,12 +4,16 @@
 package main
 
 import (
+	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"runtime/pprof"
 	"strings"
+	"syscall"
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -44,6 +48,23 @@ func run() error {
 
 	flag.Parse()
 
+	// A SIGINT cancels the in-flight indexing run so the consumer can flush its
+	// current write batch (committing partial progress) instead of being killed
+	// mid-transaction. A second SIGINT forces an immediate exit.
+	ctx, stop := context.WithCancel(context.Background())
+	defer stop()
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(sigCh)
+	go func() {
+		first := <-sigCh
+		log.Warn().Stringer("signal", first).Msg("interrupt received; flushing index and exiting")
+		stop()
+		if <-sigCh != nil {
+			os.Exit(130) // second interrupt: bail out now
+		}
+	}()
+
 	if *profile {
 		stop, err := startProfile("default.pgo")
 		if err != nil {
@@ -67,8 +88,12 @@ func run() error {
 		if err != nil {
 			return fmt.Errorf("resolve path %q: %w", *updatePath, err)
 		}
-		if err := index.Run(s, root, index.Options{Hash: !*noHash, Quick: *quick}); err != nil {
-			return err
+		if err := index.RunCtx(ctx, s, root, index.Options{Hash: !*noHash, Quick: *quick}); err != nil {
+			if errors.Is(err, context.Canceled) {
+				log.Warn().Msg("indexing interrupted; partial progress committed")
+			} else {
+				return err
+			}
 		}
 	}
 

internal/index/hash.go

@@ -2,22 +2,60 @@ package index
 
 import (
 	"fmt"
+	"io"
 	"os"
 
 	"github.com/kalafut/imohash"
 	"github.com/zeebo/xxh3"
 )
 
-// hashFile reads path once and returns its imohash and xxh3 hashes. Zero-byte
-// files are skipped (both hashes empty, nil error) since there is nothing to
-// distinguish them by content.
+// hashFile returns the imohash and xxh3 hashes of path without reading the
+// whole file into memory. imohash only samples a fixed number of bytes from
+// the file (see imohash.SampleSize/SampleThreshold) and xxh3 is computed by
+// streaming the file in chunks, so peak memory stays bounded regardless of
+// file size. Zero-byte files are skipped (both hashes empty, nil error) since
+// there is nothing to distinguish them by content.
+//
+// The hashes are byte-identical to the previous os.ReadFile-based
+// implementation, so existing database rows remain valid.
 func hashFile(path string) (imo, xxh string, err error) {
-	data, err := os.ReadFile(path)
+	f, err := os.Open(path)
 	if err != nil {
-		return "", "", fmt.Errorf("read %q: %w", path, err)
+		return "", "", fmt.Errorf("open %q: %w", path, err)
 	}
-	if len(data) == 0 {
+	defer func() {
+		if cerr := f.Close(); cerr != nil && err == nil {
+			err = fmt.Errorf("close %q: %w", path, cerr)
+		}
+	}()
+
+	fi, err := f.Stat()
+	if err != nil {
+		return "", "", fmt.Errorf("stat %q: %w", path, err)
+	}
+	if fi.Size() == 0 {
 		return "", "", nil
 	}
-	return fmt.Sprintf("%x", imohash.Sum(data)), fmt.Sprintf("%x", xxh3.Hash(data)), nil
+
+	sr := io.NewSectionReader(f, 0, fi.Size())
+
+	// imohash reads only its fixed samples from the section reader. Note that
+	// this leaves sr positioned somewhere in the file (the end for small files,
+	// the tail sample for large ones), so we must rewind before the next pass.
+	imoSum, err := imohash.SumSectionReader(sr)
+	if err != nil {
+		return "", "", fmt.Errorf("imohash %q: %w", path, err)
+	}
+
+	// Rewind to the start and stream xxh3 over the full file so the contents
+	// are never held in memory at once.
+	if _, err := sr.Seek(0, io.SeekStart); err != nil {
+		return "", "", fmt.Errorf("seek %q: %w", path, err)
+	}
+	h := xxh3.New()
+	if _, err := io.Copy(h, sr); err != nil {
+		return "", "", fmt.Errorf("xxh3 %q: %w", path, err)
+	}
+
+	return fmt.Sprintf("%x", imoSum), fmt.Sprintf("%x", h.Sum64()), nil
 }

internal/index/hash_test.go

@@ -1,9 +1,13 @@
 package index
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
+
+	"github.com/kalafut/imohash"
+	"github.com/zeebo/xxh3"
 )
 
 func writeFile(t *testing.T, dir, name, content string) string {
@@ -55,3 +59,51 @@ func TestHashFileMissing(t *testing.T) {
 		t.Fatal("hashFile on missing file should error")
 	}
 }
+
+// TestHashFileMatchesReadFileBasis verifies the streaming implementation
+// produces byte-identical hashes to the previous os.ReadFile-based approach
+// across small, medium (below imohash sample threshold), and large (above it)
+// files, so existing database rows remain valid.
+func TestHashFileMatchesReadFileBasis(t *testing.T) {
+	cases := []struct {
+		name string
+		data []byte
+	}{
+		{"small", []byte("hello world")},
+		{"medium", randomBytes(64 * 1024)}, // below imohash.SampleThreshold (128KiB)
+		{"large", randomBytes(512 * 1024)}, // above imohash.SampleThreshold
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			dir := t.TempDir()
+			p := writeFile(t, dir, "f", string(tc.data))
+
+			gotImo, gotXxh, err := hashFile(p)
+			if err != nil {
+				t.Fatalf("hashFile: %v", err)
+			}
+
+			data, err := os.ReadFile(p)
+			if err != nil {
+				t.Fatalf("read: %v", err)
+			}
+			wantImo := fmt.Sprintf("%x", imohash.Sum(data))
+			wantXxh := fmt.Sprintf("%x", xxh3.Hash(data))
+
+			if gotImo != wantImo {
+				t.Fatalf("imohash mismatch: got %q want %q", gotImo, wantImo)
+			}
+			if gotXxh != wantXxh {
+				t.Fatalf("xxh3 mismatch: got %q want %q", gotXxh, wantXxh)
+			}
+		})
+	}
+}
+
+func randomBytes(n int) []byte {
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = byte(i * 31)
+	}
+	return b
+}