fix coroutine leak

Author: yuchen-db

pkg/store/proxy_merge.go

@@ -558,13 +558,13 @@ func newAsyncRespSet(
 func (l *lazyRespSet) Close() {
 	l.bufferedResponsesMtx.Lock()
 	defer l.bufferedResponsesMtx.Unlock()
+	_ = l.cl.CloseSend()
 
 	l.closeSeries()
 	l.noMoreData = true
 	l.dataOrFinishEvent.Signal()
 
 	l.shardMatcher.Close()
-	_ = l.cl.CloseSend()
 }
 
 // eagerRespSet is a SeriesSet that blocks until all data is retrieved from
@@ -750,11 +750,11 @@ func sortWithoutLabels(set []*storepb.SeriesResponse, labelsToRemove map[string]
 }
 
 func (l *eagerRespSet) Close() {
+	_ = l.cl.CloseSend()
 	if l.closeSeries != nil {
 		l.closeSeries()
 	}
 	l.shardMatcher.Close()
-	_ = l.cl.CloseSend()
 }
 
 func (l *eagerRespSet) At() *storepb.SeriesResponse {

pkg/store/storepb/inprocess.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"io"
 	"iter"
+	"sync"
 
 	"google.golang.org/grpc"
 )
@@ -38,18 +39,22 @@ type inProcessClient struct {
 	ctx  context.Context
 	next func() (*SeriesResponse, error, bool)
 	stop func()
+	mu   sync.Mutex
 }
 
 func newInProcessClient(ctx context.Context, next func() (*SeriesResponse, error, bool), stop func()) *inProcessClient {
 	return &inProcessClient{
 		ctx:  ctx,
 		next: next,
 		stop: stop,
+		mu:   sync.Mutex{},
 	}
 }
 
 func (c *inProcessClient) Recv() (*SeriesResponse, error) {
+	c.mu.Lock()
 	resp, err, ok := c.next()
+	c.mu.Unlock()
 	if err != nil {
 		c.stop()
 		return nil, err
@@ -68,7 +73,9 @@ func (c *inProcessClient) Context() context.Context {
 }
 
 func (c *inProcessClient) CloseSend() error {
+	c.mu.Lock()
 	c.stop()
+	c.mu.Unlock()
 	return nil
 }