diff --git a/cmd/index.go b/cmd/index.go index 334b64fa..ef0e9618 100644 --- a/cmd/index.go +++ b/cmd/index.go @@ -251,11 +251,28 @@ func runIndexer(cmd *cobra.Command, cfg *config.ConfigService, emb *embedder.Fai } defer lock.Release() - // Cancel context on SIGTERM or SIGINT so the indexer stops cleanly and - // the deferred lock.Release() runs before exit. + // Install signal handling before donor seeding: copying a large sibling + // index can take seconds, and cancellation must be able to close/remove the + // seed temp file before the process exits. ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT) defer stop() + // Reuse a sibling worktree's index for a brand-new database instead of + // re-embedding from scratch. The index lock serializes CLI indexers, while + // SeedFromDonor's seed lock also serializes this copy with MCP callers. + // A forced rebuild cannot reuse the copied embeddings, so skip the donor + // copy in that mode. + force, _ := cmd.Flags().GetBool("force") + if !force { + seedFromDonorIfNew(ctx, dbPath, projectPath, emb.ModelName(), logger, seedOptions{ + status: p.Info, + }) + } + if ctx.Err() != nil { + logger.Info("indexing cancelled by signal", "project", projectPath) + return + } + idx, setupErr := setupIndexer(cfg, emb, dbPath, logger) if setupErr != nil { err = setupErr @@ -264,7 +281,7 @@ func runIndexer(cmd *cobra.Command, cfg *config.ConfigService, emb *embedder.Fai defer func() { _ = idx.Close() }() start := time.Now() - stats, err = performIndexing(ctx, cmd, idx, projectPath, p) + stats, err = performIndexing(ctx, force, idx, projectPath, p) elapsed = time.Since(start).Round(time.Millisecond) if err != nil && ctx.Err() != nil { // A signal arrived; treat as clean exit. @@ -274,9 +291,7 @@ func runIndexer(cmd *cobra.Command, cfg *config.ConfigService, emb *embedder.Fai return } -func performIndexing(ctx context.Context, cmd *cobra.Command, idx *index.Indexer, projectPath string, p *tui.Progress) (index.Stats, error) { - force, _ := cmd.Flags().GetBool("force") - +func performIndexing(ctx context.Context, force bool, idx *index.Indexer, projectPath string, p *tui.Progress) (index.Stats, error) { progress := p.AsProgressFunc() if force { diff --git a/cmd/seed.go b/cmd/seed.go new file mode 100644 index 00000000..016b2be0 --- /dev/null +++ b/cmd/seed.go @@ -0,0 +1,92 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "context" + "fmt" + "log/slog" + "os" + + "github.com/ory/lumen/internal/config" + "github.com/ory/lumen/internal/index" +) + +type seedOptions struct { + findDonor func(string, string) string + seed func(context.Context, string, string, string) (bool, error) + status func(string) +} + +// seedFromDonorIfNew seeds dbPath from a sibling worktree's index when dbPath +// does not yet exist, so a fresh git worktree reuses an already-indexed +// worktree's embeddings instead of re-embedding every file from scratch. +// +// Callers may additionally hold the index lock for dbPath, but SeedFromDonor +// has its own advisory lock so CLI and MCP callers cannot duplicate the copy. +// Seeding is best-effort — any failure is logged and indexing continues with a +// from-scratch build. The returned warning is suitable for surfacing to MCP +// clients. When dbPath already exists it is a single stat on the hot path. +func seedFromDonorIfNew(ctx context.Context, dbPath, projectPath, model string, logger *slog.Logger, opts seedOptions) string { + if _, err := os.Stat(dbPath); !os.IsNotExist(err) { + // Exists already, or stat failed for some other reason — nothing to do. + return "" + } + + findDonor := opts.findDonor + if findDonor == nil { + findDonor = config.FindDonorIndex + } + donorPath := findDonor(projectPath, model) + if donorPath == "" { + return "" + } + + logger.Info("seeding index from donor worktree", + "project_path", projectPath, + "donor_path", donorPath, + ) + if opts.status != nil { + opts.status("Seeding index from sibling worktree...") + } + + seed := opts.seed + if seed == nil { + seed = index.SeedFromDonorContext + } + seeded, err := seed(ctx, donorPath, dbPath, projectPath) + if err != nil { + logger.Warn("seed from donor worktree failed", + "project_path", projectPath, + "donor_path", donorPath, + "error", err, + ) + warning := fmt.Sprintf("index seeded from scratch (sibling copy failed: %v)", err) + if opts.status != nil { + opts.status(fmt.Sprintf("Sibling index copy failed: %v; indexing from scratch.", err)) + } + return warning + } + if seeded && opts.status != nil { + opts.status("Seeded index from sibling worktree.") + } else if opts.status != nil { + if _, statErr := os.Stat(dbPath); statErr == nil { + opts.status("Index was seeded by another process.") + } else { + opts.status("Sibling index could not be reused; indexing from scratch.") + } + } + return "" +} diff --git a/cmd/seed_test.go b/cmd/seed_test.go new file mode 100644 index 00000000..949e7c64 --- /dev/null +++ b/cmd/seed_test.go @@ -0,0 +1,126 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "context" + "errors" + "io" + "log/slog" + "os" + "path/filepath" + "testing" +) + +func seedTestLogger() *slog.Logger { + return slog.New(slog.NewTextHandler(io.Discard, nil)) +} + +func TestSeedFromDonorIfNew(t *testing.T) { + tests := []struct { + name string + setupDB bool // pre-create the destination DB + donor string + seedErr error // error returned by the seed stub + wantFind bool // donor discovery should run + wantSeed bool // seed should run + }{ + { + name: "skips when DB already exists", + setupDB: true, + donor: "/donor.db", + wantFind: false, + wantSeed: false, + }, + { + name: "seeds when DB missing and donor found", + donor: "/donor.db", + wantFind: true, + wantSeed: true, + }, + { + name: "no seed when no donor found", + donor: "", + wantFind: true, + wantSeed: false, + }, + { + name: "seed error is swallowed", + donor: "/donor.db", + seedErr: errors.New("copy failed"), + wantFind: true, + wantSeed: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "index.db") + if tt.setupDB { + if err := os.WriteFile(dbPath, []byte("existing"), 0o644); err != nil { + t.Fatal(err) + } + } + + var findCalled, seedCalled bool + var gotDonor, gotDst string + var statuses []string + warning := seedFromDonorIfNew( + context.Background(), + dbPath, + "/project", + "model", + seedTestLogger(), + seedOptions{ + findDonor: func(_, _ string) string { + findCalled = true + return tt.donor + }, + seed: func(_ context.Context, donor, dst, projectPath string) (bool, error) { + seedCalled = true + gotDonor, gotDst = donor, dst + if projectPath != "/project" { + t.Errorf("seed project path = %q, want /project", projectPath) + } + return tt.seedErr == nil, tt.seedErr + }, + status: func(message string) { + statuses = append(statuses, message) + }, + }, + ) + + if findCalled != tt.wantFind { + t.Errorf("donor discovery called = %v, want %v", findCalled, tt.wantFind) + } + if seedCalled != tt.wantSeed { + t.Errorf("seed called = %v, want %v", seedCalled, tt.wantSeed) + } + if tt.wantSeed && (gotDonor != tt.donor || gotDst != dbPath) { + t.Errorf("seed called with (%q, %q), want (%q, %q)", gotDonor, gotDst, tt.donor, dbPath) + } + if tt.seedErr != nil { + if warning == "" { + t.Error("expected warning when seed fails") + } + if len(statuses) != 2 { + t.Fatalf("expected start and failure status, got %v", statuses) + } + } else if warning != "" { + t.Errorf("unexpected warning: %q", warning) + } + }) + } +} diff --git a/cmd/stdio.go b/cmd/stdio.go index 551840d5..35114b95 100644 --- a/cmd/stdio.go +++ b/cmd/stdio.go @@ -200,7 +200,7 @@ type indexerCache struct { embedTimeout time.Duration // override for tests; 0 means defaultEmbedTimeout staleEmbedTimeout time.Duration // override for tests; 0 means defaultStaleEmbedTimeout findDonorFunc func(string, string) string // nil uses config.FindDonorIndex - seedFunc func(string, string) (bool, error) // nil uses index.SeedFromDonor + seedFunc func(context.Context, string, string, string) (bool, error) // nil uses index.SeedFromDonorContext ensureFreshFunc func(ctx context.Context, idx *index.Indexer, projectDir string, progress index.ProgressFunc) (bool, index.Stats, error) // nil uses idx.EnsureFresh log *slog.Logger wg sync.WaitGroup // tracks background reindex goroutines @@ -484,7 +484,6 @@ func (ic *indexerCache) getOrCreate(projectPath string, preferredRoot string, mo } // Seed from sibling worktree if this is a new index. - var seedWarning string isNewDB := false if _, statErr := os.Stat(dbPath); os.IsNotExist(statErr) { isNewDB = true @@ -494,29 +493,15 @@ func (ic *indexerCache) getOrCreate(projectPath string, preferredRoot string, mo "model", modelName, "index_version", config.IndexVersion, ) - findDonor := ic.findDonorFunc - if findDonor == nil { - findDonor = config.FindDonorIndex - } - if donorPath := findDonor(effectiveRoot, modelName); donorPath != "" { - ic.logger().Info("seeding index from donor worktree", - "effective_root", effectiveRoot, - "donor_path", donorPath, - ) - seedFn := ic.seedFunc - if seedFn == nil { - seedFn = index.SeedFromDonor - } - if _, seedErr := seedFn(donorPath, dbPath); seedErr != nil { - ic.logger().Warn("seed from donor worktree failed", - "effective_root", effectiveRoot, - "donor_path", donorPath, - "error", seedErr, - ) - seedWarning = fmt.Sprintf("index seeded from scratch (sibling copy failed: %v)", seedErr) - } - } } + seedCtx := ic.closeCtx + if seedCtx == nil { + seedCtx = context.Background() + } + seedWarning := seedFromDonorIfNew(seedCtx, dbPath, effectiveRoot, modelName, ic.logger(), seedOptions{ + findDonor: ic.findDonorFunc, + seed: ic.seedFunc, + }) idx, err := index.NewIndexer(dbPath, ic.embedder, ic.cfg.MaxChunkTokens()) if err != nil { diff --git a/cmd/stdio_test.go b/cmd/stdio_test.go index 2b313ff9..ea9104be 100644 --- a/cmd/stdio_test.go +++ b/cmd/stdio_test.go @@ -16,6 +16,7 @@ package cmd import ( "context" + "errors" "fmt" "io" "log/slog" @@ -1304,7 +1305,7 @@ func TestGetOrCreate_ReturnsSeedWarningWhenSeedFails(t *testing.T) { embedder: &stubEmbedder{}, cfg: newTestConfigService(t, 512), findDonorFunc: func(_, _ string) string { return "/fake/donor.db" }, - seedFunc: func(_, _ string) (bool, error) { + seedFunc: func(_ context.Context, _, _, _ string) (bool, error) { return false, fmt.Errorf("permission denied") }, } @@ -1322,6 +1323,63 @@ func TestGetOrCreate_ReturnsSeedWarningWhenSeedFails(t *testing.T) { } } +func TestGetOrCreate_SeedCancelledByClose(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("XDG_DATA_HOME", tmpDir) + + projectDir := filepath.Join(tmpDir, "project") + if err := os.MkdirAll(projectDir, 0o755); err != nil { + t.Fatal(err) + } + + closeCtx, closeFn := context.WithCancel(context.Background()) + seedStarted := make(chan struct{}) + seedResult := make(chan error, 1) + ic := &indexerCache{ + embedder: &stubEmbedder{}, + cfg: newTestConfigService(t, 512), + closeCtx: closeCtx, + closeFn: closeFn, + findDonorFunc: func(_, _ string) string { return "/fake/donor.db" }, + seedFunc: func(ctx context.Context, _, _, _ string) (bool, error) { + close(seedStarted) + <-ctx.Done() + seedResult <- ctx.Err() + return false, ctx.Err() + }, + } + + getDone := make(chan error, 1) + go func() { + _, _, _, err := ic.getOrCreate(projectDir, "") + getDone <- err + }() + + select { + case <-seedStarted: + case <-time.After(5 * time.Second): + t.Fatal("seed did not start") + } + + closeDone := make(chan struct{}) + go func() { + ic.Close() + close(closeDone) + }() + + select { + case <-closeDone: + case <-time.After(5 * time.Second): + t.Fatal("Close blocked while getOrCreate waited for the seed lock") + } + if err := <-seedResult; !errors.Is(err, context.Canceled) { + t.Fatalf("seed context error = %v, want context.Canceled", err) + } + if err := <-getDone; err != nil { + t.Fatalf("getOrCreate: %v", err) + } +} + func TestFormatSearchResults_IncludesSeedWarning(t *testing.T) { out := SemanticSearchOutput{ Results: nil, diff --git a/internal/config/seed.go b/internal/config/seed.go index cb29fb6c..404ea7ce 100644 --- a/internal/config/seed.go +++ b/internal/config/seed.go @@ -47,6 +47,18 @@ func FindDonorIndexBase(dataDir, projectPath, model string) string { // Find which worktree contains projectPath and compute the relative suffix. // This handles subdirectory effective roots (e.g., hydra-flake/backoffice) // by searching for the same subdirectory in sibling worktrees. + // + // Pick the DEEPEST (most specific) containing worktree, not the first match. + // `git worktree list` yields the main worktree first, so a first-match scan + // selects the main checkout for a worktree nested inside it — e.g. Claude + // Code's default repo/.claude/worktrees/ layout. That mis-identifies + // the current worktree, so relSuffix becomes ".claude/worktrees/", the + // donor search then looks for indexes at nonexistent + // /.claude/worktrees/ paths, and the one real donor (the + // parent repo's index) is skipped as "self" — leaving every nested worktree + // to re-embed from scratch. Every containing worktree is an ancestor of + // resolvedProject, so they form a prefix chain and the longest path is + // unambiguously the most specific. var myWorktree string for _, wt := range worktrees { resolved := wt @@ -57,8 +69,9 @@ func FindDonorIndexBase(dataDir, projectPath, model string) string { if err != nil || strings.HasPrefix(rel, "..") { continue } - myWorktree = resolved - break + if len(resolved) > len(myWorktree) { + myWorktree = resolved + } } if myWorktree == "" { return "" diff --git a/internal/config/seed_test.go b/internal/config/seed_test.go index a38e8c55..2fb66b5f 100644 --- a/internal/config/seed_test.go +++ b/internal/config/seed_test.go @@ -65,6 +65,44 @@ func TestFindDonorIndex_WithSibling(t *testing.T) { } } +func TestFindDonorIndex_NestedWorktree(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git not on PATH") + } + + // Reproduce Claude Code's layout: a worktree nested INSIDE the main repo at + //
/.claude/worktrees/. `git worktree list` reports the main + // checkout first, so donor discovery must not mistake the main checkout for + // the current (nested) worktree. + main := t.TempDir() + gitRun(t, main, "git", "init") + gitRun(t, main, "git", "commit", "--allow-empty", "-m", "init") + + nested := filepath.Join(main, ".claude", "worktrees", "feature") + gitRun(t, main, "git", "worktree", "add", nested) + + mainResolved, err := filepath.EvalSymlinks(main) + if err != nil { + t.Fatal(err) + } + + // The parent repo's index is the only real donor. + dataDir := t.TempDir() + donorDB := DBPathForProjectBase(dataDir, mainResolved, "test-model") + if err := os.MkdirAll(filepath.Dir(donorDB), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(donorDB, []byte("fake-db"), 0o644); err != nil { + t.Fatal(err) + } + + // From the nested worktree we must find the parent repo's index, not "". + result := FindDonorIndexBase(dataDir, nested, "test-model") + if result != donorDB { + t.Fatalf("expected nested worktree to seed from parent index %q, got %q", donorDB, result) + } +} + func TestFindDonorIndex_WrongModel(t *testing.T) { if _, err := exec.LookPath("git"); err != nil { t.Skip("git not on PATH") diff --git a/internal/index/seed.go b/internal/index/seed.go index bccf191e..936b45e5 100644 --- a/internal/index/seed.go +++ b/internal/index/seed.go @@ -15,22 +15,54 @@ package index import ( + "context" "database/sql" + "errors" "fmt" - "io" + "net/url" "os" "path/filepath" - _ "github.com/mattn/go-sqlite3" // register sqlite3 driver for WAL checkpoint + "github.com/ory/lumen/internal/indexlock" + + _ "github.com/mattn/go-sqlite3" // register sqlite3 driver ) -// SeedFromDonor copies the donor SQLite database to dstPath if dstPath does -// not already exist. It checkpoints the WAL first to ensure a self-contained -// copy, then performs an atomic copy (write to temp file + rename). +// SeedFromDonor snapshots the donor SQLite database to dstPath if dstPath does +// not already exist. It stamps projectPath as the new database owner, then +// atomically publishes the snapshot. +func SeedFromDonor(donorPath, dstPath, projectPath string) (bool, error) { + return SeedFromDonorContext(context.Background(), donorPath, dstPath, projectPath) +} + +// SeedFromDonorContext is SeedFromDonor with cancellation support. +// +// Seeding is safe to run concurrently from multiple processes (e.g. the +// SessionStart background indexer and the first MCP search racing to warm the +// same fresh worktree): an advisory seed lock is claimed before reading the +// donor, so only the winner copies while other callers wait and then observe +// the published destination. Advisory locks are released by the OS on process +// exit, avoiding stale claims after a crash. // // Returns (true, nil) if seeded successfully, (false, nil) if dstPath already // exists, or (false, error) on failure. -func SeedFromDonor(donorPath, dstPath string) (bool, error) { +func SeedFromDonorContext(ctx context.Context, donorPath, dstPath, projectPath string) (bool, error) { + if _, err := os.Stat(dstPath); err == nil { + return false, nil + } + + if err := os.MkdirAll(filepath.Dir(dstPath), 0o755); err != nil { + return false, fmt.Errorf("create dst directory: %w", err) + } + + seedLock, err := indexlock.Acquire(ctx, dstPath+".seed.lock") + if err != nil { + return false, fmt.Errorf("acquire seed lock: %w", err) + } + defer seedLock.Release() + + // Another seeder may have published the destination while this caller + // waited for the seed lock. Exit before opening or copying the donor. if _, err := os.Stat(dstPath); err == nil { return false, nil } @@ -39,55 +71,115 @@ func SeedFromDonor(donorPath, dstPath string) (bool, error) { // A missing or empty root_hash means the donor is still being built // (or was interrupted), so its data is incomplete and potentially // inconsistent — skip seeding to avoid inheriting corrupt state. - db, err := sql.Open("sqlite3", donorPath+"?mode=ro") + // mode=ro prevents sqlite from creating an empty donor if it disappears + // after discovery. VACUUM INTO reads a transactionally consistent snapshot, + // including committed WAL content, without modifying the live donor. + db, err := sql.Open("sqlite3", sqliteFileDSN(donorPath, "ro")) if err != nil { return false, fmt.Errorf("open donor: %w", err) } var rootHash sql.NullString - _ = db.QueryRow("SELECT value FROM project_meta WHERE key = 'root_hash'").Scan(&rootHash) - - // Checkpoint the WAL so the main DB file is self-contained. - _, _ = db.Exec("PRAGMA wal_checkpoint(TRUNCATE)") - _ = db.Close() + if err := db.QueryRowContext(ctx, "SELECT value FROM project_meta WHERE key = 'root_hash'").Scan(&rootHash); err != nil && !errors.Is(err, sql.ErrNoRows) { + _ = db.Close() + return false, fmt.Errorf("read donor metadata: %w", err) + } if !rootHash.Valid || rootHash.String == "" { + if err := db.Close(); err != nil { + return false, fmt.Errorf("close donor: %w", err) + } return false, nil } - if err := os.MkdirAll(filepath.Dir(dstPath), 0o755); err != nil { - return false, fmt.Errorf("create dst directory: %w", err) + // The seed lock makes a fixed temp name safe. A crash can leave at most this + // one file behind, and the next attempt removes and replaces it instead of + // accumulating full-size index.db.seed-* orphans. + tmp := dstPath + ".seed-tmp" + removeSQLiteFiles(tmp) + defer func() { + removeSQLiteFiles(tmp) + }() + + if _, err := db.ExecContext(ctx, "VACUUM INTO ?", tmp); err != nil { + _ = db.Close() + return false, fmt.Errorf("snapshot donor: %w", err) + } + if err := db.Close(); err != nil { + return false, fmt.Errorf("close donor: %w", err) + } + if err := os.Chmod(tmp, 0o600); err != nil { + return false, fmt.Errorf("set seed temp permissions: %w", err) } - // Atomic copy: write to temp file then rename. - tmp := dstPath + ".seed-tmp" - if err := copyFile(donorPath, tmp); err != nil { - _ = os.Remove(tmp) - return false, fmt.Errorf("copy donor: %w", err) + // Root hashes use relative paths, so an unchanged sibling worktree can + // return early from EnsureFresh without rewriting metadata. Stamp the new + // owner before publication so project-scoped purge targets the right index. + if err := setSeedProjectPath(ctx, tmp, projectPath); err != nil { + return false, fmt.Errorf("set seed project path: %w", err) } - if err := os.Rename(tmp, dstPath); err != nil { - _ = os.Remove(tmp) - return false, fmt.Errorf("rename seed: %w", err) + seeded, err := publishSeed(tmp, dstPath, os.Link, os.Rename) + if err != nil { + return false, err } + return seeded, nil +} - return true, nil +func sqliteFileDSN(path, mode string) string { + return (&url.URL{ + Scheme: "file", + Path: filepath.ToSlash(path), + RawQuery: "mode=" + mode, + }).String() } -func copyFile(src, dst string) error { - in, err := os.Open(src) +func setSeedProjectPath(ctx context.Context, dbPath, projectPath string) error { + db, err := sql.Open("sqlite3", sqliteFileDSN(dbPath, "rw")) if err != nil { return err } - defer func() { _ = in.Close() }() - - out, err := os.Create(dst) - if err != nil { + if _, err := db.ExecContext(ctx, + `INSERT INTO project_meta (key, value) VALUES ('project_path', ?) + ON CONFLICT(key) DO UPDATE SET value = excluded.value`, + projectPath, + ); err != nil { + _ = db.Close() return err } - defer func() { _ = out.Close() }() - - if _, err := io.Copy(out, in); err != nil { + if _, err := db.ExecContext(ctx, "PRAGMA wal_checkpoint(TRUNCATE)"); err != nil { + _ = db.Close() return err } - return out.Close() + return db.Close() +} + +func removeSQLiteFiles(path string) { + for _, suffix := range []string{"", "-wal", "-shm"} { + _ = os.Remove(path + suffix) + } +} + +// publishSeed prefers create-if-absent hard-link publication. Filesystems that +// do not support hard links (for example exFAT and some network/overlay mounts) +// fall back to atomic rename while the caller holds the seed lock. +func publishSeed(tmp, dst string, link, rename func(string, string) error) (bool, error) { + if err := link(tmp, dst); err == nil { + return true, nil + } else if errors.Is(err, os.ErrExist) { + return false, nil + } + + // A non-ErrExist link failure may mean the filesystem has no hard-link + // support. Do not overwrite a destination created by a non-cooperating + // process before attempting the portable rename fallback. + if _, err := os.Stat(dst); err == nil { + return false, nil + } + if err := rename(tmp, dst); err != nil { + if errors.Is(err, os.ErrExist) { + return false, nil + } + return false, fmt.Errorf("publish seed: link and rename fallback failed: %w", err) + } + return true, nil } diff --git a/internal/index/seed_test.go b/internal/index/seed_test.go index c3764b44..45f06112 100644 --- a/internal/index/seed_test.go +++ b/internal/index/seed_test.go @@ -16,9 +16,16 @@ package index import ( "context" + "database/sql" + "errors" "os" "path/filepath" + "sync" "testing" + "time" + + "github.com/ory/lumen/internal/indexlock" + "github.com/ory/lumen/internal/store" ) func TestSeedFromDonor_CopiesDB(t *testing.T) { @@ -44,7 +51,8 @@ func Hello() {} // Seed to a new path. dstPath := filepath.Join(t.TempDir(), "sub", "seeded.db") - seeded, err := SeedFromDonor(donorPath, dstPath) + seedProjectDir := t.TempDir() + seeded, err := SeedFromDonor(donorPath, dstPath, seedProjectDir) if err != nil { t.Fatal(err) } @@ -66,6 +74,92 @@ func Hello() {} if status.IndexedFiles == 0 { t.Fatal("expected seeded DB to have indexed files") } + seedMeta, err := store.ReadMetaAt(dstPath, "project_path") + if err != nil { + t.Fatal(err) + } + if seedMeta["project_path"] != seedProjectDir { + t.Fatalf("seeded project_path = %q, want %q", seedMeta["project_path"], seedProjectDir) + } +} + +func TestSeedFromDonor_SnapshotsCommittedWALWithActiveWriter(t *testing.T) { + projectDir := t.TempDir() + writeGoFile(t, projectDir, "main.go", "package main\n\nfunc Hello() {}\n") + + donorPath := filepath.Join(t.TempDir(), "donor.db") + emb := &mockEmbedder{dims: 4, model: "test-model"} + idx, err := NewIndexer(donorPath, emb, 0) + if err != nil { + t.Fatal(err) + } + if _, err := idx.Index(context.Background(), projectDir, false, nil); err != nil { + t.Fatal(err) + } + if err := idx.Close(); err != nil { + t.Fatal(err) + } + + writer, err := sql.Open("sqlite3", sqliteFileDSN(donorPath, "rw")) + if err != nil { + t.Fatal(err) + } + writer.SetMaxOpenConns(1) + t.Cleanup(func() { _ = writer.Close() }) + if _, err := writer.Exec("PRAGMA journal_mode=WAL"); err != nil { + t.Fatal(err) + } + if _, err := writer.Exec( + `INSERT INTO project_meta (key, value) VALUES ('snapshot_marker', 'committed')`, + ); err != nil { + t.Fatal(err) + } + + // Keep a second write uncommitted while seeding. The snapshot must include + // the committed WAL record without observing this in-flight transaction. + tx, err := writer.Begin() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = tx.Rollback() }) + if _, err := tx.Exec( + `INSERT INTO project_meta (key, value) VALUES ('snapshot_uncommitted', 'hidden')`, + ); err != nil { + t.Fatal(err) + } + + dstPath := filepath.Join(t.TempDir(), "seeded.db") + seeded, err := SeedFromDonor(donorPath, dstPath, projectDir) + if err != nil { + t.Fatal(err) + } + if !seeded { + t.Fatal("expected seeded=true") + } + + meta, err := store.ReadMetaAt(dstPath, "snapshot_marker", "snapshot_uncommitted") + if err != nil { + t.Fatal(err) + } + if meta["snapshot_marker"] != "committed" { + t.Fatalf("snapshot marker = %q, want committed", meta["snapshot_marker"]) + } + if _, ok := meta["snapshot_uncommitted"]; ok { + t.Fatal("snapshot included an uncommitted donor write") + } + + seedDB, err := sql.Open("sqlite3", sqliteFileDSN(dstPath, "ro")) + if err != nil { + t.Fatal(err) + } + defer func() { _ = seedDB.Close() }() + var integrity string + if err := seedDB.QueryRow("PRAGMA integrity_check").Scan(&integrity); err != nil { + t.Fatal(err) + } + if integrity != "ok" { + t.Fatalf("seed integrity_check = %q, want ok", integrity) + } } func TestSeedFromDonor_DstExists(t *testing.T) { @@ -79,7 +173,7 @@ func TestSeedFromDonor_DstExists(t *testing.T) { t.Fatal(err) } - seeded, err := SeedFromDonor(donorPath, dstPath) + seeded, err := SeedFromDonor(donorPath, dstPath, t.TempDir()) if err != nil { t.Fatal(err) } @@ -94,6 +188,184 @@ func TestSeedFromDonor_DstExists(t *testing.T) { } } +func TestSeedFromDonor_MissingDonorIsNotCreated(t *testing.T) { + donorPath := filepath.Join(t.TempDir(), "missing.db") + dstPath := filepath.Join(t.TempDir(), "seeded.db") + + seeded, err := SeedFromDonor(donorPath, dstPath, t.TempDir()) + if err == nil { + t.Fatal("expected missing donor to return an error") + } + if seeded { + t.Fatal("expected seeded=false for missing donor") + } + if _, statErr := os.Stat(donorPath); !os.IsNotExist(statErr) { + t.Fatalf("missing donor was created: stat error = %v", statErr) + } + if _, statErr := os.Stat(dstPath); !os.IsNotExist(statErr) { + t.Fatalf("destination was created: stat error = %v", statErr) + } +} + +func TestSeedFromDonor_ConcurrentSeedersExactlyOneWins(t *testing.T) { + // Build a real, complete donor DB. + projectDir := t.TempDir() + writeGoFile(t, projectDir, "main.go", `package main + +func Hello() {} +`) + + donorPath := filepath.Join(t.TempDir(), "donor.db") + emb := &mockEmbedder{dims: 4, model: "test-model"} + idx, err := NewIndexer(donorPath, emb, 0) + if err != nil { + t.Fatal(err) + } + if _, err := idx.Index(context.Background(), projectDir, false, nil); err != nil { + t.Fatal(err) + } + if err := idx.Close(); err != nil { + t.Fatal(err) + } + + // Many goroutines race to seed the same destination. + dstPath := filepath.Join(t.TempDir(), "sub", "seeded.db") + const n = 8 + var ( + wg sync.WaitGroup + mu sync.Mutex + wins int + firstErr error + ) + wg.Add(n) + for i := 0; i < n; i++ { + go func() { + defer wg.Done() + seeded, err := SeedFromDonor(donorPath, dstPath, projectDir) + mu.Lock() + defer mu.Unlock() + if err != nil && firstErr == nil { + firstErr = err + } + if seeded { + wins++ + } + }() + } + wg.Wait() + + if firstErr != nil { + t.Fatalf("concurrent SeedFromDonor returned error: %v", firstErr) + } + if wins != 1 { + t.Fatalf("expected exactly one seeder to win, got %d", wins) + } + + // No seed temp files should be left behind. The advisory lock file itself + // may remain on disk, just like the regular index lock file. + entries, err := os.ReadDir(filepath.Dir(dstPath)) + if err != nil { + t.Fatal(err) + } + for _, e := range entries { + if e.Name() == filepath.Base(dstPath)+".seed-tmp" { + t.Fatalf("unexpected leftover file after concurrent seeding: %q", e.Name()) + } + } + + // The published DB must be usable. + idx2, err := NewIndexer(dstPath, emb, 0) + if err != nil { + t.Fatalf("seeded DB is not openable: %v", err) + } + defer func() { _ = idx2.Close() }() + status, err := idx2.Status(projectDir) + if err != nil { + t.Fatal(err) + } + if status.IndexedFiles == 0 { + t.Fatal("expected seeded DB to have indexed files") + } +} + +func TestSeedFromDonor_WaitsBeforeCopying(t *testing.T) { + projectDir := t.TempDir() + writeGoFile(t, projectDir, "main.go", "package main\n\nfunc Hello() {}\n") + + donorPath := filepath.Join(t.TempDir(), "donor.db") + emb := &mockEmbedder{dims: 4, model: "test-model"} + idx, err := NewIndexer(donorPath, emb, 0) + if err != nil { + t.Fatal(err) + } + if _, err := idx.Index(context.Background(), projectDir, false, nil); err != nil { + t.Fatal(err) + } + if err := idx.Close(); err != nil { + t.Fatal(err) + } + + dstPath := filepath.Join(t.TempDir(), "seeded.db") + held, err := indexlock.Acquire(context.Background(), dstPath+".seed.lock") + if err != nil { + t.Fatal(err) + } + + done := make(chan error, 1) + go func() { + _, err := SeedFromDonor(donorPath, dstPath, projectDir) + done <- err + }() + + select { + case err := <-done: + t.Fatalf("seeder returned before lock was released: %v", err) + case <-time.After(100 * time.Millisecond): + } + if _, err := os.Stat(dstPath + ".seed-tmp"); !os.IsNotExist(err) { + t.Fatalf("seeder copied before acquiring the seed lock: stat error = %v", err) + } + + held.Release() + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(5 * time.Second): + t.Fatal("seeder did not finish after lock release") + } +} + +func TestPublishSeed_FallsBackToRename(t *testing.T) { + dir := t.TempDir() + tmp := filepath.Join(dir, "index.db.seed-tmp") + dst := filepath.Join(dir, "index.db") + if err := os.WriteFile(tmp, []byte("seed"), 0o600); err != nil { + t.Fatal(err) + } + + seeded, err := publishSeed( + tmp, + dst, + func(_, _ string) error { return errors.New("hard links unsupported") }, + os.Rename, + ) + if err != nil { + t.Fatal(err) + } + if !seeded { + t.Fatal("expected rename fallback to publish the seed") + } + content, err := os.ReadFile(dst) + if err != nil { + t.Fatal(err) + } + if string(content) != "seed" { + t.Fatalf("published content = %q, want seed", content) + } +} + func TestSeedFromDonor_IncompleteDonor(t *testing.T) { // Create a donor DB that has the schema but no root_hash (simulates // a donor whose first indexing pass hasn't finished yet). @@ -109,7 +381,7 @@ func TestSeedFromDonor_IncompleteDonor(t *testing.T) { } dstPath := filepath.Join(t.TempDir(), "seeded.db") - seeded, err := SeedFromDonor(donorPath, dstPath) + seeded, err := SeedFromDonor(donorPath, dstPath, t.TempDir()) if err != nil { t.Fatal(err) } @@ -147,7 +419,7 @@ func Hello() {} // Seed to new path. dstPath := filepath.Join(t.TempDir(), "seeded.db") - if _, err := SeedFromDonor(donorPath, dstPath); err != nil { + if _, err := SeedFromDonor(donorPath, dstPath, projectDir); err != nil { t.Fatal(err) } diff --git a/internal/indexlock/lock.go b/internal/indexlock/lock.go index ccf8da9d..ea887661 100644 --- a/internal/indexlock/lock.go +++ b/internal/indexlock/lock.go @@ -6,7 +6,10 @@ package indexlock import ( + "context" + "errors" "os" + "time" "github.com/gofrs/flock" ) @@ -41,6 +44,26 @@ func TryAcquire(lockPath string) (*Lock, error) { return &Lock{fl: fl}, nil } +// Acquire waits for an exclusive lock on lockPath or for ctx to be cancelled. +// The OS releases the lock if the process exits, so a crashed holder cannot +// leave a stale logical lock behind. +func Acquire(ctx context.Context, lockPath string) (*Lock, error) { + fl := flock.New(lockPath) + locked, err := fl.TryLockContext(ctx, 25*time.Millisecond) + if err != nil { + _ = fl.Close() + return nil, err + } + if !locked { + _ = fl.Close() + if err := ctx.Err(); err != nil { + return nil, err + } + return nil, errors.New("lock not acquired") + } + return &Lock{fl: fl}, nil +} + // IsHeld reports whether another process currently holds an exclusive lock on // lockPath. Returns true on any error (fail-closed: callers skip work rather // than risk concurrent writes). Does NOT create the lock file — if it doesn't