Configure gRPC clients for query-frontend->query-scheduler communicat…

…ions Signed-off-by: Yuri Nikolic <[email protected]>
grafana · Mar 6, 2025 · ed4205b · ed4205b
1 parent 9bddb85
commit ed4205b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/pkg/frontend/v2/frontend_scheduler_worker.go b/pkg/frontend/v2/frontend_scheduler_worker.go
@@ -16,6 +16,7 @@ import (
 	"github.com/grafana/dskit/backoff"
 	"github.com/grafana/dskit/cancellation"
 	"github.com/grafana/dskit/httpgrpc"
+	"github.com/grafana/dskit/middleware"
 	"github.com/grafana/dskit/servicediscovery"
 	"github.com/grafana/dskit/services"
 	"github.com/pkg/errors"
@@ -58,7 +59,8 @@ type frontendSchedulerWorkers struct {
 	// Set to nil when stop is called... no more workers are created afterwards.
 	workers map[string]*frontendSchedulerWorker
 
-	enqueueDuration *prometheus.HistogramVec
+	enqueueDuration          *prometheus.HistogramVec
+	invalidClusterValidation *prometheus.CounterVec
 }
 
 func newFrontendSchedulerWorkers(
@@ -84,6 +86,7 @@ func newFrontendSchedulerWorkers(
 			// track 1ms latency too and removing any bucket bigger than 1s.
 			Buckets: []float64{.001, .005, .01, .025, .05, .1, .25, .5, 1},
 		}, []string{schedulerAddressLabel}),
+		invalidClusterValidation: util.NewRequestInvalidClusterValidationLabelsTotalCounter(reg, "query-scheduler", util.GRPCProtocol),
 	}
 
 	var err error
@@ -219,7 +222,11 @@ func (f *frontendSchedulerWorkers) getWorkersCount() int {
 
 func (f *frontendSchedulerWorkers) connectToScheduler(ctx context.Context, address string) (*grpc.ClientConn, error) {
 	// Because we only use single long-running method, it doesn't make sense to inject user ID, send over tracing or add metrics.
-	opts, err := f.cfg.GRPCClientConfig.DialOption(nil, nil)
+	var unary []grpc.UnaryClientInterceptor
+	if f.cfg.GRPCClientConfig.ClusterValidation.Label != "" {
+		unary = []grpc.UnaryClientInterceptor{middleware.ClusterUnaryClientInterceptor(f.cfg.GRPCClientConfig.ClusterValidation.Label, f.invalidClusterValidation, f.log)}
+	}
+	opts, err := f.cfg.GRPCClientConfig.DialOption(unary, nil)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/mimir/mimir.go b/pkg/mimir/mimir.go
@@ -229,6 +229,7 @@ func (c *Config) CommonConfigInheritance() CommonConfigInheritance {
 			"frontend_worker_frontend_client":  &c.Worker.QueryFrontendGRPCClientConfig.ClusterValidation,
 			"frontend_worker_scheduler_client": &c.Worker.QuerySchedulerGRPCClientConfig.ClusterValidation,
 			"block_builder_scheduler_client":   &c.BlockBuilder.SchedulerConfig.GRPCClientConfig.ClusterValidation,
+			"frontend_query_scheduler_client":  &c.Frontend.FrontendV2.GRPCClientConfig.ClusterValidation,
 		},
 	}
 }

diff --git a/pkg/mimir/mimir_config_test.go b/pkg/mimir/mimir_config_test.go
@@ -39,6 +39,7 @@ func TestCommonConfigCanBeExtended(t *testing.T) {
 		require.Equal(t, "client-cluster", cfg.MimirConfig.Worker.QueryFrontendGRPCClientConfig.ClusterValidation.Label)
 		require.Equal(t, "client-cluster", cfg.MimirConfig.Worker.QuerySchedulerGRPCClientConfig.ClusterValidation.Label)
 		require.Equal(t, "client-cluster", cfg.MimirConfig.BlockBuilder.SchedulerConfig.GRPCClientConfig.ClusterValidation.Label)
+		require.Equal(t, "client-cluster", cfg.MimirConfig.Frontend.FrontendV2.GRPCClientConfig.ClusterValidation.Label)
 	})
 
 	t.Run("yaml inheritance", func(t *testing.T) {
@@ -67,6 +68,7 @@ common:
 		require.Equal(t, "client-cluster", cfg.MimirConfig.Worker.QueryFrontendGRPCClientConfig.ClusterValidation.Label)
 		require.Equal(t, "client-cluster", cfg.MimirConfig.Worker.QuerySchedulerGRPCClientConfig.ClusterValidation.Label)
 		require.Equal(t, "client-cluster", cfg.MimirConfig.BlockBuilder.SchedulerConfig.GRPCClientConfig.ClusterValidation.Label)
+		require.Equal(t, "client-cluster", cfg.MimirConfig.Frontend.FrontendV2.GRPCClientConfig.ClusterValidation.Label)
 	})
 }