feat: add subscriptions hook

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

@@ -9,6 +9,7 @@ import (
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/astparser"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/asttransform"
 	grpcdatasource "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/grpc_datasource"
+	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/federation"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/operationreport"
 )
@@ -103,6 +104,13 @@ type SingleTypeField struct {
 	FieldName string
 }
 
+// SubscriptionOnStartFn defines a hook function that is called when a subscription starts.
+// It receives the resolve context and the input of the subscription.
+// The function can return a boolean indicating if the subscription should be closed, and an error.
+// The error is propagated to the client.
+// If close is true, the subscription is closed.
+type SubscriptionOnStartFn func(ctx *resolve.Context, input []byte) (close bool, err error)
+
 type SubscriptionConfiguration struct {
 	URL           string
 	Header        http.Header
@@ -119,6 +127,7 @@ type SubscriptionConfiguration struct {
 	// these headers by itself.
 	ForwardedClientHeaderRegularExpressions []RegularExpression
 	WsSubProtocol                           string
+	SubscriptionOnStartFns                  []SubscriptionOnStartFn
 }
 
 type FetchConfiguration struct {

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

@@ -450,7 +450,8 @@ func (p *Planner[T]) ConfigureSubscription() plan.SubscriptionConfiguration {
 	return plan.SubscriptionConfiguration{
 		Input: string(input),
 		DataSource: &SubscriptionSource{
-			client: p.subscriptionClient,
+			client:                 p.subscriptionClient,
+			subscriptionOnStartFns: p.config.subscription.SubscriptionOnStartFns,
 		},
 		Variables:      p.variables,
 		PostProcessing: DefaultPostProcessingConfiguration,
@@ -1939,7 +1940,8 @@ type RegularExpression struct {
 }
 
 type SubscriptionSource struct {
-	client GraphQLSubscriptionClient
+	client                 GraphQLSubscriptionClient
+	subscriptionOnStartFns []SubscriptionOnStartFn
 }
 
 func (s *SubscriptionSource) AsyncStart(ctx *resolve.Context, id uint64, input []byte, updater resolve.SubscriptionUpdater) error {
@@ -1989,3 +1991,13 @@ func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte,
 	}
 	return s.client.UniqueRequestID(ctx, options, xxh)
 }
+
+func (s *SubscriptionSource) SubscriptionOnStart(ctx *resolve.Context, input []byte) (close bool, err error) {
+	for _, fn := range s.subscriptionOnStartFns {
+		close, err = fn(ctx, input)
+		if err != nil || close {
+			return close, err
+		}
+	}
+	return
+}

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

@@ -4019,7 +4019,7 @@ func TestGraphQLDataSource(t *testing.T) {
 				Trigger: resolve.GraphQLSubscriptionTrigger{
 					Input: []byte(`{"url":"wss://swapi.com/graphql","body":{"query":"subscription{remainingJedis}"}}`),
 					Source: &SubscriptionSource{
-						NewGraphQLSubscriptionClient(http.DefaultClient, http.DefaultClient, ctx),
+						client: NewGraphQLSubscriptionClient(http.DefaultClient, http.DefaultClient, ctx),
 					},
 					PostProcessing: DefaultPostProcessingConfiguration,
 				},
@@ -8902,6 +8902,59 @@ func TestSanitizeKey(t *testing.T) {
 	}
 }
 
+func TestSubscriptionSource_SubscriptionOnStart(t *testing.T) {
+
+	t.Run("SubscriptionOnStart calls subscriptionOnStartFns", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		type fnData struct {
+			ctx   *resolve.Context
+			input []byte
+		}
+
+		startFnCalled := make(chan fnData, 1)
+		subscriptionSource := SubscriptionSource{
+			subscriptionOnStartFns: []SubscriptionOnStartFn{
+				func(ctx *resolve.Context, input []byte) (bool, error) {
+					startFnCalled <- fnData{ctx, input}
+					return false, nil
+				},
+			},
+		}
+
+		close, err := subscriptionSource.SubscriptionOnStart(ctx, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`))
+		require.NoError(t, err)
+		assert.False(t, close)
+		var called fnData
+		select {
+		case called = <-startFnCalled:
+		case <-time.After(1 * time.Second):
+			t.Fatal("SubscriptionOnStartFn was not called")
+		}
+		assert.Equal(t, ctx, called.ctx)
+		assert.Equal(t, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`), called.input)
+	})
+
+	t.Run("SubscriptionOnStart calls subscriptionOnStartFns and returns error if one of the functions returns an error", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		subscriptionSource := SubscriptionSource{
+			subscriptionOnStartFns: []SubscriptionOnStartFn{
+				func(ctx *resolve.Context, input []byte) (bool, error) {
+					return false, errors.New("test error")
+				},
+			},
+		}
+
+		close, err := subscriptionSource.SubscriptionOnStart(ctx, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`))
+		assert.False(t, close)
+		require.Error(t, err)
+		assert.ErrorContains(t, err, "test error")
+	})
+}
+
 const interfaceSelectionSchema = `
 
 scalar String

v2/pkg/engine/resolve/context.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"io"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/wundergraph/astjson"
@@ -32,6 +33,9 @@ type Context struct {
 	fieldRenderer FieldValueRenderer
 
 	subgraphErrors error
+
+	emitEventRWMutex *sync.RWMutex
+	emitEventFn      func(data []byte) bool
 }
 
 type ExecutionOptions struct {
@@ -104,6 +108,22 @@ func (c *Context) SetEngineLoaderHooks(hooks LoaderHooks) {
 	c.LoaderHooks = hooks
 }
 
+// TryEmitSubscriptionUpdate emits a subscription update to the client
+// Returns true if the update was emitted.
+func (c *Context) TryEmitSubscriptionUpdate(data []byte) bool {
+	if c.emitEventRWMutex == nil {
+		return false
+	}
+	c.emitEventRWMutex.RLock()
+	emitEventFn := c.emitEventFn
+	c.emitEventRWMutex.RUnlock()
+	if emitEventFn == nil {
+		return false
+	}
+
+	return emitEventFn(data)
+}
+
 type RateLimitOptions struct {
 	// Enable switches rate limiting on or off
 	Enable bool
@@ -205,6 +225,8 @@ func (c *Context) Free() {
 	c.subgraphErrors = nil
 	c.authorizer = nil
 	c.LoaderHooks = nil
+	c.emitEventRWMutex = nil
+	c.emitEventFn = nil
 }
 
 type traceStartKey struct{}

v2/pkg/engine/resolve/datasource.go

@@ -26,3 +26,16 @@ type AsyncSubscriptionDataSource interface {
 	AsyncStop(id uint64)
 	UniqueRequestID(ctx *Context, input []byte, xxh *xxhash.Digest) (err error)
 }
+
+// HookableSubscriptionDataSource is a hookable interface for subscription data sources.
+// It is used to call a function when a subscription is started.
+// This is useful for data sources that need to do some work when a subscription is started,
+// e.g. to establish a connection to the data source or to emit updates to the client.
+// The function is called with the context and the input of the subscription.
+// The function is called before the subscription is started and can be used to emit updates to the client.
+type HookableSubscriptionDataSource interface {
+	// SubscriptionOnStart is called when a new subscription is created
+	// If close is true, the subscription is closed.
+	// If an error is returned, the error is propagated to the client.
+	SubscriptionOnStart(ctx *Context, input []byte) (close bool, err error)
+}

v2/pkg/engine/resolve/resolve.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"sync"
 	"time"
 
 	"github.com/buger/jsonparser"
@@ -488,7 +489,7 @@ func (r *Resolver) handleEvent(event subscriptionEvent) {
 	case subscriptionEventKindRemoveClient:
 		r.handleRemoveClient(event.id.ConnectionID)
 	case subscriptionEventKindTriggerUpdate:
-		r.handleTriggerUpdate(event.triggerID, event.data)
+		r.handleTriggerUpdate(event.triggerID, event.data, event.pinnedSubscription)
 	case subscriptionEventKindTriggerComplete:
 		r.handleTriggerComplete(event.triggerID)
 	case subscriptionEventKindTriggerInitialized:
@@ -569,6 +570,56 @@ func (r *Resolver) handleTriggerComplete(triggerID uint64) {
 	r.completeTrigger(triggerID)
 }
 
+func (r *Resolver) handleAddSubscriptionInitialHook(triggerID uint64, add *addSubscription, initialHookCtx *Context, s *sub) {
+	// Set the emitEventFn to synchronously emit the event to the subscription
+	if initialHookCtx.emitEventRWMutex == nil {
+		initialHookCtx.emitEventRWMutex = &sync.RWMutex{}
+	}
+	initialHookCtx.emitEventRWMutex.Lock()
+	initialHookCtx.emitEventFn = func(data []byte) bool {
+		if r.shouldSkipEvent(initialHookCtx, s, data) {
+			return true
+		}
+		r.executeSubscriptionUpdate(initialHookCtx, s, data)
+		return true
+	}
+	initialHookCtx.emitEventRWMutex.Unlock()
+
+	// After the initial hook is executed, there could be other work items in the workChan that need to be processed.
+	// But if the hook starts a go routine, it could emit events while other work items are being processed.
+	// Also the workChan could be closed while the go routine is running.
+	// We have to avoid concurrent executeSubscriptionUpdate to the same subscription.
+	// To avoid this, after the initial hook is executed, we add the work item to the events channel with a pinned subscription.
+	defer func() {
+		initialHookCtx.emitEventRWMutex.Lock()
+		initialHookCtx.emitEventFn = func(data []byte) bool {
+			r.events <- subscriptionEvent{
+				triggerID:          triggerID,
+				kind:               subscriptionEventKindTriggerUpdate,
+				data:               data,
+				pinnedSubscription: s,
+			}
+			return true
+		}
+		initialHookCtx.emitEventRWMutex.Unlock()
+	}()
+
+	hook, ok := add.resolve.Trigger.Source.(HookableSubscriptionDataSource)
+	if !ok {
+		return
+	}
+
+	close, errStartHook := hook.SubscriptionOnStart(initialHookCtx, add.input)
+	if errStartHook != nil {
+		r.asyncErrorWriter.WriteError(initialHookCtx, errStartHook, add.resolve.Response, add.writer)
+	}
+	if close {
+		r.closeTrigger(triggerID, SubscriptionCloseKindNormal)
+		return
+	}
+
+}
+
 func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription) {
 	var (
 		err error
@@ -590,6 +641,20 @@ func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription)
 		s.heartbeat = true
 	}
 
+	// Clone the context to avoid modifying the original context.
+	// This is important because the original context is changed in the parent thread
+	// and could cause data races.
+	initialHookCtx := add.ctx.clone(add.ctx.ctx)
+
+	// Send handleAddSubscriptionInitialHook as the first work sent to the worker to process the hooks
+	// before getting updates from the subscription. This operation is never blocking because workChan
+	// is buffered and this is the first work item, before anyone is allowed to send other work items.
+	// This is important because the hooks need to be executed before the other updates are processed.
+	s.workChan <- workItem{
+		fn:    func() { r.handleAddSubscriptionInitialHook(triggerID, add, initialHookCtx, s) },
+		final: false,
+	}
+
 	// Start the dedicated worker goroutine where the subscription updates are processed
 	// and writes are written to the client in a single threaded manner
 	go s.startWorker()
@@ -757,7 +822,42 @@ func (r *Resolver) handleRemoveClient(id int64) {
 	}
 }
 
-func (r *Resolver) handleTriggerUpdate(id uint64, data []byte) {
+func (r *Resolver) shouldSkipEvent(ctx *Context, sub *sub, data []byte) bool {
+	if err := ctx.ctx.Err(); err != nil {
+		return true // no need to schedule an event update when the client already disconnected
+	}
+	skip, err := sub.resolve.Filter.SkipEvent(ctx, data, r.triggerUpdateBuf)
+	if err != nil {
+		r.asyncErrorWriter.WriteError(ctx, err, sub.resolve.Response, sub.writer)
+		return true
+	}
+	if skip {
+		return true
+	}
+	return false
+}
+
+func (r *Resolver) handleTriggerUpdateQueueWorkItem(c *Context, s *sub, data []byte) {
+	if r.shouldSkipEvent(c, s, data) {
+		return
+	}
+
+	fn := func() {
+		r.executeSubscriptionUpdate(c, s, data)
+	}
+
+	select {
+	case <-r.ctx.Done():
+		// Skip sending all events if the resolver is shutting down
+		return
+	case <-c.ctx.Done():
+		// Skip sending the event if the client disconnected
+	case s.workChan <- workItem{fn, false}:
+		// Send the event to the subscription worker
+	}
+}
+
+func (r *Resolver) handleTriggerUpdate(id uint64, data []byte, pinnedSubscription *sub) {
 	trig, ok := r.triggers[id]
 	if !ok {
 		return
@@ -766,33 +866,13 @@ func (r *Resolver) handleTriggerUpdate(id uint64, data []byte) {
 		fmt.Printf("resolver:trigger:update:%d\n", id)
 	}
 
-	for c, s := range trig.subscriptions {
-		c, s := c, s
-		if err := c.ctx.Err(); err != nil {
-			continue // no need to schedule an event update when the client already disconnected
-		}
-		skip, err := s.resolve.Filter.SkipEvent(c, data, r.triggerUpdateBuf)
-		if err != nil {
-			r.asyncErrorWriter.WriteError(c, err, s.resolve.Response, s.writer)
-			continue
-		}
-		if skip {
-			continue
-		}
-
-		fn := func() {
-			r.executeSubscriptionUpdate(c, s, data)
-		}
+	if pinnedSubscription != nil {
+		r.handleTriggerUpdateQueueWorkItem(pinnedSubscription.ctx, pinnedSubscription, data)
+		return
+	}
 
-		select {
-		case <-r.ctx.Done():
-			// Skip sending all events if the resolver is shutting down
-			return
-		case <-c.ctx.Done():
-			// Skip sending the event if the client disconnected
-		case s.workChan <- workItem{fn, false}:
-			// Send the event to the subscription worker
-		}
+	for c, s := range trig.subscriptions {
+		r.handleTriggerUpdateQueueWorkItem(c, s, data)
 	}
 }
 
@@ -1231,12 +1311,13 @@ func (s *subscriptionUpdater) Close(kind SubscriptionCloseKind) {
 }
 
 type subscriptionEvent struct {
-	triggerID       uint64
-	id              SubscriptionIdentifier
-	kind            subscriptionEventKind
-	data            []byte
-	addSubscription *addSubscription
-	closeKind       SubscriptionCloseKind
+	triggerID          uint64
+	id                 SubscriptionIdentifier
+	kind               subscriptionEventKind
+	data               []byte
+	addSubscription    *addSubscription
+	pinnedSubscription *sub
+	closeKind          SubscriptionCloseKind
 }
 
 type addSubscription struct {