Add crates index manager abstraction

The full crates.io index corpus is too large and growing too fast to
realistically assume we can have access to it all at once for inference
purposes. This new component enables us to keep a fixed-size cache of
snapshots on disk while rotating the least used snapshots out when other
snapshots are requested. We unconditionally keep the current (i.e.
non-archived) index cloned given it is the newest and most likely to be
requested during inference.

The interface also supports loading index portions from an existing
filesystem which means it can also be used transiently to mediate access
to these resources.

Author: msuozzo

internal/safememfs/safememfs.go

@@ -0,0 +1,138 @@
+// Copyright 2025 Google LLC
+// SPDX-License-Identifier: Apache-2.0
+
+package safememfs
+
+import (
+	"os"
+	"sync"
+
+	"github.com/go-git/go-billy/v5"
+	"github.com/go-git/go-billy/v5/memfs"
+)
+
+// SafeMemory is a thread-safe wrapper for any billy.Filesystem.
+type SafeMemory struct {
+	fs billy.Filesystem
+	mu *sync.RWMutex
+}
+
+// New creates a new thread-safe in-memory filesystem.
+func New() *SafeMemory {
+	return &SafeMemory{
+		fs: memfs.New(),
+		mu: &sync.RWMutex{},
+	}
+}
+
+func (s *SafeMemory) Chroot(path string) (billy.Filesystem, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	newFs, err := s.fs.Chroot(path)
+	if err != nil {
+		return nil, err
+	}
+	return &SafeMemory{
+		fs: newFs,
+		mu: s.mu, // NOTE: same mutex
+	}, nil
+}
+
+func (s *SafeMemory) Root() string {
+	return "/"
+}
+
+// --- Write Operations (use exclusive Lock) ---
+
+func (s *SafeMemory) Create(filename string) (billy.File, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.Create(filename)
+}
+
+func (s *SafeMemory) OpenFile(filename string, flag int, perm os.FileMode) (billy.File, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.OpenFile(filename, flag, perm)
+}
+
+func (s *SafeMemory) MkdirAll(path string, perm os.FileMode) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.MkdirAll(path, perm)
+}
+
+func (s *SafeMemory) Rename(from, to string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.Rename(from, to)
+}
+
+func (s *SafeMemory) Remove(filename string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.Remove(filename)
+}
+
+func (s *SafeMemory) Symlink(target, link string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.Symlink(target, link)
+}
+
+func (s *SafeMemory) TempFile(dir, prefix string) (billy.File, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.fs.TempFile(dir, prefix)
+}
+
+// --- Read Operations (use shared RLock) ---
+
+func (s *SafeMemory) Open(filename string) (billy.File, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.fs.OpenFile(filename, os.O_RDONLY, 0)
+}
+
+func (s *SafeMemory) Stat(filename string) (os.FileInfo, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.fs.Stat(filename)
+}
+
+func (s *SafeMemory) Lstat(filename string) (os.FileInfo, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.fs.Lstat(filename)
+}
+
+func (s *SafeMemory) ReadDir(path string) ([]os.FileInfo, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.fs.ReadDir(path)
+}
+
+func (s *SafeMemory) Readlink(link string) (string, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.fs.Readlink(link)
+}
+
+func (s *SafeMemory) Join(elem ...string) string {
+	// No lock needed as this is a pure function
+	return s.fs.Join(elem...)
+}
+
+// Capabilities forwards the call to the underlying filesystem.
+// No lock needed as it doesn't modify state.
+func (s *SafeMemory) Capabilities() billy.Capability {
+	if capable, ok := s.fs.(billy.Capable); ok {
+		return capable.Capabilities()
+	}
+	// Default capabilities if the underlying fs doesn't implement billy.Capable
+	return 0
+}
+
+// Ensure the wrapper implements the full interface.
+var _ billy.Filesystem = (*SafeMemory)(nil)
+var _ billy.Capable = (*SafeMemory)(nil)

pkg/registry/cratesio/index/fetch.go

