Propagate the stats via gRPC

Only partially done, still need to merge & record the results in the query frontend.

Signed-off-by: Tom Wilkie <[email protected]>

Author: tomwilkie

pkg/frontend/v1/frontendv1pb/frontend.pb.go

@@ -8,6 +8,7 @@ option go_package = "frontendv1pb";
 
 import "github.com/gogo/protobuf/gogoproto/gogo.proto";
 import "github.com/weaveworks/common/httpgrpc/httpgrpc.proto";
+import "github.com/cortexproject/cortex/pkg/querier/stats/stats.proto";
 
 option (gogoproto.marshaler_all) = true;
 option (gogoproto.unmarshaler_all) = true;
@@ -31,4 +32,5 @@ message FrontendToClient {
 message ClientToFrontend {
   httpgrpc.HTTPResponse httpResponse = 1;
   string clientID = 2;
+  stats.Stats stats = 3;
 }

pkg/frontend/v1/frontendv1pb/frontend.proto

@@ -20,6 +20,7 @@ import (
 	"go.uber.org/atomic"
 
 	"github.com/cortexproject/cortex/pkg/frontend/v2/frontendv2pb"
+	"github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	"github.com/cortexproject/cortex/pkg/util/grpcutil"
@@ -79,7 +80,7 @@ type frontendRequest struct {
 	cancel context.CancelFunc
 
 	enqueue  chan enqueueResult
-	response chan *httpgrpc.HTTPResponse
+	response chan *frontendv2pb.QueryResultRequest
 }
 
 type enqueueStatus int
@@ -178,7 +179,7 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest)
 		// Buffer of 1 to ensure response or error can be written to the channel
 		// even if this goroutine goes away due to client context cancellation.
 		enqueue:  make(chan enqueueResult, 1),
-		response: make(chan *httpgrpc.HTTPResponse, 1),
+		response: make(chan *frontendv2pb.QueryResultRequest, 1),
 	}
 
 	f.requests.put(freq)
@@ -228,7 +229,9 @@ enqueueAgain:
 		return nil, ctx.Err()
 
 	case resp := <-freq.response:
-		return resp, nil
+		stats := stats.FromContext(ctx)
+		stats.Merge(resp.Stats)
+		return resp.HttpResponse, nil
 	}
 }
 
@@ -244,7 +247,7 @@ func (f *Frontend) QueryResult(ctx context.Context, qrReq *frontendv2pb.QueryRes
 	// To avoid mixing results from different queries, we randomize queryID counter on start.
 	if req != nil && req.userID == userID {
 		select {
-		case req.response <- qrReq.HttpResponse:
+		case req.response <- qrReq:
 			// Should always be possible, unless QueryResult is called multiple times with the same queryID.
 		default:
 			level.Warn(f.log).Log("msg", "failed to write query result to the response channel", "queryID", qrReq.QueryID, "user", userID)

pkg/frontend/v2/frontend.go

@@ -12,6 +12,7 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"google.golang.org/grpc"
 
+	"github.com/cortexproject/cortex/pkg/frontend/v2/frontendv2pb"
 	"github.com/cortexproject/cortex/pkg/scheduler/schedulerpb"
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/services"
@@ -285,16 +286,20 @@ func (w *frontendSchedulerWorker) schedulerLoop(loop schedulerpb.SchedulerForFro
 
 			case schedulerpb.ERROR:
 				req.enqueue <- enqueueResult{status: waitForResponse}
-				req.response <- &httpgrpc.HTTPResponse{
-					Code: http.StatusInternalServerError,
-					Body: []byte(err.Error()),
+				req.response <- &frontendv2pb.QueryResultRequest{
+					HttpResponse: &httpgrpc.HTTPResponse{
+						Code: http.StatusInternalServerError,
+						Body: []byte(err.Error()),
+					},
 				}
 
 			case schedulerpb.TOO_MANY_REQUESTS_PER_TENANT:
 				req.enqueue <- enqueueResult{status: waitForResponse}
-				req.response <- &httpgrpc.HTTPResponse{
-					Code: http.StatusTooManyRequests,
-					Body: []byte("too many outstanding requests"),
+				req.response <- &frontendv2pb.QueryResultRequest{
+					HttpResponse: &httpgrpc.HTTPResponse{
+						Code: http.StatusTooManyRequests,
+						Body: []byte("too many outstanding requests"),
+					},
 				}
 			}
 

pkg/frontend/v2/frontend_scheduler_worker.go

@@ -17,6 +17,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/cortexproject/cortex/pkg/frontend/v2/frontendv2pb"
+	"github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/scheduler/schedulerpb"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/services"
@@ -87,6 +88,7 @@ func sendResponseWithDelay(f *Frontend, delay time.Duration, userID string, quer
 	_, _ = f.QueryResult(ctx, &frontendv2pb.QueryResultRequest{
 		QueryID:      queryID,
 		HttpResponse: resp,
+		Stats:        &stats.Stats{},
 	})
 }