Merge pull request #1138 from gouthamve/up-max-size

Make the max grpc receive message size configurable, in Bigtable client

Author: tomwilkie

pkg/chunk/gcp/bigtable_index_client.go

@@ -10,10 +10,12 @@ import (
 	"cloud.google.com/go/bigtable"
 	ot "github.com/opentracing/opentracing-go"
 	otlog "github.com/opentracing/opentracing-go/log"
+	"google.golang.org/api/option"
 
 	"github.com/cortexproject/cortex/pkg/chunk"
 	chunk_util "github.com/cortexproject/cortex/pkg/chunk/util"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	"github.com/pkg/errors"
 )
 
@@ -30,12 +32,21 @@ const (
 type Config struct {
 	Project  string `yaml:"project"`
 	Instance string `yaml:"instance"`
+
+	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config"`
+
+	ColumnKey bool
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Project, "bigtable.project", "", "Bigtable project ID.")
 	f.StringVar(&cfg.Instance, "bigtable.instance", "", "Bigtable instance ID.")
+
+	cfg.GRPCClientConfig.RegisterFlags("bigtable", f)
+
+	// Deprecated.
+	f.Int("bigtable.max-recv-msg-size", 100<<20, "DEPRECATED. Bigtable grpc max receive message size.")
 }
 
 // storageClientColumnKey implements chunk.storageClient for GCP.
