Query Frontend Refactor: Fix goroutine leaks (#3495)

* drain channels

Signed-off-by: Joe Elliott <[email protected]>

* remove duplicate test

Signed-off-by: Joe Elliott <[email protected]>

* lint

Signed-off-by: Joe Elliott <[email protected]>

* fix

Signed-off-by: Joe Elliott <[email protected]>

* remove handlers test

Signed-off-by: Joe Elliott <[email protected]>

* simplify + tests pass

Signed-off-by: Joe Elliott <[email protected]>

* lint

Signed-off-by: Joe Elliott <[email protected]>

* drain async for proper cancellation

Signed-off-by: Joe Elliott <[email protected]>

* Add context to send

Signed-off-by: Joe Elliott <[email protected]>

---------

Signed-off-by: Joe Elliott <[email protected]>

modules/frontend/pipeline/async_handler_multitenant.go

@@ -50,7 +50,7 @@ func (t *tenantRoundTripper) RoundTrip(req *http.Request) (Responses[*http.Respo
 	// join tenants for logger because list value type is unsupported.
 	_ = level.Debug(t.logger).Log("msg", "handling multi-tenant query", "tenants", strings.Join(tenants, ","))
 
-	return NewAsyncSharderFunc(0, len(tenants), func(tenantIdx int) *http.Request {
+	return NewAsyncSharderFunc(req.Context(), 0, len(tenants), func(tenantIdx int) *http.Request {
 		if tenantIdx >= len(tenants) {
 			return nil
 		}

modules/frontend/pipeline/async_sharding.go

@@ -1,6 +1,7 @@
 package pipeline
 
 import (
+	"context"
 	"net/http"
 	"sync"
 
@@ -15,7 +16,7 @@ type waitGroup interface {
 
 // NewAsyncSharderFunc creates a new AsyncResponse that shards requests to the next AsyncRoundTripper[*http.Response]. It creates one
 // goroutine per concurrent request.
-func NewAsyncSharderFunc(concurrentReqs, totalReqs int, reqFn func(i int) *http.Request, next AsyncRoundTripper[*http.Response]) Responses[*http.Response] {
+func NewAsyncSharderFunc(ctx context.Context, concurrentReqs, totalReqs int, reqFn func(i int) *http.Request, next AsyncRoundTripper[*http.Response]) Responses[*http.Response] {
 	var wg waitGroup
 	if concurrentReqs <= 0 {
 		wg = &sync.WaitGroup{}
@@ -26,7 +27,7 @@ func NewAsyncSharderFunc(concurrentReqs, totalReqs int, reqFn func(i int) *http.
 	asyncResp := newAsyncResponse()
 
 	go func() {
-		defer asyncResp.done()
+		defer asyncResp.SendComplete()
 
 		for i := 0; i < totalReqs; i++ {
 			req := reqFn(i)
@@ -50,7 +51,7 @@ func NewAsyncSharderFunc(concurrentReqs, totalReqs int, reqFn func(i int) *http.
 					return
 				}
 
-				asyncResp.Send(resp)
+				asyncResp.Send(ctx, resp)
 			}(req)
 		}
 
@@ -61,7 +62,7 @@ func NewAsyncSharderFunc(concurrentReqs, totalReqs int, reqFn func(i int) *http.
 }
 
 // NewAsyncSharderChan creates a new AsyncResponse that shards requests to the next AsyncRoundTripper[*http.Response] using a limited number of goroutines.
-func NewAsyncSharderChan(concurrentReqs int, reqs <-chan *http.Request, resps Responses[*http.Response], next AsyncRoundTripper[*http.Response]) Responses[*http.Response] {
+func NewAsyncSharderChan(ctx context.Context, concurrentReqs int, reqs <-chan *http.Request, resps Responses[*http.Response], next AsyncRoundTripper[*http.Response]) Responses[*http.Response] {
 	if concurrentReqs == 0 {
 		panic("NewAsyncSharderChan: concurrentReqs must be greater than 0")
 	}
@@ -85,20 +86,19 @@ func NewAsyncSharderChan(concurrentReqs int, reqs <-chan *http.Request, resps Re
 					continue
 				}
 
-				asyncResp.Send(resp)
+				asyncResp.Send(ctx, resp)
 			}
 		}()
 	}
 
 	go func() {
 		// send any responses back the caller would like to send
 		if resps != nil {
-			asyncResp.Send(resps)
+			asyncResp.Send(ctx, resps)
 		}
 
-		// and wait for all the workers to finish
 		wg.Wait()
-		asyncResp.done()
+		asyncResp.SendComplete()
 	}()
 
 	return asyncResp

modules/frontend/pipeline/async_sharding_test.go

@@ -21,7 +21,7 @@ func TestAsyncSharders(t *testing.T) {
 		{
 			name: "AsyncSharder",
 			responseFn: func(next AsyncRoundTripper[*http.Response]) *asyncResponse {
-				return NewAsyncSharderFunc(10, expectedRequestCount, func(i int) *http.Request {
+				return NewAsyncSharderFunc(context.Background(), 10, expectedRequestCount, func(i int) *http.Request {
 					if i >= expectedRequestCount {
 						return nil
 					}
@@ -32,7 +32,7 @@ func TestAsyncSharders(t *testing.T) {
 		{
 			name: "AsyncSharder - no limit",
 			responseFn: func(next AsyncRoundTripper[*http.Response]) *asyncResponse {
-				return NewAsyncSharderFunc(0, expectedRequestCount, func(i int) *http.Request {
+				return NewAsyncSharderFunc(context.Background(), 0, expectedRequestCount, func(i int) *http.Request {
 					if i >= expectedRequestCount {
 						return nil
 					}
@@ -51,7 +51,7 @@ func TestAsyncSharders(t *testing.T) {
 					close(reqChan)
 				}()
 
-				return NewAsyncSharderChan(10, reqChan, nil, next).(*asyncResponse)
+				return NewAsyncSharderChan(context.Background(), 10, reqChan, nil, next).(*asyncResponse)
 			},
 		},
 	}

modules/frontend/pipeline/collector_grpc.go

@@ -99,6 +99,7 @@ func (c GRPCCollector[T]) RoundTrip(req *http.Request) error {
 	if err != nil {
 		return grpcError(err)
 	}
+
 	return nil
 }
 

modules/frontend/pipeline/collector_http.go

@@ -58,5 +58,6 @@ func (r httpCollector) RoundTrip(req *http.Request) (*http.Response, error) {
 	}
 
 	resp, err := r.combiner.HTTPFinal()
+
 	return resp, err
 }

modules/frontend/pipeline/responses.go

@@ -10,6 +10,7 @@ import (
 )
 
 type Responses[T any] interface {
+	// Next returns the next response or an error if one is available. It always prefers an error over a response.
 	Next(context.Context) (T, bool, error) // bool = done
 }
 
@@ -72,8 +73,11 @@ func newAsyncResponse() *asyncResponse {
 	}
 }
 
-func (a *asyncResponse) Send(r Responses[*http.Response]) {
-	a.respChan <- r
+func (a *asyncResponse) Send(ctx context.Context, r Responses[*http.Response]) {
+	select {
+	case <-ctx.Done():
+	case a.respChan <- r:
+	}
 }
 
 // SendError sends an error to the asyncResponse. This will cause the asyncResponse to return the error on the next call to Next.
@@ -87,7 +91,8 @@ func (a *asyncResponse) SendError(err error) {
 	}
 }
 
-func (a *asyncResponse) done() {
+// SendComplete indicates the sender is done. We close the channel to give a clear signal to the consumer
+func (a *asyncResponse) SendComplete() {
 	close(a.respChan)
 }