fix(pubsub): fix inefficient trigger id generation

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go

@@ -14,6 +14,7 @@ import (
 	"unicode"
 
 	"github.com/buger/jsonparser"
+	"github.com/cespare/xxhash/v2"
 	"github.com/jensneuse/abstractlogger"
 	"github.com/pkg/errors"
 	"github.com/tidwall/sjson"
@@ -1997,6 +1998,11 @@ type SubscriptionSource struct {
 	subscriptionOnStartFns []SubscriptionOnStartFn
 }
 
+func (s *SubscriptionSource) HashTriggerInput(input []byte, xxh *xxhash.Digest) error {
+	_, err := xxh.Write(input)
+	return err
+}
+
 // Start the subscription. The updater is called on new events. Start needs to be called in a separate goroutine.
 func (s *SubscriptionSource) Start(ctx *resolve.Context, headers http.Header, input []byte, updater resolve.SubscriptionUpdater) error {
 	var options GraphQLSubscriptionOptions

v2/pkg/engine/resolve/datasource.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"net/http"
 
+	"github.com/cespare/xxhash/v2"
+
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient"
 )
 
@@ -16,6 +18,9 @@ type SubscriptionDataSource interface {
 	// Start is called when a new subscription is created. It establishes the connection to the data source.
 	// The updater is used to send updates to the client. Deduplication of the request must be done before calling this method.
 	Start(ctx *Context, headers http.Header, input []byte, updater SubscriptionUpdater) error
+	// HashTriggerInput writes identity-relevant fields of the subscription into xxh.
+	// The resolver appends the subgraph headers hash afterward to produce the final trigger ID.
+	HashTriggerInput(input []byte, xxh *xxhash.Digest) error
 }
 
 // HookableSubscriptionDataSource is a hookable interface for subscription data sources.

v2/pkg/engine/resolve/resolve.go

@@ -1274,9 +1274,15 @@ func (r *Resolver) UnsubscribeClient(connectionID ConnectionID) error {
 // prepareTrigger safely gets the headers for the trigger Subgraph and computes the hash across headers and input
 // the generated hash is the unique triggerID
 // the headers must be forwarded to the DataSource to create the trigger
-func (r *Resolver) prepareTrigger(ctx *Context, sourceName string, input []byte) (headers http.Header, triggerID uint64) {
+func (r *Resolver) prepareTrigger(ctx *Context, sourceName string, input []byte, source SubscriptionDataSource) (
+	headers http.Header, triggerID uint64, err error) {
 	keyGen := pool.Hash64.Get()
-	_, _ = keyGen.Write(input)
+	defer pool.Hash64.Put(keyGen)
+
+	if err = source.HashTriggerInput(input, keyGen); err != nil {
+		return nil, 0, err
+	}
+
 	if ctx.SubgraphHeadersBuilder != nil {
 		var headersHash uint64
 		headers, headersHash = ctx.SubgraphHeadersBuilder.HeadersForSubgraph(sourceName)
@@ -1286,9 +1292,10 @@ func (r *Resolver) prepareTrigger(ctx *Context, sourceName string, input []byte)
 			_, _ = keyGen.Write(b[:])
 		}
 	}
+
 	triggerID = keyGen.Sum64()
-	pool.Hash64.Put(keyGen)
-	return headers, triggerID
+
+	return headers, triggerID, nil
 }
 
 func (r *Resolver) ResolveGraphQLSubscription(ctx *Context, subscription *GraphQLSubscription, writer SubscriptionResponseWriter) error {
@@ -1328,7 +1335,11 @@ func (r *Resolver) ResolveGraphQLSubscription(ctx *Context, subscription *GraphQ
 		return nil
 	}
 
-	headers, triggerID := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input)
+	headers, triggerID, err := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input, subscription.Trigger.Source)
+	if err != nil {
+		msg := []byte(`{"errors":[{"message":"failed to prepare subscription trigger"}]}`)
+		return writeFlushComplete(writer, msg)
+	}
 	id := SubscriptionIdentifier{
 		ConnectionID:   NewConnectionID(),
 		SubscriptionID: 0,
@@ -1422,7 +1433,11 @@ func (r *Resolver) AsyncResolveGraphQLSubscription(ctx *Context, subscription *G
 		return err
 	}
 
-	headers, triggerID := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input)
+	headers, triggerID, err := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input, subscription.Trigger.Source)
+	if err != nil {
+		msg := []byte(`{"errors":[{"message":"failed to prepare subscription trigger"}]}`)
+		return writeFlushComplete(writer, msg)
+	}
 
 	return r.addSubscription(triggerID, &addSubscription{
 		ctx:        ctx,

v2/pkg/engine/resolve/resolve_test.go

@@ -14,6 +14,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cespare/xxhash/v2"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -5766,6 +5767,11 @@ type _fakeStream struct {
 	subscriptionOnStartFn func(ctx StartupHookContext, input []byte) (err error)
 }
 
+func (f *_fakeStream) HashTriggerInput(input []byte, xxh *xxhash.Digest) error {
+	_, err := xxh.Write(input)
+	return err
+}
+
 func (f *_fakeStream) SubscriptionOnStart(ctx StartupHookContext, input []byte) (err error) {
 	if f.subscriptionOnStartFn == nil {
 		return nil
@@ -8020,6 +8026,11 @@ type startFailStream struct {
 	subBReady chan struct{}
 }
 
+func (s *startFailStream) HashTriggerInput(input []byte, xxh *xxhash.Digest) error {
+	_, err := xxh.Write(input)
+	return err
+}
+
 func (s *startFailStream) Start(_ *Context, _ http.Header, _ []byte, _ SubscriptionUpdater) error {
 	<-s.subBReady
 	return errors.New("connection refused")
@@ -8104,6 +8115,11 @@ type hookFailStream struct {
 	sourceStarted  atomic.Bool
 }
 
+func (s *hookFailStream) HashTriggerInput(input []byte, xxh *xxhash.Digest) error {
+	_, err := xxh.Write(input)
+	return err
+}
+
 func (s *hookFailStream) Start(_ *Context, _ http.Header, _ []byte, _ SubscriptionUpdater) error {
 	s.sourceStarted.Store(true)
 	select {}

v2/pkg/engine/resolve/resolver_subscription_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cespare/xxhash/v2"
 	"github.com/stretchr/testify/require"
 )
 
@@ -66,6 +67,11 @@ type FakeSource struct {
 	interval time.Duration
 }
 
+func (f *FakeSource) HashTriggerInput(input []byte, xxh *xxhash.Digest) error {
+	_, err := xxh.Write(input)
+	return err
+}
+
 func (f *FakeSource) Start(ctx *Context, headers http.Header, input []byte, updater SubscriptionUpdater) error {
 	go func() {
 		for i, u := range f.updates {