fix(cache): make disk ranged reads revision-atomic

Disk.Open opened the file descriptor and read the ETag from BoltDB
without synchronisation, while diskWriter.Close renamed the body then
wrote metadata as two separate steps. A replacement interleaving those
reads returned one revision's bytes with another's ETag. With
ParallelGet this passes the chunk-0 consistency check yet splices old
and new bytes across later chunks, so per-chunk ETag pinning did not
actually bind the served body to the returned metadata.

Guard each key's body swap with a striped RWMutex: Open captures the
(descriptor, headers) pair under a read lock, and Close/Delete perform
the rename/remove plus metadata write under the write lock. os.Open
binds the descriptor to a stable inode, so the captured pair stays
revision-consistent for the descriptor's lifetime.

Closes #348

Co-authored-by: Amp <amp@ampcode.com>
Signed-off-by: Elizabeth Worstell <worstell@squareup.com>
Amp-Thread-ID: https://ampcode.com/threads/T-019f3d90-8c5a-757c-87ff-308404f3b2bd

Author: worstell

internal/cache/disk.go

@@ -2,6 +2,7 @@ package cache
 
 import (
 	"context"
+	"hash/fnv"
 	"io"
 	"io/fs"
 	"log/slog"
@@ -10,6 +11,7 @@ import (
 	"path/filepath"
 	"sort"
 	"strconv"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -45,8 +47,17 @@ type Disk struct {
 	runEviction  chan struct{}
 	stop         context.CancelFunc
 	evictionDone chan struct{}
+	// locks serialises each key's body swap against reads so Open captures a
+	// revision-consistent (file descriptor, headers) pair. Striped to keep
+	// cross-key contention negligible. Shared across namespaced copies (the
+	// slice header is copied, the mutexes are not).
+	locks []sync.RWMutex
 }
 
+// diskLockStripes bounds the number of per-key locks. A power of two so the
+// hash reduction is a mask.
+const diskLockStripes = 1024
+
 var _ Cache = (*Disk)(nil)
 
 // NewDisk creates a new disk-based cache instance.
@@ -119,6 +130,7 @@ func NewDisk(ctx context.Context, config DiskConfig) (*Disk, error) {
 		runEviction:  make(chan struct{}),
 		stop:         stop,
 		evictionDone: make(chan struct{}),
+		locks:        make([]sync.RWMutex, diskLockStripes),
 	}
 	disk.size.Store(size)
 
@@ -127,6 +139,16 @@ func NewDisk(ctx context.Context, config DiskConfig) (*Disk, error) {
 	return disk, nil
 }
 
+// keyLock returns the stripe guarding a given (namespace, key) body/metadata
+// pair. It is derived deterministically so every namespaced view of the same
+// disk agrees on the same lock.
+func (d *Disk) keyLock(namespace Namespace, key Key) *sync.RWMutex {
+	h := fnv.New32a()
+	_, _ = h.Write([]byte(namespace))
+	_, _ = h.Write(key[:])
+	return &d.locks[h.Sum32()&(diskLockStripes-1)]
+}
+
 func (d *Disk) String() string { return "disk:" + d.config.Root }
 
 func (d *Disk) Close() error {
@@ -211,19 +233,24 @@ func (d *Disk) Delete(_ context.Context, key Key) error {
 		expired = true
 	}
 
+	// Remove the body and its metadata as one critical section, excluding any
+	// concurrent Open's snapshot capture (#348).
+	lock := d.keyLock(d.namespace, key)
+	lock.Lock()
 	info, err := os.Stat(fullPath)
 	if err != nil {
+		lock.Unlock()
 		return errors.Errorf("failed to stat file: %w", err)
 	}
-
 	if err := os.Remove(fullPath); err != nil {
+		lock.Unlock()
 		return errors.Errorf("failed to remove file: %w", err)
 	}
-
-	// Remove metadata
 	if err := d.db.delete(d.namespace, key); err != nil {
+		lock.Unlock()
 		return errors.Errorf("failed to delete TTL metadata: %w", err)
 	}
+	lock.Unlock()
 
 	d.size.Add(-info.Size())
 
@@ -275,25 +302,42 @@ func (d *Disk) Open(ctx context.Context, key Key, opts ...Option) (io.ReadCloser
 	path := d.keyToPath(d.namespace, key)
 	fullPath := filepath.Join(d.config.Root, path)
 
+	// Capture the file descriptor and its metadata under the key's read lock,
+	// excluding the writer's rename+db.set critical section. os.Open binds the
+	// descriptor to a specific inode (stable across a later rename), so the
+	// pair captured here stays revision-consistent for the descriptor's whole
+	// lifetime — without this, a replacement landing between os.Open and
+	// getHeaders returns one revision's bytes with another's ETag, silently
+	// splicing versions across ParallelGet chunks (#348).
+	lock := d.keyLock(d.namespace, key)
+	lock.RLock()
 	f, err := os.Open(fullPath)
 	if err != nil {
+		lock.RUnlock()
 		return nil, nil, errors.Errorf("failed to open file: %w", err)
 	}
-
 	expiresAt, err := d.db.getTTL(d.namespace, key)
 	if err != nil {
+		lock.RUnlock()
 		return nil, nil, errors.Join(err, f.Close())
 	}
 
+	// Delete takes the write lock, so an expired entry must be released before
+	// cleanup. Checking expiry before reading headers keeps the original
+	// semantics: an expired entry (even one with bad/missing header metadata)
+	// is treated as a miss and deleted, rather than surfacing a hard error.
 	now := time.Now()
 	if now.After(expiresAt) {
+		lock.RUnlock()
 		return nil, nil, errors.Join(fs.ErrNotExist, f.Close(), d.Delete(ctx, key))
 	}
 
 	headers, err := d.db.getHeaders(d.namespace, key)
 	if err != nil {
+		lock.RUnlock()
 		return nil, nil, errors.Join(errors.Errorf("failed to get headers: %w", err), f.Close())
 	}
+	lock.RUnlock()
 
 	finfo, err := f.Stat()
 	if err != nil {
@@ -505,18 +549,24 @@ func (w *diskWriter) Close() error {
 		return errors.Errorf("failed to create directory: %w", err)
 	}
 
-	// Check if we're overwriting an existing file and subtract its size
+	// Swap the body and publish its metadata as one critical section so a
+	// concurrent Open never captures the new bytes with the old ETag, or vice
+	// versa (#348).
+	lock := w.disk.keyLock(w.namespace, w.key)
+	lock.Lock()
+	// Check if we're overwriting an existing file and subtract its size.
 	if info, err := os.Stat(w.path); err == nil {
 		w.disk.size.Add(-info.Size())
 	}
-
 	if err := os.Rename(w.tempPath, w.path); err != nil {
+		lock.Unlock()
 		return errors.Errorf("failed to rename temp file: %w", err)
 	}
-
 	if err := w.disk.db.set(w.key, w.namespace, w.expiresAt, w.headers); err != nil {
+		lock.Unlock()
 		return errors.Join(errors.Errorf("failed to set metadata: %w", err), os.Remove(w.path))
 	}
+	lock.Unlock()
 
 	w.disk.size.Add(w.size)
 

internal/cache/disk_test.go

@@ -1,12 +1,19 @@
 package cache_test
 
 import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
 	"log/slog"
+	"net/http"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/alecthomas/assert/v2"
+	"github.com/alecthomas/errors"
 
 	"github.com/block/cachew/internal/cache"
 	"github.com/block/cachew/internal/cache/cachetest"
@@ -52,3 +59,93 @@ func TestDiskCacheSoak(t *testing.T) {
 		TTL:              5 * time.Minute,
 	})
 }
+
+// TestDiskOpenRevisionAtomic asserts that Open never returns one revision's
+// bytes paired with another revision's ETag while the same key is concurrently
+// replaced (#348). The invariant checked mirrors production: the ETag returned
+// with a body must equal the hash of that body. Run with -race for best
+// coverage.
+func TestDiskOpenRevisionAtomic(t *testing.T) {
+	dir := t.TempDir()
+	_, ctx := logging.Configure(t.Context(), logging.Config{Level: slog.LevelError})
+	c, err := cache.NewDisk(ctx, cache.DiskConfig{Root: dir, MaxTTL: time.Hour})
+	assert.NoError(t, err)
+	defer c.Close()
+
+	key := cache.NewKey("revision-atomic")
+
+	// Each revision's raw ETag is the hash of its body, exactly as the higher
+	// cache layers derive content-hash ETags. A revision-consistent Open must
+	// therefore return headers whose ETag re-hashes to the served body.
+	write := func(rev int) {
+		body := bytes.Repeat([]byte{byte(rev)}, 4096+rev)
+		sum := sha256.Sum256(body)
+		raw := hex.EncodeToString(sum[:])
+		w, err := c.Create(ctx, key, http.Header{}, time.Hour, cache.WithETag(raw))
+		if err != nil {
+			t.Errorf("create: %v", err)
+			return
+		}
+		if _, err := w.Write(body); err != nil {
+			t.Errorf("write: %v", errors.Join(err, w.Abort(err)))
+			return
+		}
+		if err := w.Close(); err != nil {
+			t.Errorf("close: %v", err)
+		}
+	}
+
+	write(0) // seed so readers don't race an empty key
+
+	const writers, readsPerReader, readers = 4, 500, 4
+	var wg sync.WaitGroup
+	stop := make(chan struct{})
+
+	for range writers {
+		wg.Go(func() {
+			for rev := 1; ; rev++ {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				write(rev % 251) // stay in a single byte; distinct bodies
+			}
+		})
+	}
+
+	for range readers {
+		wg.Go(func() {
+			for range readsPerReader {
+				r, headers, err := c.Open(ctx, key)
+				if err != nil {
+					t.Errorf("open: %v", err)
+					continue
+				}
+				body, err := io.ReadAll(r)
+				_ = r.Close()
+				if err != nil {
+					t.Errorf("read: %v", err)
+					continue
+				}
+				sum := sha256.Sum256(body)
+				want, err := cache.FormatETag(hex.EncodeToString(sum[:]))
+				if err != nil {
+					t.Errorf("format etag: %v", err)
+					continue
+				}
+				if got := headers.Get(cache.ETagKey); got != want {
+					t.Errorf("etag/body mismatch: header %s does not match body hash %s (spliced revision)", got, want)
+				}
+			}
+		})
+	}
+
+	// Bound the writers so the test terminates even if the readers finish
+	// first; readers stop on their own after readsPerReader each.
+	go func() {
+		time.Sleep(time.Second)
+		close(stop)
+	}()
+	wg.Wait()
+}