fix: delete leftover heartbeat connections

v2/pkg/engine/resolve/resolve.go

@@ -51,6 +51,7 @@ type Resolver struct {
 	maxConcurrency chan struct{}
 
 	triggers               map[uint64]*trigger
+	heartbeatSubLock       *sync.Mutex
 	heartbeatSubscriptions map[*Context]*sub
 	events                 chan subscriptionEvent
 	triggerEventsSem       *semaphore.Weighted
@@ -189,6 +190,7 @@ func New(ctx context.Context, options ResolverOptions) *Resolver {
 		propagateSubgraphStatusCodes:  options.PropagateSubgraphStatusCodes,
 		events:                        make(chan subscriptionEvent),
 		triggers:                      make(map[uint64]*trigger),
+		heartbeatSubLock:              &sync.Mutex{},
 		heartbeatSubscriptions:        make(map[*Context]*sub),
 		reporter:                      options.Reporter,
 		asyncErrorWriter:              options.AsyncErrorWriter,
@@ -407,6 +409,9 @@ func (r *Resolver) handleEvent(event subscriptionEvent) {
 }
 
 func (r *Resolver) handleHeartbeat(data []byte) {
+	r.heartbeatSubLock.Lock()
+	defer r.heartbeatSubLock.Unlock()
+
 	if r.options.Debug {
 		fmt.Printf("resolver:heartbeat:%d\n", len(r.heartbeatSubscriptions))
 	}
@@ -417,7 +422,7 @@ func (r *Resolver) handleHeartbeat(data []byte) {
 		s.mux.Lock()
 		skipHeartbeat := now.Sub(s.lastWrite) < r.multipartSubHeartbeatInterval
 		s.mux.Unlock()
-		if skipHeartbeat {
+		if skipHeartbeat || (c.Context().Err() != nil && errors.Is(c.Context().Err(), context.Canceled)) {
 			continue
 		}
 
@@ -428,6 +433,12 @@ func (r *Resolver) handleHeartbeat(data []byte) {
 
 			s.mux.Lock()
 			if _, err := s.writer.Write(data); err != nil {
+				if errors.Is(err, context.Canceled) {
+					// client disconnected
+					s.mux.Unlock()
+					_ = r.AsyncUnsubscribeSubscription(s.id)
+					return
+				}
 				r.asyncErrorWriter.WriteError(c, err, nil, s.writer)
 			}
 			err := s.writer.Flush()
@@ -468,30 +479,7 @@ func (r *Resolver) handleTriggerInitialized(triggerID uint64) {
 }
 
 func (r *Resolver) handleTriggerDone(triggerID uint64) {
-	trig, ok := r.triggers[triggerID]
-	if !ok {
-		return
-	}
-	isInitialized := trig.initialized
-	wg := trig.inFlight
-	subscriptionCount := len(trig.subscriptions)
-
-	delete(r.triggers, triggerID)
-
-	go func() {
-		if wg != nil {
-			wg.Wait()
-		}
-		for _, s := range trig.subscriptions {
-			s.writer.Complete()
-		}
-		if r.reporter != nil {
-			r.reporter.SubscriptionCountDec(subscriptionCount)
-			if isInitialized {
-				r.reporter.TriggerCountDec(1)
-			}
-		}
-	}()
+	r.shutdownTrigger(triggerID)
 }
 
 func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription) {
@@ -510,7 +498,9 @@ func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription)
 		executor:  add.executor,
 	}
 	if add.ctx.ExecutionOptions.SendHeartbeat {
+		r.heartbeatSubLock.Lock()
 		r.heartbeatSubscriptions[add.ctx] = s
+		r.heartbeatSubLock.Unlock()
 	}
 	trig, ok := r.triggers[triggerID]
 	if ok {
@@ -636,20 +626,9 @@ func (r *Resolver) handleRemoveSubscription(id SubscriptionIdentifier) {
 	removed := 0
 	for u := range r.triggers {
 		trig := r.triggers[u]
-		for ctx, s := range trig.subscriptions {
-			if s.id == id {
-
-				if ctx.Context().Err() == nil {
-					s.writer.Complete()
-				}
-				delete(r.heartbeatSubscriptions, ctx)
-				delete(trig.subscriptions, ctx)
-				if r.options.Debug {
-					fmt.Printf("resolver:trigger:subscription:removed:%d:%d\n", trig.id, id.SubscriptionID)
-				}
-				removed++
-			}
-		}
+		removed += r.shutdownTriggerSubscriptions(u, func(sID SubscriptionIdentifier) bool {
+			return sID == id
+		})
 		if len(trig.subscriptions) == 0 {
 			r.shutdownTrigger(trig.id)
 		}
@@ -665,20 +644,9 @@ func (r *Resolver) handleRemoveClient(id int64) {
 	}
 	removed := 0
 	for u := range r.triggers {
-		for c, s := range r.triggers[u].subscriptions {
-			if s.id.ConnectionID == id && !s.id.internal {
-
-				if c.Context().Err() == nil {
-					s.writer.Complete()
-				}
-
-				delete(r.triggers[u].subscriptions, c)
-				if r.options.Debug {
-					fmt.Printf("resolver:trigger:subscription:done:%d:%d\n", u, s.id.SubscriptionID)
-				}
-				removed++
-			}
-		}
+		removed += r.shutdownTriggerSubscriptions(u, func(sID SubscriptionIdentifier) bool {
+			return sID.ConnectionID == id && !sID.internal
+		})
 		if len(r.triggers[u].subscriptions) == 0 {
 			r.shutdownTrigger(r.triggers[u].id)
 		}
@@ -739,30 +707,46 @@ func (r *Resolver) shutdownTrigger(id uint64) {
 		return
 	}
 	count := len(trig.subscriptions)
+	r.shutdownTriggerSubscriptions(id, nil)
+	trig.cancel()
+	delete(r.triggers, id)
+	if r.options.Debug {
+		fmt.Printf("resolver:trigger:done:%d\n", trig.id)
+	}
+	if r.reporter != nil {
+		r.reporter.SubscriptionCountDec(count)
+		if trig.initialized {
+			r.reporter.TriggerCountDec(1)
+		}
+	}
+}
+
+func (r *Resolver) shutdownTriggerSubscriptions(id uint64, shutdownMatcher func(a SubscriptionIdentifier) bool) int {
+	trig, ok := r.triggers[id]
+	if !ok {
+		return 0
+	}
+	removed := 0
 	for c, s := range trig.subscriptions {
+		if shutdownMatcher != nil && !shutdownMatcher(s.id) {
+			continue
+		}
 		if c.Context().Err() == nil {
 			s.writer.Complete()
 		}
 		if s.completed != nil {
 			close(s.completed)
 		}
+		r.heartbeatSubLock.Lock()
 		delete(r.heartbeatSubscriptions, c)
+		r.heartbeatSubLock.Unlock()
 		delete(trig.subscriptions, c)
 		if r.options.Debug {
 			fmt.Printf("resolver:trigger:subscription:done:%d:%d\n", trig.id, s.id.SubscriptionID)
 		}
+		removed++
 	}
-	trig.cancel()
-	delete(r.triggers, id)
-	if r.options.Debug {
-		fmt.Printf("resolver:trigger:done:%d\n", trig.id)
-	}
-	if r.reporter != nil {
-		r.reporter.SubscriptionCountDec(count)
-		if trig.initialized {
-			r.reporter.TriggerCountDec(1)
-		}
-	}
+	return removed
 }
 
 func (r *Resolver) handleShutdown() {

v2/pkg/engine/resolve/resolve_test.go

@@ -86,13 +86,16 @@ func (t *TestErrorWriter) WriteError(ctx *Context, err error, res *GraphQLRespon
 	}
 }
 
+var multipartSubHeartbeatInterval = 15 * time.Millisecond
+
 func newResolver(ctx context.Context) *Resolver {
 	return New(ctx, ResolverOptions{
-		MaxConcurrency:               1024,
-		Debug:                        false,
-		PropagateSubgraphErrors:      true,
-		PropagateSubgraphStatusCodes: true,
-		AsyncErrorWriter:             &TestErrorWriter{},
+		MaxConcurrency:                1024,
+		Debug:                         false,
+		PropagateSubgraphErrors:       true,
+		PropagateSubgraphStatusCodes:  true,
+		AsyncErrorWriter:              &TestErrorWriter{},
+		MultipartSubHeartbeatInterval: multipartSubHeartbeatInterval,
 	})
 }
 
@@ -5164,19 +5167,67 @@ func TestResolver_ResolveGraphQLSubscription(t *testing.T) {
 
 		ctx := &Context{
 			ctx: context.Background(),
+			ExecutionOptions: ExecutionOptions{
+				SendHeartbeat: true,
+			},
 		}
 
 		err := resolver.AsyncResolveGraphQLSubscription(ctx, plan, recorder, id)
 		assert.NoError(t, err)
+
 		recorder.AwaitComplete(t, defaultTimeout)
 		assert.Equal(t, 3, len(recorder.Messages()))
+		time.Sleep(2 * resolver.multipartSubHeartbeatInterval)
+		// Validate that despite the time, we don't see any heartbeats sent
 		assert.ElementsMatch(t, []string{
 			`{"data":{"counter":0}}`,
 			`{"data":{"counter":1}}`,
 			`{"data":{"counter":2}}`,
 		}, recorder.Messages())
 	})
 
+	t.Run("should successfully delete multiple finished subscriptions", func(t *testing.T) {
+		c, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		fakeStream := createFakeStream(func(counter int) (message string, done bool) {
+			return fmt.Sprintf(`{"data":{"counter":%d}}`, counter), counter == 1
+		}, 1*time.Millisecond, func(input []byte) {
+			assert.Equal(t, `{"method":"POST","url":"http://localhost:4000","body":{"query":"subscription { counter }"}}`, string(input))
+		})
+
+		resolver, plan, recorder, id := setup(c, fakeStream)
+
+		ctx := &Context{
+			ctx: context.Background(),
+			ExecutionOptions: ExecutionOptions{
+				SendHeartbeat: true,
+			},
+		}
+
+		for i := 1; i <= 10; i++ {
+			id.ConnectionID = int64(i)
+			id.SubscriptionID = int64(i)
+			recorder.complete.Store(false)
+			err := resolver.AsyncResolveGraphQLSubscription(ctx, plan, recorder, id)
+			assert.NoError(t, err)
+			recorder.AwaitComplete(t, defaultTimeout)
+		}
+
+		recorder.AwaitComplete(t, defaultTimeout)
+		time.Sleep(2 * resolver.multipartSubHeartbeatInterval)
+
+		assert.Equal(t, 20, len(recorder.Messages()))
+		// Validate that despite the time, we don't see any heartbeats sent
+		assert.ElementsMatch(t, []string{
+			`{"data":{"counter":0}}`, `{"data":{"counter":1}}`, `{"data":{"counter":0}}`, `{"data":{"counter":1}}`,
+			`{"data":{"counter":0}}`, `{"data":{"counter":1}}`, `{"data":{"counter":0}}`, `{"data":{"counter":1}}`,
+			`{"data":{"counter":0}}`, `{"data":{"counter":1}}`, `{"data":{"counter":0}}`, `{"data":{"counter":1}}`,
+			`{"data":{"counter":0}}`, `{"data":{"counter":1}}`, `{"data":{"counter":0}}`, `{"data":{"counter":1}}`,
+			`{"data":{"counter":0}}`, `{"data":{"counter":1}}`, `{"data":{"counter":0}}`, `{"data":{"counter":1}}`,
+		}, recorder.Messages())
+	})
+
 	t.Run("should propagate extensions to stream", func(t *testing.T) {
 		c, cancel := context.WithCancel(context.Background())
 		defer cancel()