Fix gateway registration wait lifetime

control/gateway/gateway.go

@@ -2,6 +2,7 @@ package gateway
 
 import (
 	"context"
+	"time"
 
 	"github.com/moby/buildkit/client/buildid"
 	"github.com/moby/buildkit/frontend/gateway"
@@ -40,6 +41,10 @@ func (gwf *GatewayForwarder) lookupForwarder(ctx context.Context) (gateway.LLBBr
 		return nil, errors.New("no buildid found in context")
 	}
 
+	// Match the gateway client's initial Ping budget. A late Solve must not lose
+	// its registration after 3s, but unknown build IDs must still be bounded.
+	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
+	defer cancel()
 	fwd, err := gwf.registrar.Get(ctx, bid)
 	if err != nil {
 		if errors.Is(err, context.Canceled) {

control/gateway/gateway_test.go

@@ -0,0 +1,59 @@
+package gateway
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/moby/buildkit/client/buildid"
+	"github.com/moby/buildkit/frontend/gateway"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestLookupRegistration(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		register bool
+		deadline time.Duration
+		elapsed  time.Duration
+	}{
+		{name: "late solve", register: true, elapsed: 4 * time.Second},
+		{name: "missing solve", elapsed: 15 * time.Second},
+		{name: "shorter caller deadline", deadline: time.Second, elapsed: time.Second},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				gw := NewGatewayForwarder()
+				outgoing := buildid.AppendToOutgoingContext(t.Context(), "build")
+				md, _ := metadata.FromOutgoingContext(outgoing)
+				ctx := metadata.NewIncomingContext(t.Context(), md)
+				if tc.deadline != 0 {
+					var cancel context.CancelFunc
+					ctx, cancel = context.WithTimeout(ctx, tc.deadline)
+					defer cancel()
+				}
+				bridge := &struct{ gateway.LLBBridgeForwarder }{}
+				if tc.register {
+					go func() {
+						time.Sleep(4 * time.Second)
+						gw.RegisterBuild(ctx, "build", bridge)
+					}()
+				}
+				start := time.Now()
+				got, err := gw.lookupForwarder(ctx)
+				if tc.register {
+					if err != nil || got != bridge {
+						t.Errorf("late registration: bridge=%v err=%v", got, err)
+					}
+				} else if !errors.Is(err, context.DeadlineExceeded) {
+					t.Errorf("missing registration: expected deadline, got %v", err)
+				}
+				if elapsed := time.Since(start); elapsed != tc.elapsed {
+					t.Errorf("lookup waited %s, want %s", elapsed, tc.elapsed)
+				}
+			})
+		})
+	}
+}

solver/llbsolver/solver.go

@@ -274,9 +274,8 @@ func (s *Solver) Solve(ctx context.Context, id string, sessionID string, req fro
 		fwd = gateway.NewBridgeForwarder(ctx, br, br, s.workerController.Infos(), req.FrontendInputs, sessionID, s.sm)
 		defer fwd.Discard()
 		// Register build before calling s.recordBuildHistory, because
-		// s.recordBuildHistory can block for several seconds on
-		// LeaseManager calls, and there is a fixed 3s timeout in
-		// GatewayForwarder on build registration.
+		// s.recordBuildHistory can block for several seconds on LeaseManager
+		// calls while the gateway client is waiting for its initial Ping.
 		s.gatewayForwarder.RegisterBuild(ctx, id, fwd)
 		defer s.gatewayForwarder.UnregisterBuild(context.Background(), id)
 	}

util/registrar/registrar.go