@@ -53,7 +64,10 @@ type storageClientV1 struct {
 
 // NewStorageClientV1 returns a new v1 StorageClient.
 func NewStorageClientV1(ctx context.Context, cfg Config, schemaCfg chunk.SchemaConfig) (chunk.IndexClient, error) {
-	client, err := bigtable.NewClient(ctx, cfg.Project, cfg.Instance, instrumentation()...)
+	opts := instrumentation()
+	opts = append(opts, option.WithGRPCDialOption(cfg.GRPCClientConfig.DialOption()))
+
+	client, err := bigtable.NewClient(ctx, cfg.Project, cfg.Instance, opts...)
 	if err != nil {
 		return nil, err
 	}

pkg/ingester/client/client.go

@@ -12,6 +12,7 @@ import (
 	_ "google.golang.org/grpc/encoding/gzip" // get gzip compressor registered
 	"google.golang.org/grpc/health/grpc_health_v1"
 
+	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	cortex_middleware "github.com/cortexproject/cortex/pkg/util/middleware"
 	"github.com/weaveworks/common/middleware"
 )
@@ -49,13 +50,7 @@ func MakeIngesterClient(addr string, cfg Config) (HealthAndIngesterClient, error
 			middleware.StreamClientUserHeaderInterceptor,
 			cortex_middleware.PrometheusGRPCStreamInstrumentation(ingesterClientRequestDuration),
 		)),
-		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(cfg.MaxRecvMsgSize)),
-	}
-	if cfg.legacyCompressToIngester {
-		cfg.CompressToIngester = true
-	}
-	if cfg.CompressToIngester {
-		opts = append(opts, grpc.WithDefaultCallOptions(grpc.UseCompressor("gzip")))
+		cfg.GRPCClientConfig.DialOption(),
 	}
 	conn, err := grpc.Dial(addr, opts...)
 	if err != nil {
@@ -74,16 +69,15 @@ func (c *closableHealthAndIngesterClient) Close() error {
 
 // Config is the configuration struct for the ingester client
 type Config struct {
-	MaxRecvMsgSize           int
-	CompressToIngester       bool
-	legacyCompressToIngester bool
+	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config"`
 }
 
 // RegisterFlags registers configuration settings used by the ingester client config.
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
-	// We have seen 20MB returns from queries - add a bit of headroom
-	f.IntVar(&cfg.MaxRecvMsgSize, "ingester.client.max-recv-message-size", 64*1024*1024, "Maximum message size, in bytes, this client will receive.")
-	f.BoolVar(&cfg.CompressToIngester, "ingester.client.compress-to-ingester", false, "Compress data in calls to ingesters.")
-	// moved from distributor pkg, but flag prefix left as back compat fallback for existing users.
-	f.BoolVar(&cfg.legacyCompressToIngester, "distributor.compress-to-ingester", false, "Compress data in calls to ingesters. (DEPRECATED: use ingester.client.compress-to-ingester instead")
+	cfg.GRPCClientConfig.RegisterFlags("ingester.client", f)
+
+	// Deprecated.
+	f.Int("ingester.client.max-recv-message-size", 64*1024*1024, "DEPRECATED. Maximum message size, in bytes, this client will receive.")
+	f.Bool("ingester.client.compress-to-ingester", false, "DEPRECATED. Compress data in calls to ingesters.")
+	f.Bool("distributor.compress-to-ingester", false, "DEPRECATED. Compress data in calls to ingesters. (DEPRECATED: use ingester.client.compress-to-ingester instead")
 }

pkg/querier/frontend/worker.go

@@ -14,6 +14,7 @@ import (
 	"google.golang.org/grpc/naming"
 
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/httpgrpc/server"
 	"github.com/weaveworks/common/middleware"
@@ -31,13 +32,17 @@ type WorkerConfig struct {
 	Address           string
 	Parallelism       int
 	DNSLookupDuration time.Duration
+
+	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config"`
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet.
 func (cfg *WorkerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Address, "querier.frontend-address", "", "Address of query frontend service.")
 	f.IntVar(&cfg.Parallelism, "querier.worker-parallelism", 10, "Number of simultaneous queries to process.")
 	f.DurationVar(&cfg.DNSLookupDuration, "querier.dns-lookup-period", 10*time.Second, "How often to query DNS.")
+
+	cfg.GRPCClientConfig.RegisterFlags("querier.frontend-client", f)
 }
 
 // Worker is the counter-part to the frontend, actually processing requests.
@@ -142,7 +147,7 @@ func (w *worker) watchDNSLoop() {
 
 // runMany starts N runOne loops for a given address.
 func (w *worker) runMany(ctx context.Context, address string) {
-	client, err := connect(address)
+	client, err := w.connect(address)
 	if err != nil {
 		level.Error(w.log).Log("msg", "error connecting", "addr", address, "err", err)
 		return
@@ -206,13 +211,14 @@ func (w *worker) process(ctx context.Context, c Frontend_ProcessClient) error {
 	}
 }
 
-func connect(address string) (FrontendClient, error) {
+func (w *worker) connect(address string) (FrontendClient, error) {
 	conn, err := grpc.Dial(
 		address,
 		grpc.WithInsecure(),
 		grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
 			middleware.ClientUserHeaderInterceptor,
 		)),
+		w.cfg.GRPCClientConfig.DialOption(),
 	)
 	if err != nil {
 		return nil, err

pkg/util/grpcclient/grpcclient.go

@@ -0,0 +1,37 @@
+package grpcclient
+
+import (
+	"flag"
+
+	"google.golang.org/grpc"
+)
+
+// Config for a gRPC client.
+type Config struct {
+	MaxRecvMsgSize     int  `yaml:"max_recv_msg_size"`
+	MaxSendMsgSize     int  `yaml:"max_send_msg_size"`
+	UseGzipCompression bool `yaml:"use_gzip_compression"`
+}
+
+// RegisterFlags registers flags.
+func (cfg *Config) RegisterFlags(prefix string, f *flag.FlagSet) {
+	f.IntVar(&cfg.MaxRecvMsgSize, prefix+".grpc-max-recv-msg-size", 100<<20, "gRPC client max receive message size (bytes).")
+	f.IntVar(&cfg.MaxSendMsgSize, prefix+".grpc-max-send-msg-size", 16<<20, "gRPC client max send message size (bytes).")
+	f.BoolVar(&cfg.UseGzipCompression, prefix+".grpc-use-gzip-compression", false, "Use compression when sending messages.")
+}
+
+// CallOptions returns the config in terms of CallOptions.
+func (cfg *Config) CallOptions() []grpc.CallOption {
+	var opts []grpc.CallOption
+	opts = append(opts, grpc.MaxCallRecvMsgSize(cfg.MaxRecvMsgSize))
+	opts = append(opts, grpc.MaxCallSendMsgSize(cfg.MaxSendMsgSize))
+	if cfg.UseGzipCompression {
+		opts = append(opts, grpc.UseCompressor("gzip"))
+	}
+	return opts
+}
+
+// DialOption returns the config as a grpc.DialOptions.
+func (cfg *Config) DialOption() grpc.DialOption {
+	return grpc.WithDefaultCallOptions(cfg.CallOptions()...)
+}