Gate blob reads on in-progress writes to prevent spurious NotFound

.plzconfig

@@ -2,7 +2,7 @@
 version = >=17.0.0
 
 [build]
-path = /usr/local/go/bin:/usr/local/bin:/usr/bin:/bin
+path = /usr/local/go/bin:/usr/local/bin:/usr/bin:/bin:/sbin:/usr/sbin
 
 [buildconfig]
 local-host = 127.0.0.1

elan/rpc/BUILD

@@ -57,6 +57,15 @@ go_test(
     ],
 )
 
+go_test(
+    name = "inflight_test",
+    srcs = ["inflight_test.go"],
+    deps = [
+        ":rpc",
+        "///third_party/go/github.com_stretchr_testify//require",
+    ],
+)
+
 genrule(
     name = "test_data",
     cmd = [

elan/rpc/inflight.go

@@ -0,0 +1,55 @@
+package rpc
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+// defaultWriteTimeout is the maximum time a reader will wait for an in-progress
+// write to complete before proceeding anyway.
+const defaultWriteTimeout = 10 * time.Minute
+
+// inflightWrites tracks blob writes that are currently in progress so that
+// concurrent readers can block until the write completes rather than getting
+// a spurious NotFound.
+type inflightWrites struct {
+	mu struct {
+		sync.Mutex
+		blobs map[string]context.Context
+	}
+}
+
+func newInflightWrites() *inflightWrites {
+	w := &inflightWrites{}
+	w.mu.blobs = make(map[string]context.Context)
+	return w
+}
+
+// startWrite registers a blob write in progress. The returned cancel function
+// must be called when the write completes (success or failure) — use defer.
+// The write context inherits from the caller so parent cancellation propagates.
+// A default timeout of defaultWriteTimeout is applied as a safety net.
+func (w *inflightWrites) startWrite(ctx context.Context, hash string) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithTimeout(ctx, defaultWriteTimeout)
+	w.mu.Lock()
+	w.mu.blobs[hash] = ctx
+	w.mu.Unlock()
+	return ctx, func() {
+		w.mu.Lock()
+		delete(w.mu.blobs, hash)
+		w.mu.Unlock()
+		cancel()
+	}
+}
+
+// waitForWrite blocks until any in-progress write for the given hash completes.
+// If no write is in progress, it returns immediately.
+func (w *inflightWrites) waitForWrite(hash string) {
+	w.mu.Lock()
+	ctx, ok := w.mu.blobs[hash]
+	w.mu.Unlock()
+	if ok {
+		<-ctx.Done()
+	}
+}

elan/rpc/inflight_test.go

@@ -0,0 +1,129 @@
+package rpc
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWaitForWrite_NoInflight(t *testing.T) {
+	w := newInflightWrites()
+	// Should return immediately when nothing is in flight.
+	doneCtx, doneFunc := context.WithCancel(t.Context())
+	go func() {
+		w.waitForWrite("abc123")
+		doneFunc()
+	}()
+	select {
+	case <-doneCtx.Done():
+	case <-time.After(120 * time.Second):
+		t.Fatal("waitForWrite blocked when no write was in progress")
+	}
+}
+
+func TestWaitForWrite_BlocksUntilDone(t *testing.T) {
+	w := newInflightWrites()
+	_, finish := w.startWrite(t.Context(), "abc123")
+
+	var order []string
+	var mu sync.Mutex
+	record := func(s string) {
+		mu.Lock()
+		order = append(order, s)
+		mu.Unlock()
+	}
+
+	doneCtx, doneFunc := context.WithCancel(t.Context())
+	go func() {
+		w.waitForWrite("abc123")
+		record("read")
+		doneFunc()
+	}()
+
+	// Give the reader goroutine time to block.
+	time.Sleep(50 * time.Millisecond)
+	record("write")
+	finish()
+
+	select {
+	case <-doneCtx.Done():
+	case <-t.Context().Done():
+		t.Fatal("reader never unblocked")
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+	require.Equal(t, []string{"write", "read"}, order)
+}
+
+func TestWaitForWrite_DifferentDigests(t *testing.T) {
+	w := newInflightWrites()
+	_, finish := w.startWrite(t.Context(), "abc123")
+	defer finish()
+
+	// A different digest should not block.
+	doneCtx, doneFunc := context.WithCancel(t.Context())
+	go func() {
+		w.waitForWrite("def456")
+		doneFunc()
+	}()
+	select {
+	case <-doneCtx.Done():
+	case <-t.Context().Done():
+		t.Fatal("waitForWrite blocked on a different digest")
+	}
+}
+
+func TestWaitForWrite_ParentCancellation(t *testing.T) {
+	w := newInflightWrites()
+	ctx, cancel := context.WithCancel(t.Context())
+	_, finish := w.startWrite(ctx, "abc123")
+	defer finish()
+
+	doneCtx, doneFunc := context.WithCancel(t.Context())
+	go func() {
+		w.waitForWrite("abc123")
+		doneFunc()
+	}()
+
+	// Cancel the parent context — reader should unblock.
+	cancel()
+	select {
+	case <-doneCtx.Done():
+	case <-t.Context().Done():
+		t.Fatal("reader did not unblock after parent context cancellation")
+	}
+}
+
+func TestWaitForWrite_MultipleReaders(t *testing.T) {
+	w := newInflightWrites()
+	_, finish := w.startWrite(t.Context(), "abc123")
+
+	const numReaders = 10
+	doneCtx, doneFunc := context.WithCancel(t.Context())
+	var wg sync.WaitGroup
+	wg.Add(numReaders)
+	for range numReaders {
+		go func() {
+			w.waitForWrite("abc123")
+			wg.Done()
+		}()
+	}
+
+	// All readers should be blocked. Finish the write.
+	time.Sleep(50 * time.Millisecond)
+	finish()
+
+	go func() {
+		wg.Wait()
+		doneFunc()
+	}()
+	select {
+	case <-doneCtx.Done():
+	case <-t.Context().Done():
+		t.Fatal("not all readers unblocked")
+	}
+}

elan/rpc/rpc.go

@@ -179,6 +179,7 @@ func createServer(storage string, parallelism int, maxDirCacheSize, maxKnownBlob
 		decompressor:   dec,
 		readRedis:      readRedis,
 		largeBlobSize:  largeBlobSize,
+		inflight:       newInflightWrites(),
 	}
 }
 
@@ -222,6 +223,7 @@ type server struct {
 	decompressor             *zstd.Decoder
 	readRedis                *redis.Client
 	largeBlobSize            int64
+	inflight                 *inflightWrites
 }
 
 func (s *server) GetCapabilities(ctx context.Context, req *pb.GetCapabilitiesRequest) (*pb.ServerCapabilities, error) {
@@ -549,6 +551,7 @@ func (s *server) Read(req *bs.ReadRequest, srv bs.ByteStream_ReadServer) error {
 }
 
 func (s *server) readCompressed(ctx context.Context, prefix string, digest *pb.Digest, compressed bool, offset, limit int64) (io.ReadCloser, bool, error) {
+	s.inflight.waitForWrite(digest.Hash)
 	if prefix != "cas" {
 		if compressed {
 			return nil, false, fmt.Errorf("Attempted to do a compressed read for non-CAS prefix %s", prefix) // This is a programming error and shouldn't happen.
@@ -687,6 +690,7 @@ func (s *server) readAllBlobBatched(ctx context.Context, prefix string, digest *
 }
 
 func (s *server) readAllBlobCompressed(ctx context.Context, digest *pb.Digest, key string, batched, compressed bool) ([]byte, error) {
+	s.inflight.waitForWrite(digest.Hash)
 	if digest.SizeBytes > s.largeBlobSize {
 		s.limiter <- struct{}{}
 		defer func() { <-s.limiter }()
@@ -724,6 +728,8 @@ func (s *server) compressedKey(prefix string, digest *pb.Digest, compressed bool
 }
 
 func (s *server) writeBlob(ctx context.Context, prefix string, digest *pb.Digest, r io.Reader, compressed bool) error {
+	_, done := s.inflight.startWrite(ctx, digest.Hash)
+	defer done()
 	key := s.compressedKey(prefix, digest, compressed)
 	if s.isEmpty(digest) || s.blobExists(ctx, prefix, digest, compressed, true) {
 		// Read and discard entire content; there is no need to update.
@@ -774,6 +780,8 @@ func (s *server) writeBlob(ctx context.Context, prefix string, digest *pb.Digest
 }
 
 func (s *server) writeAll(ctx context.Context, digest *pb.Digest, data []byte, compressed bool) error {
+	_, done := s.inflight.startWrite(ctx, digest.Hash)
+	defer done()
 	if digest.SizeBytes > s.largeBlobSize {
 		s.limiter <- struct{}{}
 		defer func() { <-s.limiter }()

go.mod

@@ -1,44 +1,44 @@
 module github.com/thought-machine/please-servers
 
-go 1.24.0
+go 1.25.0
 
 require (
 	cloud.google.com/go/profiler v0.4.0
 	cloud.google.com/go/pubsub v1.33.0
 	cloud.google.com/go/storage v1.36.0
 	github.com/bazelbuild/remote-apis v0.0.0-20230411132548-35aee1c4a425
 	github.com/bazelbuild/remote-apis-sdks v0.0.0-20230419185642-269815af5db1
-	github.com/dgraph-io/ristretto v0.1.1
+	github.com/dgraph-io/ristretto v0.2.0
 	github.com/dustin/go-humanize v1.0.1
 	github.com/go-redis/redis/v8 v8.11.5
 	github.com/golang/protobuf v1.5.4
 	github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.0
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.1
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/hashicorp/go-retryablehttp v0.7.7
-	github.com/klauspost/compress v1.17.4
-	github.com/mostynb/go-grpc-compression v1.2.2
+	github.com/klauspost/compress v1.18.6
+	github.com/mostynb/go-grpc-compression v1.2.3
 	github.com/peterebden/go-cli-init/v4 v4.0.2
 	github.com/peterebden/go-copyfile v0.0.0-20200424115000-bc0baf74909c
 	github.com/peterebden/go-sri v1.1.1
 	github.com/prometheus/client_golang v1.18.0
 	github.com/prometheus/common v0.45.0
 	github.com/shirou/gopsutil v3.21.11+incompatible
-	github.com/sirupsen/logrus v1.9.3
+	github.com/sirupsen/logrus v1.9.4
 	github.com/stretchr/testify v1.11.1
 	github.com/thought-machine/http-admin v1.1.1
-	go.uber.org/automaxprocs v1.5.3
+	go.uber.org/automaxprocs v1.6.0
 	gocloud.dev v0.36.0
-	golang.org/x/crypto v0.46.0
-	golang.org/x/exp v0.0.0-20240103183307-be819d1f06fc
-	golang.org/x/sync v0.19.0
-	golang.org/x/time v0.5.0
+	golang.org/x/crypto v0.53.0
+	golang.org/x/exp v0.0.0-20260611194520-c48552f49976
+	golang.org/x/sync v0.21.0
+	golang.org/x/time v0.15.0
 	google.golang.org/api v0.155.0
 	google.golang.org/genproto v0.0.0-20240102182953-50ed04b92917
 	google.golang.org/genproto/googleapis/bytestream v0.0.0-20240102182953-50ed04b92917
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217
 	google.golang.org/grpc v1.79.3
-	google.golang.org/protobuf v1.36.10
+	google.golang.org/protobuf v1.36.11
 	gopkg.in/op/go-logging.v1 v1.0.0-20160211212156-b2cb9fa56473
 )
 
@@ -83,11 +83,11 @@ require (
 	go.opentelemetry.io/otel v1.39.0 // indirect
 	go.opentelemetry.io/otel/metric v1.39.0 // indirect
 	go.opentelemetry.io/otel/trace v1.39.0 // indirect
-	golang.org/x/net v0.48.0 // indirect
+	golang.org/x/net v0.55.0 // indirect
 	golang.org/x/oauth2 v0.34.0 // indirect
-	golang.org/x/sys v0.39.0 // indirect
-	golang.org/x/term v0.38.0 // indirect
-	golang.org/x/text v0.32.0 // indirect
+	golang.org/x/sys v0.46.0 // indirect
+	golang.org/x/term v0.44.0 // indirect
+	golang.org/x/text v0.38.0 // indirect
 	golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect