feat: add epoll support to ws origin requests

v2/go.mod

@@ -37,6 +37,7 @@ require (
 
 require (
 	github.com/agnivade/levenshtein v1.1.1 // indirect
+	github.com/coder/websocket v1.8.12 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect

v2/go.sum

@@ -13,6 +13,8 @@ github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMU
 github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
 github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo=
+github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=

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

@@ -1738,6 +1738,8 @@ type GraphQLSubscriptionClient interface {
 	// Subscribe to the origin source. The implementation must not block the calling goroutine.
 	Subscribe(ctx *resolve.Context, options GraphQLSubscriptionOptions, updater resolve.SubscriptionUpdater) error
 	UniqueRequestID(ctx *resolve.Context, options GraphQLSubscriptionOptions, hash *xxhash.Digest) (err error)
+	SubscribeAsync(ctx *resolve.Context, id uint64, options GraphQLSubscriptionOptions, updater resolve.SubscriptionUpdater) error
+	Unsubscribe(id uint64)
 }
 
 type GraphQLSubscriptionOptions struct {
@@ -1763,6 +1765,22 @@ type SubscriptionSource struct {
 	client GraphQLSubscriptionClient
 }
 
+func (s *SubscriptionSource) AsyncStart(ctx *resolve.Context, id uint64, input []byte, updater resolve.SubscriptionUpdater) error {
+	var options GraphQLSubscriptionOptions
+	err := json.Unmarshal(input, &options)
+	if err != nil {
+		return err
+	}
+	if options.Body.Query == "" {
+		return resolve.ErrUnableToResolve
+	}
+	return s.client.SubscribeAsync(ctx, id, options, updater)
+}
+
+func (s *SubscriptionSource) AsyncStop(id uint64) {
+	s.client.Unsubscribe(id)
+}
+
 // 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, input []byte, updater resolve.SubscriptionUpdater) error {
 	var options GraphQLSubscriptionOptions

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

@@ -8834,6 +8834,14 @@ var errSubscriptionClientFail = errors.New("subscription client fail error")
 
 type FailingSubscriptionClient struct{}
 
+func (f *FailingSubscriptionClient) SubscribeAsync(ctx *resolve.Context, id uint64, options GraphQLSubscriptionOptions, updater resolve.SubscriptionUpdater) error {
+	return errSubscriptionClientFail
+}
+
+func (f *FailingSubscriptionClient) Unsubscribe(id uint64) {
+
+}
+
 func (f *FailingSubscriptionClient) Subscribe(ctx *resolve.Context, options GraphQLSubscriptionOptions, updater resolve.SubscriptionUpdater) error {
 	return errSubscriptionClientFail
 }

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

@@ -25,65 +25,68 @@ var (
 )
 
 type gqlSSEConnectionHandler struct {
-	conn    *http.Client
-	ctx     context.Context
-	log     log.Logger
-	options GraphQLSubscriptionOptions
+	conn                          *http.Client
+	requestContext, engineContext context.Context
+	log                           log.Logger
+	options                       GraphQLSubscriptionOptions
+	updater                       resolve.SubscriptionUpdater
 }
 
-func newSSEConnectionHandler(ctx *resolve.Context, conn *http.Client, opts GraphQLSubscriptionOptions, l log.Logger) *gqlSSEConnectionHandler {
+func newSSEConnectionHandler(requestContext, engineContext context.Context, conn *http.Client, updater resolve.SubscriptionUpdater, options GraphQLSubscriptionOptions, l log.Logger) *gqlSSEConnectionHandler {
 	return &gqlSSEConnectionHandler{
-		conn:    conn,
-		ctx:     ctx.Context(),
-		log:     l,
-		options: opts,
+		conn:           conn,
+		requestContext: requestContext,
+		engineContext:  engineContext,
+		log:            l,
+		updater:        updater,
+		options:        options,
 	}
 }
 
-func (h *gqlSSEConnectionHandler) StartBlocking(sub Subscription) {
-	reqCtx := sub.ctx
-
+func (h *gqlSSEConnectionHandler) StartBlocking() {
 	dataCh := make(chan []byte)
 	errCh := make(chan []byte)
 	defer func() {
 		close(dataCh)
 		close(errCh)
-		sub.updater.Done()
+		h.updater.Done()
 	}()
 
-	go h.subscribe(reqCtx, sub, dataCh, errCh)
+	go h.subscribe(dataCh, errCh)
 
 	ticker := time.NewTicker(resolve.HearbeatInterval)
 	defer ticker.Stop()
 
 	for {
 		select {
 		case <-ticker.C:
-			sub.updater.Heartbeat()
+			h.updater.Heartbeat()
 		case data := <-dataCh:
 			ticker.Reset(resolve.HearbeatInterval)
-			sub.updater.Update(data)
+			h.updater.Update(data)
 		case data := <-errCh:
 			ticker.Reset(resolve.HearbeatInterval)
-			sub.updater.Update(data)
+			h.updater.Update(data)
+			return
+		case <-h.requestContext.Done():
 			return
-		case <-reqCtx.Done():
+		case <-h.engineContext.Done():
 			return
 		}
 	}
 }
 
-func (h *gqlSSEConnectionHandler) subscribe(ctx context.Context, sub Subscription, dataCh, errCh chan []byte) {
-	resp, err := h.performSubscriptionRequest(ctx)
+func (h *gqlSSEConnectionHandler) subscribe(dataCh, errCh chan []byte) {
+	resp, err := h.performSubscriptionRequest()
 	if err != nil {
 		h.log.Error("failed to perform subscription request", log.Error(err))
 
-		if ctx.Err() != nil {
+		if h.requestContext.Err() != nil {
 			// request context was canceled do not send an error as channel will be closed
 			return
 		}
 
-		sub.updater.Update([]byte(internalError))
+		h.updater.Update([]byte(internalError))
 
 		return
 	}
@@ -94,8 +97,12 @@ func (h *gqlSSEConnectionHandler) subscribe(ctx context.Context, sub Subscriptio
 	reader := sse.NewEventStreamReader(resp.Body, math.MaxInt)
 
 	for {
-		if ctx.Err() != nil {
+		select {
+		case <-h.requestContext.Done():
+			return
+		case <-h.engineContext.Done():
 			return
+		default:
 		}
 
 		msg, err := reader.ReadEvent()
@@ -126,7 +133,7 @@ func (h *gqlSSEConnectionHandler) subscribe(ctx context.Context, sub Subscriptio
 					continue
 				}
 
-				if ctx.Err() != nil {
+				if h.requestContext.Err() != nil {
 					// request context was canceled do not send an error as channel will be closed
 					return
 				}
@@ -205,16 +212,16 @@ func trim(data []byte) []byte {
 	return data
 }
 
-func (h *gqlSSEConnectionHandler) performSubscriptionRequest(ctx context.Context) (*http.Response, error) {
+func (h *gqlSSEConnectionHandler) performSubscriptionRequest() (*http.Response, error) {
 
 	var req *http.Request
 	var err error
 
 	// default to GET requests when SSEMethodPost is not enabled in the SubscriptionConfiguration
 	if h.options.SSEMethodPost {
-		req, err = h.buildPOSTRequest(ctx)
+		req, err = h.buildPOSTRequest(h.requestContext)
 	} else {
-		req, err = h.buildGETRequest(ctx)
+		req, err = h.buildGETRequest(h.requestContext)
 	}
 
 	if err != nil {