registry: detach the token singleflight from the leader's context

docs/registry.md

@@ -166,6 +166,13 @@ Behavior:
   is consulted before any challenge is seen, so only a request-computable key
   keeps store and lookup symmetric; keying on the advertised scope would make a
   registry that formats it differently miss the cache on every block.
+- **Token exchanges are singleflighted** per cache key, and the shared flight
+  runs on a **bounded background context** (1 minute) — the same semantics as
+  `fetch.Coalescing`: a caller's own context bounds only how long that caller
+  waits, never the exchange itself. A cancelled leader therefore cannot poison
+  still-live followers with `context.Canceled` (issue #51), a completed flight
+  warms the cache even if no waiter remains, and a stalled token endpoint
+  cannot pin the cache key beyond the bound.
 - **Near-expiry tokens are not reused** (30 s leeway), and a token the registry
   stops accepting triggers re-authentication rather than a failed read.
 - **A `Basic` challenge is satisfied directly**, no token endpoint (Harbor and
@@ -250,14 +257,17 @@ asserted.
 | `TestProxyForwardsAuthorization` | client credentials reach the upstream on the pass-through path |
 | `TestClientAuthorizationTakesPrecedence` | a client's own Authorization is never overwritten/substituted by the operator credential |
 | `TestConcurrentColdRequestsShareOneExchange` | token exchange is singleflighted per repository |
+| **`TestLeaderCancellationDoesNotPoisonFollowers`** | **a cancelled singleflight leader cannot fail still-live followers; the cohort still costs one exchange** |
+| `TestCancelledFollowerDoesNotAbortFlight` | a departing follower does not cancel the shared exchange |
 | `TestExpiredTokensPurged` | the token cache sweeps expired entries and stays bounded |
 | `TestPercentEncodedPathReachesUpstreamVerbatim` | a %3F in a repository name is never decoded into a query |
 | `TestPullOnly` | POST/PUT/PATCH/DELETE → `405` with `Allow`, and never reach the upstream |
 | `TestNonV2PathRejected` | non-registry paths → `404` |
 | `TestUpstreamPathPrefixPreserved` | an upstream subpath is prepended |
 | **`TestBlobDedupAcrossRegistries`** | **the same digest via two mirrors/repos is fetched once** |
 
-Authentication (`auth_test.go`), against a fake registry that speaks the token
+Authentication (`auth_test.go`, plus the issue-pinning `issue16_test.go` and
+`issue51_test.go`), against a fake registry that speaks the token
 flow and derives its advertised scope per repository:
 
 | Test | Property guarded |

internal/registry/auth.go

@@ -351,7 +351,11 @@ func (a *AuthTransport) dropTokenIf(key, rejected string) {
 
 // fetchTokenShared singleflights the exchange per cache key: concurrent cold
 // pulls of one repository share one token endpoint round trip. The first
-// caller leads; followers wait for its result.
+// caller starts the flight; every caller — the leader included — then waits
+// for its result, and each caller's own ctx only bounds how long THAT caller
+// waits. The flight itself runs on a bounded background context
+// (runTokenExchange): one caller's cancellation must not poison the other
+// waiters (issue #51) — the same flight semantics as fetch.Coalescing.
 func (a *AuthTransport) fetchTokenShared(ctx context.Context, key string, ch challenge, cred Credential) (string, time.Time, error) {
 	a.mu.Lock()
 	// Double-check the cache: a sibling request may have completed the exchange
@@ -360,36 +364,57 @@ func (a *AuthTransport) fetchTokenShared(ctx context.Context, key string, ch cha
 		a.mu.Unlock()
 		return tok.value, tok.expiresAt, nil
 	}
-	if c, ok := a.inflight[key]; ok {
-		a.mu.Unlock()
-		select {
-		case <-c.done:
-			return c.value, c.expiresAt, c.err
-		case <-ctx.Done():
-			return "", time.Time{}, ctx.Err()
+	c, joined := a.inflight[key]
+	if !joined {
+		c = &tokenCall{done: make(chan struct{})}
+		if a.inflight == nil {
+			a.inflight = make(map[string]*tokenCall)
 		}
+		a.inflight[key] = c
 	}
-	c := &tokenCall{done: make(chan struct{})}
-	if a.inflight == nil {
-		a.inflight = make(map[string]*tokenCall)
-	}
-	a.inflight[key] = c
 	a.mu.Unlock()
+	if !joined {
+		go a.runTokenExchange(key, c, ch, cred)
+	}
+	select {
+	case <-c.done:
+		return c.value, c.expiresAt, c.err
+	case <-ctx.Done():
+		return "", time.Time{}, ctx.Err()
+	}
+}
 
-	c.value, c.expiresAt, c.err = a.fetchToken(ctx, ch, cred)
+// maxTokenExchange bounds one shared token-exchange flight: a stalled token
+// endpoint (accepted, then silent) must not pin the cache key forever. The
+// bound is generous — the exchange is one small JSON GET — and expiring it
+// fails the flight, unpublishes the key, and lets the next caller lead a
+// fresh exchange.
+const maxTokenExchange = 1 * time.Minute
+
+// runTokenExchange completes one in-flight exchange and publishes the result.
+// It runs detached from every caller's context (bounded by maxTokenExchange)
+// so a cancelled caller cannot abort work its peers still wait on; a
+// completed flight still warms the cache even if no waiter remains.
+//
+// The result goes to the cache BEFORE the in-flight marker is removed: a
+// follower arriving after the delete but before a later store would see
+// neither and lead a second exchange. Under this one critical section a new
+// entrant always sees either the in-flight call or the warm cache. done is
+// closed only after the result fields are written and the lock is released,
+// so waiters observe a happens-before edge on the result and nothing runs
+// under a.mu.
+func (a *AuthTransport) runTokenExchange(key string, c *tokenCall, ch challenge, cred Credential) {
+	flightCtx, cancel := context.WithTimeout(context.Background(), maxTokenExchange)
+	c.value, c.expiresAt, c.err = a.fetchToken(flightCtx, ch, cred)
+	cancel()
 
-	// Publish the result to the cache BEFORE removing the in-flight marker:
-	// a follower arriving after the delete but before a later store would see
-	// neither and lead a second exchange. Under this one critical section a
-	// new entrant always sees either the in-flight call or the warm cache.
 	a.mu.Lock()
 	if c.err == nil {
 		a.storeTokenLocked(key, c.value, c.expiresAt)
 	}
 	delete(a.inflight, key)
 	a.mu.Unlock()
 	close(c.done)
-	return c.value, c.expiresAt, c.err
 }
 
 func tokenKey(host, scope string) string { return host + "|" + scope }

internal/registry/issue51_test.go

@@ -0,0 +1,182 @@
+package registry
+
+// Regression tests for issue #51: a cancelled singleflight leader used to
+// poison every still-live follower with context.Canceled, because the shared
+// token exchange ran on the leader's request context.
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"runtime"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// stalledTokenTransport stands in for the token endpoint: the exchange stalls
+// until release is closed. Like a real transport it honors request-context
+// cancellation — that fidelity is what makes the test fail on the old code,
+// where the flight ran on the leader's context.
+type stalledTokenTransport struct {
+	entered   chan struct{} // closed when the first exchange request arrives
+	release   chan struct{} // closing lets the exchange complete
+	exchanges atomic.Int64
+}
+
+func (s *stalledTokenTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	if s.exchanges.Add(1) == 1 {
+		close(s.entered)
+	}
+	select {
+	case <-s.release:
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Status:     "200 OK",
+			Header:     make(http.Header),
+			Body:       io.NopCloser(strings.NewReader(`{"token":"flight-token","expires_in":300}`)),
+		}, nil
+	case <-req.Context().Done():
+		return nil, req.Context().Err()
+	}
+}
+
+// TestLeaderCancellationDoesNotPoisonFollowers orchestrates the exact issue
+// scenario: the elected leader is cancelled while the token endpoint is
+// blocked, followers with live contexts have joined the same flight, and
+// every follower must still complete with the flight's token.
+//
+// Looped: a follower is observably parked on the flight's done channel only
+// from inside, so the test relies on a scheduler-yield rendezvous and repeats
+// across scheduling windows — the same pattern as
+// TestConcurrentColdRequestsShareOneExchange. On the old code the cohort
+// either inherits context.Canceled or pays a second exchange, so a round
+// fails as soon as any follower was parked before the cancel.
+func TestLeaderCancellationDoesNotPoisonFollowers(t *testing.T) {
+	for round := 0; round < 25; round++ {
+		st := &stalledTokenTransport{entered: make(chan struct{}), release: make(chan struct{})}
+		at := NewAuthTransport(st, nil)
+		ch := challenge{scheme: "bearer", realm: "http://token.test/token", service: "svc", scope: "repository:x:pull"}
+		const key = "token.test|repository:x:pull"
+
+		// The leader leads the exchange; entered proves the in-flight call is
+		// published (fetchTokenShared registers it before calling fetchToken).
+		leaderCtx, cancelLeader := context.WithCancel(context.Background())
+		leaderErr := make(chan error, 1)
+		go func() {
+			_, _, err := at.fetchTokenShared(leaderCtx, key, ch, Credential{})
+			leaderErr <- err
+		}()
+		<-st.entered
+
+		// Followers with live contexts join the same flight.
+		const followers = 4
+		var started sync.WaitGroup
+		followerErrs := make([]chan error, followers)
+		for i := range followerErrs {
+			followerErrs[i] = make(chan error, 1)
+			started.Add(1)
+			go func(out chan<- error) {
+				started.Done()
+				tok, _, err := at.fetchTokenShared(context.Background(), key, ch, Credential{})
+				if err == nil && tok != "flight-token" {
+					err = fmt.Errorf("token = %q, want the flight's token", tok)
+				}
+				out <- err
+			}(followerErrs[i])
+		}
+		started.Wait()
+		// Rendezvous: let the followers reach the wait on the flight's done
+		// channel before the leader is cancelled.
+		for i := 0; i < 100; i++ {
+			runtime.Gosched()
+		}
+
+		// Cancel the leader mid-exchange. The leader itself must observe its
+		// own cancellation...
+		cancelLeader()
+		select {
+		case err := <-leaderErr:
+			if !errors.Is(err, context.Canceled) {
+				t.Fatalf("round %d: leader err = %v, want context.Canceled", round, err)
+			}
+		case <-time.After(5 * time.Second):
+			t.Fatalf("round %d: leader did not observe its own cancellation", round)
+		}
+
+		// ...but no still-live follower may inherit it: once the endpoint
+		// responds, every follower completes with the shared token.
+		close(st.release)
+		for i, out := range followerErrs {
+			select {
+			case err := <-out:
+				if err != nil {
+					t.Fatalf("round %d: follower %d poisoned by the leader's cancellation: %v", round, i, err)
+				}
+			case <-time.After(5 * time.Second):
+				t.Fatalf("round %d: follower %d did not complete after the exchange was released", round, i)
+			}
+		}
+
+		// The cohort still cost exactly one exchange, and its result is cached.
+		if got := st.exchanges.Load(); got != 1 {
+			t.Fatalf("round %d: %d token exchanges, want 1 (singleflight intact)", round, got)
+		}
+		if tok, ok := at.cachedToken(key); !ok || tok != "flight-token" {
+			t.Fatalf("round %d: cached token = %q, %v; want the completed flight's token", round, tok, ok)
+		}
+	}
+}
+
+// TestCancelledFollowerDoesNotAbortFlight pins the symmetric direction: a
+// follower leaving early must not cancel the shared exchange either — the
+// leader and any remaining followers still get the token.
+func TestCancelledFollowerDoesNotAbortFlight(t *testing.T) {
+	st := &stalledTokenTransport{entered: make(chan struct{}), release: make(chan struct{})}
+	at := NewAuthTransport(st, nil)
+	ch := challenge{scheme: "bearer", realm: "http://token.test/token", service: "svc", scope: "repository:x:pull"}
+	const key = "token.test|repository:x:pull"
+
+	leaderRes := make(chan error, 1)
+	go func() {
+		tok, _, err := at.fetchTokenShared(context.Background(), key, ch, Credential{})
+		if err == nil && tok != "flight-token" {
+			err = fmt.Errorf("token = %q, want the flight's token", tok)
+		}
+		leaderRes <- err
+	}()
+	<-st.entered
+
+	quitterCtx, cancelQuitter := context.WithCancel(context.Background())
+	quitterErr := make(chan error, 1)
+	go func() {
+		_, _, err := at.fetchTokenShared(quitterCtx, key, ch, Credential{})
+		quitterErr <- err
+	}()
+	cancelQuitter()
+	select {
+	case err := <-quitterErr:
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("departing follower: err = %v, want context.Canceled", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("departing follower did not observe its own cancellation")
+	}
+
+	close(st.release)
+	select {
+	case err := <-leaderRes:
+		if err != nil {
+			t.Fatalf("leader failed after a follower left: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("leader did not complete after the exchange was released")
+	}
+	if got := st.exchanges.Load(); got != 1 {
+		t.Fatalf("%d token exchanges, want 1", got)
+	}
+}