@@ -0,0 +1,115 @@
+// Copyright 2025 Google LLC
+// SPDX-License-Identifier: Apache-2.0
+
+package index
+
+import (
+	"context"
+	"regexp"
+	"sort"
+
+	"github.com/go-git/go-billy/v5"
+	"github.com/go-git/go-git/v5"
+	"github.com/go-git/go-git/v5/config"
+	"github.com/go-git/go-git/v5/plumbing"
+	"github.com/go-git/go-git/v5/plumbing/cache"
+	"github.com/go-git/go-git/v5/storage/filesystem"
+	"github.com/pkg/errors"
+)
+
+var (
+	currentIndexURL = "https://github.com/rust-lang/crates.io-index.git"
+	archiveIndexURL = "https://github.com/rust-lang/crates.io-index-archive.git"
+)
+
+var snapshotBranchRegex = regexp.MustCompile(`^refs/heads/snapshot-(\d{4}-\d{2}-\d{2})$`)
+
+// ListAvailableSnapshots queries the archive repository for available snapshots
+// Snapshots are returned as their associated RFC3339 date e.g. 2025-06-14.
+func ListAvailableSnapshots(ctx context.Context) ([]string, error) {
+	// Create a remote reference to list branches
+	rem := git.NewRemote(nil, &config.RemoteConfig{URLs: []string{archiveIndexURL}})
+	// List the references
+	refs, err := rem.ListContext(ctx, &git.ListOptions{})
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to list remote refs")
+	}
+	var snapshots []string
+	for _, ref := range refs {
+		if matches := snapshotBranchRegex.FindStringSubmatch(ref.Name().String()); matches != nil {
+			snapshots = append(snapshots, matches[1])
+		}
+	}
+	sort.Strings(snapshots)
+	return snapshots, nil
+}
+
+// Fetcher defines how to fetch and update a repository index
+type Fetcher interface {
+	// Fetch clones the repository into the given filesystem
+	Fetch(ctx context.Context, fs billy.Filesystem) error
+	// Update updates an existing repository in the filesystem
+	Update(ctx context.Context, fs billy.Filesystem) error
+}
+
+// CurrentIndexFetcher fetches the current crates.io index
+type CurrentIndexFetcher struct{}
+
+func (f *CurrentIndexFetcher) Fetch(ctx context.Context, fs billy.Filesystem) error {
+	storer := filesystem.NewStorage(fs, cache.NewObjectLRUDefault())
+	_, err := git.CloneContext(ctx, storer, nil, &git.CloneOptions{
+		URL:           currentIndexURL,
+		ReferenceName: plumbing.Master,
+		SingleBranch:  true,
+		NoCheckout:    true,
+	})
+	if err != nil {
+		return errors.Wrap(err, "failed to clone current index")
+	}
+	// Nice-to-have: Set HEAD to track the remote since it will remain up-to-date on Update.
+	remoteMain := plumbing.NewRemoteReferenceName(git.DefaultRemoteName, "master")
+	err = storer.SetReference(plumbing.NewSymbolicReference(plumbing.HEAD, remoteMain))
+	if err != nil {
+		return errors.Wrap(err, "failed to configure HEAD")
+	}
+	return nil
+}
+
+func (f *CurrentIndexFetcher) Update(ctx context.Context, fs billy.Filesystem) error {
+	storer := filesystem.NewStorage(fs, cache.NewObjectLRUDefault())
+	repo, err := git.Open(storer, nil)
+	if err != nil {
+		return errors.Wrap(err, "failed to open repository")
+	}
+	err = repo.FetchContext(ctx, &git.FetchOptions{Force: true})
+	if err == git.NoErrAlreadyUpToDate {
+		return nil
+	} else if err != nil {
+		return errors.Wrap(err, "failed to fetch updates")
+	}
+	return nil
+}
+
+// SnapshotIndexFetcher fetches a specific snapshot branch
+type SnapshotIndexFetcher struct {
+	Date string
+}
+
+func (f *SnapshotIndexFetcher) Fetch(ctx context.Context, fs billy.Filesystem) error {
+	storer := filesystem.NewStorage(fs, cache.NewObjectLRUDefault())
+	_, err := git.CloneContext(ctx, storer, nil, &git.CloneOptions{
+		URL:           archiveIndexURL,
+		ReferenceName: plumbing.NewBranchReferenceName("snapshot-" + f.Date),
+		SingleBranch:  true,
+		NoCheckout:    true,
+	})
+	if err != nil {
+		return errors.Wrapf(err, "failed to clone snapshot %s", f.Date)
+	}
+	return nil
+}
+
+func (f *SnapshotIndexFetcher) Update(ctx context.Context, fs billy.Filesystem) error {
+	// Snapshots are immutable
+	return nil
+}