@@ -3,7 +3,6 @@ package registrar
 import (
 	"context"
 	"sync"
-	"time"
 )
 
 type Registrar[K comparable, V any] struct {
@@ -20,23 +19,30 @@ func New[K comparable, V any]() *Registrar[K, V] {
 // Register will register the value with the given id.
 // This value will persist until Discard is called with the same id.
 func (r *Registrar[K, V]) Register(id K, val V) {
-	reg := r.getOrCreateRegistrar(id, nil)
-	reg.Register(val, nil)
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.getOrCreateRegistrar(id).register(val, nil)
 }
 
-// Get will retrieve a registered value and will wait a small time period for that
-// value to appear if it hasn't been registered yet.
+// Get retrieves a registered value, waiting until it appears, is discarded, or
+// ctx is canceled. An unregistered value is removed when its last waiter leaves.
 func (r *Registrar[K, V]) Get(ctx context.Context, id K) (v V, _ error) {
-	onCreate := func(reg *registrarValue[V]) {
-		select {
-		case <-reg.notifyCh:
-			return
-		case <-time.After(3 * time.Second):
-			r.Discard(id)
-		}
+	if err := context.Cause(ctx); err != nil {
+		return v, err
 	}
-
-	reg := r.getOrCreateRegistrar(id, onCreate)
+	r.mu.Lock()
+	reg := r.getOrCreateRegistrar(id)
+	reg.waiters++
+	r.mu.Unlock()
+	defer func() {
+		r.mu.Lock()
+		defer r.mu.Unlock()
+		reg.waiters--
+		// Discard may have removed this entry and a new request reused its ID.
+		if reg.waiters == 0 && !reg.isSet && r.values[id] == reg {
+			delete(r.values, id)
+		}
+	}()
 
 	select {
 	case <-ctx.Done():
@@ -50,35 +56,24 @@ func (r *Registrar[K, V]) Get(ctx context.Context, id K) (v V, _ error) {
 // with Register.
 func (r *Registrar[K, V]) Discard(id K) {
 	r.mu.Lock()
+	defer r.mu.Unlock()
 	reg, ok := r.values[id]
 	delete(r.values, id)
-	r.mu.Unlock()
 
 	if ok {
 		var value V
-		reg.Register(value, context.Canceled)
+		reg.register(value, context.Canceled)
 	}
 }
 
-// getOrCreateRegistrar will create a registrar with the given id to be retrieved at a later time.
-// The same id will return the same registrar.
-//
-// If the registrar is newly created, the onCreate function is invoked in a separate goroutine
-// if it is present. If nil, this function is ignored.
-func (r *Registrar[K, V]) getOrCreateRegistrar(id K, onCreate func(*registrarValue[V])) *registrarValue[V] {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
+// getOrCreateRegistrar requires r.mu to be held.
+func (r *Registrar[K, V]) getOrCreateRegistrar(id K) *registrarValue[V] {
 	reg, ok := r.values[id]
 	if !ok {
 		reg = &registrarValue[V]{
 			notifyCh: make(chan struct{}),
 		}
 		r.values[id] = reg
-
-		if onCreate != nil {
-			go onCreate(reg)
-		}
 	}
 	return reg
 }
@@ -88,17 +83,15 @@ type registrarValue[V any] struct {
 	// the bridge is registered.
 	notifyCh chan struct{}
 
-	value V
-	err   error
-	isSet bool
-
-	mu sync.Mutex
+	value   V
+	err     error
+	isSet   bool
+	waiters int
 }
 
-func (r *registrarValue[V]) Register(value V, err error) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
+// register requires the owning Registrar's mutex to be held. Published values
+// never change; closing notifyCh makes them visible to waiting readers.
+func (r *registrarValue[V]) register(value V, err error) {
 	if r.isSet {
 		return
 	}

util/registrar/registrar_test.go

@@ -0,0 +1,131 @@
+package registrar
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"testing/synctest"
+	"time"
+)
+
+func TestDelayedRegistration(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		r := New[string, string]()
+		registered := make(chan struct{})
+		go func() {
+			time.Sleep(4 * time.Second)
+			r.Register("build", "bridge")
+			close(registered)
+		}()
+		ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
+		defer cancel()
+		value, err := r.Get(ctx, "build")
+		<-registered
+		if err != nil || value != "bridge" {
+			t.Fatalf("live request lost its pending registration: value=%q err=%v", value, err)
+		}
+	})
+}
+
+func TestAbandonedLookupIsRemoved(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		r := New[string, string]()
+		ctx, cancel := context.WithTimeout(t.Context(), time.Second)
+		defer cancel()
+		_, err := r.Get(ctx, "missing")
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("expected request deadline, got %v", err)
+		}
+		retained := len(r.values)
+		r.Discard("missing")
+		synctest.Wait()
+		if retained != 0 {
+			t.Fatal("abandoned lookup retained a registration")
+		}
+	})
+}
+
+func TestOneCanceledWaiterDoesNotCancelAnother(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		r := New[string, string]()
+		ctx, cancel := context.WithCancel(t.Context())
+		go func() {
+			_, err := r.Get(ctx, "build")
+			if !errors.Is(err, context.Canceled) {
+				t.Errorf("expected cancellation, got %v", err)
+			}
+		}()
+		go func() {
+			value, err := r.Get(t.Context(), "build")
+			if err != nil || value != "bridge" {
+				t.Errorf("surviving waiter: value=%q err=%v", value, err)
+			}
+		}()
+		synctest.Wait()
+		cancel()
+		synctest.Wait()
+		r.Register("build", "bridge")
+		synctest.Wait()
+	})
+}
+
+func TestDiscardWakesWaitersWithoutDeletingReplacement(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		r := New[string, string]()
+		go func() {
+			_, err := r.Get(t.Context(), "build")
+			if !errors.Is(err, context.Canceled) {
+				t.Errorf("expected discarded build, got %v", err)
+			}
+		}()
+		synctest.Wait()
+		r.Discard("build")
+		r.Register("build", "replacement")
+		synctest.Wait()
+		value, err := r.Get(t.Context(), "build")
+		if err != nil || value != "replacement" {
+			t.Fatalf("replacement registration lost: value=%q err=%v", value, err)
+		}
+	})
+}
+
+func TestRegisteredValuePersistsUntilDiscard(t *testing.T) {
+	r := New[string, string]()
+	r.Register("build", "bridge")
+	r.Register("build", "must-not-replace")
+	for range 2 {
+		value, err := r.Get(t.Context(), "build")
+		if err != nil || value != "bridge" {
+			t.Fatalf("registered value changed: value=%q err=%v", value, err)
+		}
+	}
+	r.Discard("build")
+	if len(r.values) != 0 {
+		t.Fatal("discarded registration retained")
+	}
+}
+
+func TestConcurrentCancelAndRegister(t *testing.T) {
+	for range 100 {
+		r := New[string, string]()
+		ctx, cancel := context.WithCancel(t.Context())
+		var wg sync.WaitGroup
+		wg.Go(func() { r.Register("build", "bridge") })
+		wg.Go(cancel)
+		value, err := r.Get(ctx, "build")
+		if err != nil && !errors.Is(err, context.Canceled) {
+			t.Fatalf("lookup failed: %v", err)
+		}
+		if err == nil && value != "bridge" {
+			t.Fatalf("lookup returned unpublished value: %q", value)
+		}
+		wg.Wait()
+		lookup, stop := context.WithTimeout(t.Context(), time.Second)
+		value, err = r.Get(lookup, "build")
+		stop()
+		if err != nil || value != "bridge" {
+			t.Fatalf("canceled waiter removed completed registration: value=%q err=%v", value, err)
+		}
+	}
+}