Fix coroutine leak (thanos-io#7821)

* Fix coroutine leak The in-process client uses a pull based iterator which needs to be closed, otherwise it will leak the underlying coroutine. When this happens, the tsdb reader will remain open which blocks head compaction indefinitely. Signed-off-by: Filip Petkovski <[email protected]> * Fix race condition Signed-off-by: Filip Petkovski <[email protected]> * Fix CHANGELOG Signed-off-by: Filip Petkovski <[email protected]> * Improve tests Signed-off-by: Filip Petkovski <[email protected]> * Fix blockSeriesClient Signed-off-by: Filip Petkovski <[email protected]> * Fix unit test Signed-off-by: Filip Petkovski <[email protected]> * Fix another unit test Signed-off-by: Filip Petkovski <[email protected]> --------- Signed-off-by: Filip Petkovski <[email protected]>
jnyi · Oct 17, 2024 · 5419ead · 5419ead
1 parent 53ffb04
commit 5419ead
Show file tree

Hide file tree

Showing 6 changed files with 163 additions and 68 deletions.
diff --git a/pkg/store/bucket.go b/pkg/store/bucket.go
@@ -1584,8 +1584,6 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store
 				tenant,
 			)
 
-			defer blockClient.Close()
-
 			g.Go(func() error {
 
 				span, _ := tracing.StartSpan(gctx, "bucket_store_block_series", tracing.Tags{
@@ -3366,6 +3364,11 @@ func (r *bucketIndexReader) Close() error {
 	return nil
 }
 
+func (b *blockSeriesClient) CloseSend() error {
+	b.Close()
+	return nil
+}
+
 // LookupLabelsSymbols allows populates label set strings from symbolized label set.
 func (r *bucketIndexReader) LookupLabelsSymbols(ctx context.Context, symbolized []symbolizedLabel, b *labels.Builder) error {
 	b.Reset(labels.EmptyLabels())

diff --git a/pkg/store/proxy_merge.go b/pkg/store/proxy_merge.go
@@ -238,6 +238,7 @@ func (l *lazyRespSet) StoreLabels() map[string]struct{} {
 type lazyRespSet struct {
 	// Generic parameters.
 	span           opentracing.Span
+	cl             storepb.Store_SeriesClient
 	closeSeries    context.CancelFunc
 	storeName      string
 	storeLabelSets []labels.Labels
@@ -337,6 +338,7 @@ func newLazyRespSet(
 		frameTimeout:         frameTimeout,
 		storeName:            storeName,
 		storeLabelSets:       storeLabelSets,
+		cl:                   cl,
 		closeSeries:          closeSeries,
 		span:                 span,
 		dataOrFinishEvent:    dataAvailable,
@@ -468,8 +470,10 @@ func newAsyncRespSet(
 	emptyStreamResponses prometheus.Counter,
 ) (respSet, error) {
 
-	var span opentracing.Span
-	var closeSeries context.CancelFunc
+	var (
+		span   opentracing.Span
+		cancel context.CancelFunc
+	)
 
 	storeID, storeAddr, isLocalStore := storeInfo(st)
 	seriesCtx := grpc_opentracing.ClientAddContextTags(ctx, opentracing.Tags{
@@ -481,7 +485,7 @@ func newAsyncRespSet(
 		"store.addr":     storeAddr,
 	})
 
-	seriesCtx, closeSeries = context.WithCancel(seriesCtx)
+	seriesCtx, cancel = context.WithCancel(seriesCtx)
 
 	shardMatcher := shardInfo.Matcher(buffers)
 
@@ -496,7 +500,7 @@ func newAsyncRespSet(
 
 		span.SetTag("err", err.Error())
 		span.Finish()
-		closeSeries()
+		cancel()
 		return nil, err
 	}
 
@@ -520,7 +524,7 @@ func newAsyncRespSet(
 			frameTimeout,
 			st.String(),
 			st.LabelSets(),
-			closeSeries,
+			cancel,
 			cl,
 			shardMatcher,
 			applySharding,
@@ -533,7 +537,7 @@ func newAsyncRespSet(
 			frameTimeout,
 			st.String(),
 			st.LabelSets(),
-			closeSeries,
+			cancel,
 			cl,
 			shardMatcher,
 			applySharding,
@@ -554,6 +558,7 @@ func (l *lazyRespSet) Close() {
 	l.dataOrFinishEvent.Signal()
 
 	l.shardMatcher.Close()
+	_ = l.cl.CloseSend()
 }
 
 // eagerRespSet is a SeriesSet that blocks until all data is retrieved from
@@ -563,6 +568,7 @@ type eagerRespSet struct {
 	// Generic parameters.
 	span opentracing.Span
 
+	cl           storepb.Store_SeriesClient
 	closeSeries  context.CancelFunc
 	frameTimeout time.Duration
 
@@ -593,6 +599,7 @@ func newEagerRespSet(
 ) respSet {
 	ret := &eagerRespSet{
 		span:              span,
+		cl:                cl,
 		closeSeries:       closeSeries,
 		frameTimeout:      frameTimeout,
 		bufferedResponses: []*storepb.SeriesResponse{},
@@ -744,6 +751,7 @@ func (l *eagerRespSet) Close() {
 		l.closeSeries()
 	}
 	l.shardMatcher.Close()
+	_ = l.cl.CloseSend()
 }
 
 func (l *eagerRespSet) At() *storepb.SeriesResponse {

diff --git a/pkg/store/proxy_merge_test.go b/pkg/store/proxy_merge_test.go
@@ -35,6 +35,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 			input: []respSet{
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
@@ -43,6 +44,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 				},
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "4", "e", "5")),
@@ -62,6 +64,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 			input: []respSet{
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
@@ -71,6 +74,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 				},
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("d", "4", "e", "5")),
@@ -91,6 +95,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 			input: []respSet{
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "ext2", "9")),
@@ -100,6 +105,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 				},
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "ext1", "5", "ext2", "9")),
@@ -120,6 +126,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 			input: []respSet{
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
@@ -129,6 +136,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 				},
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "1", "c", "3")),
@@ -149,6 +157,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 			input: []respSet{
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "ext2", "9")),
@@ -158,6 +167,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 				},
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "ext1", "5", "ext2", "9")),
@@ -178,6 +188,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 			input: []respSet{
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("cluster", "beam-platform", "instance", "10.70.13.3:15692", "prometheus", "telemetry/observe-prometheus", "receive", "true", "tenant_id", "default-tenant")),
@@ -188,6 +199,7 @@ func TestProxyResponseTreeSort(t *testing.T) {
 				},
 				&eagerRespSet{
 					closeSeries: func() {},
+					cl:          nopClientSendCloser{},
 					wg:          &sync.WaitGroup{},
 					bufferedResponses: []*storepb.SeriesResponse{
 						storeSeriesResponse(t, labelsFromStrings("cluster", "beam-platform", "instance", "10.70.13.3:15692", "prometheus", "telemetry/observe-prometheus", "receive", "true", "tenant_id", "default-tenant")),
@@ -219,6 +231,12 @@ func TestProxyResponseTreeSort(t *testing.T) {
 	}
 }
 
+type nopClientSendCloser struct {
+	storepb.Store_SeriesClient
+}
+
+func (c nopClientSendCloser) CloseSend() error { return nil }
+
 func TestSortWithoutLabels(t *testing.T) {
 	for _, tcase := range []struct {
 		input       []*storepb.SeriesResponse
@@ -341,6 +359,7 @@ func BenchmarkKWayMerge(b *testing.B) {
 		for j := 0; j < 1000; j++ {
 			respSets = append(respSets, &eagerRespSet{
 				closeSeries: func() {},
+				cl:          nopClientSendCloser{},
 				wg:          &sync.WaitGroup{},
 				bufferedResponses: []*storepb.SeriesResponse{
 					storeSeriesResponse(b, labelsFromStrings("a", "1", "b", fmt.Sprintf("replica-%d", j), "c", "3")),