fix(resolve): fix flaky singleflight deduplication tests

v2/pkg/engine/resolve/inbound_request_singleflight.go

@@ -3,6 +3,7 @@ package resolve
 import (
 	"encoding/binary"
 	"sync"
+	"sync/atomic"
 
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/pool"
 )
@@ -46,8 +47,7 @@ type InflightRequest struct {
 	Err        error
 	ID         uint64
 
-	HasFollowers bool
-	Mu           sync.Mutex
+	followerCount atomic.Int32
 }
 
 // GetOrCreate creates a new InflightRequest or returns an existing (shared) one
@@ -90,9 +90,7 @@ func (r *InboundRequestSingleFlight) GetOrCreate(ctx *Context, response *GraphQL
 	inflight, shared := shard.m.LoadOrStore(key, request)
 	if shared {
 		request = inflight.(*InflightRequest)
-		request.Mu.Lock()
-		request.HasFollowers = true
-		request.Mu.Unlock()
+		request.followerCount.Add(1)
 		select {
 		case <-request.Done:
 			if request.Err != nil {
@@ -113,10 +111,7 @@ func (r *InboundRequestSingleFlight) FinishOk(req *InflightRequest, data []byte)
 	}
 	shard := r.shardFor(req.ID)
 	shard.m.Delete(req.ID)
-	req.Mu.Lock()
-	hasFollowers := req.HasFollowers
-	req.Mu.Unlock()
-	if hasFollowers {
+	if req.followerCount.Load() > 0 {
 		// optimization to only copy when we actually have to
 		req.Data = make([]byte, len(data))
 		copy(req.Data, data)

v2/pkg/engine/resolve/inbound_request_singleflight_test.go

@@ -2,8 +2,10 @@ package resolve
 
 import (
 	"context"
+	"runtime"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
 )
@@ -75,8 +77,7 @@ func TestInboundSingleFlight_FollowerReceivesLeaderError(t *testing.T) {
 	}
 
 	// The follower calls GetOrCreate which blocks on inflight.Done.
-	// We wait for HasFollowers to be set before calling FinishErr.
-	followerReady := make(chan struct{})
+	// We wait for followerCount to confirm it has entered before calling FinishErr.
 	var wg sync.WaitGroup
 	wg.Add(1)
 
@@ -85,25 +86,20 @@ func TestInboundSingleFlight_FollowerReceivesLeaderError(t *testing.T) {
 		followerCtx := NewContext(context.Background())
 		followerCtx.Request.ID = 2
 
-		// Signal that we're about to enter GetOrCreate. HasFollowers will be
-		// set inside GetOrCreate before the select blocks, so closing
-		// followerReady here is slightly early, but we poll HasFollowers below.
-		close(followerReady)
-
 		_, followerErr := sf.GetOrCreate(followerCtx, response)
 		if followerErr == nil {
 			t.Error("expected error from follower after leader FinishErr")
 		}
 	}()
 
-	<-followerReady
-	// Spin until the follower has actually registered (set HasFollowers)
-	for {
-		inflight.Mu.Lock()
-		ready := inflight.HasFollowers
-		inflight.Mu.Unlock()
-		if ready {
-			break
+	// Poll until the follower has actually registered inside GetOrCreate.
+	deadline := time.After(time.Second)
+	for inflight.followerCount.Load() < 1 {
+		select {
+		case <-deadline:
+			t.Fatal("timeout waiting for follower to enter singleflight")
+		default:
+			runtime.Gosched()
 		}
 	}
 

v2/pkg/engine/resolve/resolve_test.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -149,6 +150,40 @@ func (w *blockingWriter) String() string {
 	return w.buf.String()
 }
 
+// findAnyInflight iterates through all singleflight shards and returns
+// the first inflight request found. Used in tests to poll followerCount.
+func findAnyInflight(r *Resolver) *InflightRequest {
+	for i := range r.inboundRequestSingleFlight.shards {
+		var found *InflightRequest
+		r.inboundRequestSingleFlight.shards[i].m.Range(func(_, value any) bool {
+			found = value.(*InflightRequest)
+			return false
+		})
+		if found != nil {
+			return found
+		}
+	}
+	return nil
+}
+
+// waitForFollowerCount polls until the inflight request has at least count followers registered.
+func waitForFollowerCount(t *testing.T, r *Resolver, count int32) {
+	t.Helper()
+	deadline := time.After(time.Second)
+	for {
+		inflight := findAnyInflight(r)
+		if inflight != nil && inflight.followerCount.Load() >= count {
+			return
+		}
+		select {
+		case <-deadline:
+			t.Fatal("timeout waiting for followers to enter singleflight")
+		default:
+			runtime.Gosched()
+		}
+	}
+}
+
 type TestErrorWriter struct {
 }
 
@@ -4694,30 +4729,20 @@ func TestResolver_ArenaResolveGraphQLResponse_RequestDeduplication(t *testing.T)
 		t.Fatalf("timeout waiting for leader data source load")
 	}
 
-	startFollowers := make(chan struct{})
-	followersEntered := make(chan struct{}, requestCount-1)
-
 	for i := 1; i < requestCount; i++ {
 		go func(i int) {
 			defer wg.Done()
 			ctx := ctxTemplate
-			<-startFollowers
-			followersEntered <- struct{}{}
 			buf := &bytes.Buffer{}
 			info, err := r.ArenaResolveGraphQLResponse(&ctx, response, buf)
 			results[i] = result{info: info, output: buf.String(), err: err}
 		}(i)
 	}
 
-	close(startFollowers)
-
-	for i := 1; i < requestCount; i++ {
-		select {
-		case <-followersEntered:
-		case <-time.After(time.Second):
-			t.Fatalf("timeout waiting for follower %d to start", i)
-		}
-	}
+	// Wait until all followers have entered the singleflight (called LoadOrStore)
+	// before releasing the data source. This guarantees they join the leader's
+	// inflight request rather than creating their own.
+	waitForFollowerCount(t, r, int32(requestCount-1))
 
 	ds.Release()
 
@@ -4823,9 +4848,6 @@ func TestResolver_ArenaResolveGraphQLResponse_RequestDeduplication_SharedData(t
 		t.Fatalf("timeout waiting for leader data source load")
 	}
 
-	startFollowers := make(chan struct{})
-	followersEntered := make(chan struct{}, requestCount-1)
-
 	for i := 1; i < requestCount; i++ {
 		go func(i int) {
 			defer wg.Done()
@@ -4838,23 +4860,16 @@ func TestResolver_ArenaResolveGraphQLResponse_RequestDeduplication_SharedData(t
 					followerData.Store(i, data)
 				},
 			)
-			<-startFollowers
-			followersEntered <- struct{}{}
 			buf := &bytes.Buffer{}
 			info, err := r.ArenaResolveGraphQLResponse(&ctx, response, buf)
 			results[i] = result{info: info, output: buf.String(), err: err}
 		}(i)
 	}
 
-	close(startFollowers)
-
-	for i := 1; i < requestCount; i++ {
-		select {
-		case <-followersEntered:
-		case <-time.After(time.Second):
-			t.Fatalf("timeout waiting for follower %d to start", i)
-		}
-	}
+	// Wait until all followers have entered the singleflight (called LoadOrStore)
+	// before releasing the data source. This guarantees they join the leader's
+	// inflight request rather than creating their own.
+	waitForFollowerCount(t, r, int32(requestCount-1))
 
 	ds.Release()