Implement stream connection for remote write

docs/configuration/config-file-reference.md

@@ -2657,6 +2657,11 @@ ha_tracker:
 # CLI flag: -distributor.sign-write-requests
 [sign_write_requests: <boolean> | default = false]
 
+# EXPERIMENTAL: If enabled, distributor would use stream connection to send
+# requests to ingesters.
+# CLI flag: -distributor.use-stream-push
+[use_stream_push: <boolean> | default = false]
+
 ring:
   kvstore:
     # Backend storage to use for the ring. Supported values are: consul, etcd,

pkg/cortexpb/cortex.proto

@@ -21,6 +21,11 @@ message WriteRequest {
   bool skip_label_name_validation = 1000; //set intentionally high to keep WriteRequest compatible with upstream Prometheus
 }
 
+message StreamWriteRequest {
+  string TenantID = 1;
+  WriteRequest Request = 2;
+}
+
 message WriteResponse {}
 
 message TimeSeries {

pkg/distributor/distributor.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	io "io"
+	"io"
 	"net/http"
 	"sort"
 	"strings"
@@ -150,6 +150,7 @@ type Config struct {
 	ShardByAllLabels         bool   `yaml:"shard_by_all_labels"`
 	ExtendWrites             bool   `yaml:"extend_writes"`
 	SignWriteRequestsEnabled bool   `yaml:"sign_write_requests"`
+	UseStreamPush            bool   `yaml:"use_stream_push"`
 
 	// Distributors ring
 	DistributorRing RingConfig `yaml:"ring"`
@@ -204,6 +205,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.ExtraQueryDelay, "distributor.extra-query-delay", 0, "Time to wait before sending more than the minimum successful query requests.")
 	f.BoolVar(&cfg.ShardByAllLabels, "distributor.shard-by-all-labels", false, "Distribute samples based on all labels, as opposed to solely by user and metric name.")
 	f.BoolVar(&cfg.SignWriteRequestsEnabled, "distributor.sign-write-requests", false, "EXPERIMENTAL: If enabled, sign the write request between distributors and ingesters.")
+	f.BoolVar(&cfg.UseStreamPush, "distributor.use-stream-push", false, "EXPERIMENTAL: If enabled, distributor would use stream connection to send requests to ingesters.")
 	f.StringVar(&cfg.ShardingStrategy, "distributor.sharding-strategy", util.ShardingStrategyDefault, fmt.Sprintf("The sharding strategy to use. Supported values are: %s.", strings.Join(supportedShardingStrategies, ", ")))
 	f.BoolVar(&cfg.ExtendWrites, "distributor.extend-writes", true, "Try writing to an additional ingester in the presence of an ingester not in the ACTIVE state. It is useful to disable this along with -ingester.unregister-on-shutdown=false in order to not spread samples to extra ingesters during rolling restarts with consistent naming.")
 	f.BoolVar(&cfg.ZoneResultsQuorumMetadata, "distributor.zone-results-quorum-metadata", false, "Experimental, this flag may change in the future. If zone awareness and this both enabled, when querying metadata APIs (labels names and values for now), only results from quorum number of zones will be included.")
@@ -242,7 +244,7 @@ const (
 func New(cfg Config, clientConfig ingester_client.Config, limits *validation.Overrides, ingestersRing ring.ReadRing, canJoinDistributorsRing bool, reg prometheus.Registerer, log log.Logger) (*Distributor, error) {
 	if cfg.IngesterClientFactory == nil {
 		cfg.IngesterClientFactory = func(addr string) (ring_client.PoolClient, error) {
-			return ingester_client.MakeIngesterClient(addr, clientConfig)
+			return ingester_client.MakeIngesterClient(addr, clientConfig, cfg.UseStreamPush)
 		}
 	}
 
@@ -1134,20 +1136,29 @@ func (d *Distributor) send(ctx context.Context, ingester ring.InstanceDesc, time
 
 	c := h.(ingester_client.HealthAndIngesterClient)
 
-	req := cortexpb.PreallocWriteRequestFromPool()
-	req.Timeseries = timeseries
-	req.Metadata = metadata
-	req.Source = source
-
 	d.inflightClientRequests.Inc()
 	defer d.inflightClientRequests.Dec()
 
-	_, err = c.PushPreAlloc(ctx, req)
+	if d.cfg.UseStreamPush {
+		req := &cortexpb.WriteRequest{
+			Timeseries: timeseries,
+			Metadata:   metadata,
+			Source:     source,
+		}
+		_, err = c.PushStreamConnection(ctx, req)
+	} else {
+		req := cortexpb.PreallocWriteRequestFromPool()
+		req.Timeseries = timeseries
+		req.Metadata = metadata
+		req.Source = source
 
-	// We should not reuse the req in case of errors:
-	// See: https://github.com/grpc/grpc-go/issues/6355
-	if err == nil {
-		cortexpb.ReuseWriteRequest(req)
+		_, err = c.PushPreAlloc(ctx, req)
+
+		// We should not reuse the req in case of errors:
+		// See: https://github.com/grpc/grpc-go/issues/6355
+		if err == nil {
+			cortexpb.ReuseWriteRequest(req)
+		}
 	}
 
 	if len(metadata) > 0 {

pkg/distributor/distributor_test.go

@@ -3307,6 +3307,10 @@ func (i *mockIngester) PushPreAlloc(ctx context.Context, in *cortexpb.PreallocWr
 	return i.Push(ctx, &in.WriteRequest, opts...)
 }
 
+func (i *mockIngester) PushStreamConnection(ctx context.Context, in *cortexpb.WriteRequest, opts ...grpc.CallOption) (*cortexpb.WriteResponse, error) {
+	return i.Push(ctx, in, opts...)
+}
+
 func (i *mockIngester) Push(ctx context.Context, req *cortexpb.WriteRequest, opts ...grpc.CallOption) (*cortexpb.WriteResponse, error) {
 	i.Lock()
 	defer i.Unlock()

pkg/ingester/client/client.go

@@ -3,15 +3,19 @@ package client
 import (
 	"context"
 	"flag"
+	"fmt"
+	"io"
 
 	"github.com/cortexproject/cortex/pkg/cortexpb"
+	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	"github.com/cortexproject/cortex/pkg/util/grpcencoding/snappyblock"
 
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/weaveworks/common/user"
 	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
@@ -31,6 +35,8 @@ var ingesterClientInflightPushRequests = promauto.NewGaugeVec(prometheus.GaugeOp
 
 var errTooManyInflightPushRequests = errors.New("too many inflight push requests in ingester client")
 
+const INGESTER_CLIENT_STREAM_WORKER_COUNT = 1024
+
 // ClosableClientConn is grpc.ClientConnInterface with Close function
 type ClosableClientConn interface {
 	grpc.ClientConnInterface
@@ -43,6 +49,15 @@ type HealthAndIngesterClient interface {
 	grpc_health_v1.HealthClient
 	Close() error
 	PushPreAlloc(ctx context.Context, in *cortexpb.PreallocWriteRequest, opts ...grpc.CallOption) (*cortexpb.WriteResponse, error)
+	PushStreamConnection(ctx context.Context, in *cortexpb.WriteRequest, opts ...grpc.CallOption) (*cortexpb.WriteResponse, error)
+}
+
+type streamWriteJob struct {
+	req    *cortexpb.StreamWriteRequest
+	resp   *cortexpb.WriteResponse
+	ctx    context.Context
+	cancel context.CancelFunc
+	err    error
 }
 
 type closableHealthAndIngesterClient struct {
@@ -53,6 +68,9 @@ type closableHealthAndIngesterClient struct {
 	maxInflightPushRequests int64
 	inflightRequests        atomic.Int64
 	inflightPushRequests    *prometheus.GaugeVec
+	streamPushChan          chan *streamWriteJob
+	streamCtx               context.Context
+	streamCancel            context.CancelFunc
 }
 
 func (c *closableHealthAndIngesterClient) PushPreAlloc(ctx context.Context, in *cortexpb.PreallocWriteRequest, opts ...grpc.CallOption) (*cortexpb.WriteResponse, error) {
@@ -72,6 +90,32 @@ func (c *closableHealthAndIngesterClient) Push(ctx context.Context, in *cortexpb
 	})
 }
 
+func (c *closableHealthAndIngesterClient) PushStreamConnection(ctx context.Context, in *cortexpb.WriteRequest, opts ...grpc.CallOption) (*cortexpb.WriteResponse, error) {
+	return c.handlePushRequest(func() (*cortexpb.WriteResponse, error) {
+		tenantID, err := tenant.TenantID(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		streamReq := &cortexpb.StreamWriteRequest{
+			TenantID: tenantID,
+			Request:  in,
+		}
+
+		reqCtx, reqCancel := context.WithCancel(ctx)
+		defer reqCancel()
+
+		job := &streamWriteJob{
+			req:    streamReq,
+			ctx:    reqCtx,
+			cancel: reqCancel,
+		}
+		c.streamPushChan <- job
+		<-reqCtx.Done()
+		return job.resp, job.err
+	})
+}
+
 func (c *closableHealthAndIngesterClient) handlePushRequest(mainFunc func() (*cortexpb.WriteResponse, error)) (*cortexpb.WriteResponse, error) {
 	currentInflight := c.inflightRequests.Inc()
 	c.inflightPushRequests.WithLabelValues(c.addr).Set(float64(currentInflight))
@@ -85,7 +129,7 @@ func (c *closableHealthAndIngesterClient) handlePushRequest(mainFunc func() (*co
 }
 
 // MakeIngesterClient makes a new IngesterClient
-func MakeIngesterClient(addr string, cfg Config) (HealthAndIngesterClient, error) {
+func MakeIngesterClient(addr string, cfg Config, useStreamConnection bool) (HealthAndIngesterClient, error) {
 	dialOpts, err := cfg.GRPCClientConfig.DialOption(grpcclient.Instrument(ingesterClientRequestDuration))
 	if err != nil {
 		return nil, err
@@ -94,21 +138,88 @@ func MakeIngesterClient(addr string, cfg Config) (HealthAndIngesterClient, error
 	if err != nil {
 		return nil, err
 	}
-	return &closableHealthAndIngesterClient{
+	streamCtx, streamCancel := context.WithCancel(context.Background())
+	c := &closableHealthAndIngesterClient{
 		IngesterClient:          NewIngesterClient(conn),
 		HealthClient:            grpc_health_v1.NewHealthClient(conn),
 		conn:                    conn,
 		addr:                    addr,
 		maxInflightPushRequests: cfg.MaxInflightPushRequests,
 		inflightPushRequests:    ingesterClientInflightPushRequests,
-	}, nil
+		streamPushChan:          make(chan *streamWriteJob, INGESTER_CLIENT_STREAM_WORKER_COUNT),
+		streamCtx:               streamCtx,
+		streamCancel:            streamCancel,
+	}
+	if useStreamConnection {
+		err = c.Run()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return c, nil
 }
 
 func (c *closableHealthAndIngesterClient) Close() error {
 	c.inflightPushRequests.DeleteLabelValues(c.addr)
+	c.streamCancel()
 	return c.conn.Close()
 }
 
+func (c *closableHealthAndIngesterClient) Run() error {
+	var err error
+	for i := 0; i < INGESTER_CLIENT_STREAM_WORKER_COUNT; i++ {
+		workerCtx := user.InjectOrgID(c.streamCtx, fmt.Sprintf("stream-worker-%d", i))
+		go func() {
+			for {
+				select {
+				case <-workerCtx.Done():
+					return
+				default:
+					err = c.worker(workerCtx)
+					if err != nil {
+						return
+					}
+				}
+			}
+		}()
+	}
+	return err
+}
+
+func (c *closableHealthAndIngesterClient) worker(ctx context.Context) error {
+	stream, err := c.PushStream(ctx)
+	if err != nil {
+		return err
+	}
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		case job := <-c.streamPushChan:
+			err = stream.Send(job.req)
+			if err == io.EOF {
+				job.resp = &cortexpb.WriteResponse{}
+				job.cancel()
+				return nil
+			}
+			if err != nil {
+				job.err = err
+				job.cancel()
+				continue
+			}
+			resp, err := stream.Recv()
+			if err == io.EOF {
+				job.resp = &cortexpb.WriteResponse{}
+				job.cancel()
+				return nil
+			}
+			job.resp = resp
+			job.err = err
+			job.cancel()
+		}
+	}
+}
+
 // Config is the configuration struct for the ingester client
 type Config struct {
 	GRPCClientConfig        grpcclient.ConfigWithHealthCheck `yaml:"grpc_client_config"`

pkg/ingester/client/cortex_mock_test.go

@@ -17,6 +17,11 @@ func (m *IngesterServerMock) Push(ctx context.Context, r *cortexpb.WriteRequest)
 	return args.Get(0).(*cortexpb.WriteResponse), args.Error(1)
 }
 
+func (m *IngesterServerMock) PushStream(srv Ingester_PushStreamServer) error {
+	args := m.Called(srv)
+	return args.Error(0)
+}
+
 func (m *IngesterServerMock) Query(ctx context.Context, r *QueryRequest) (*QueryResponse, error) {
 	args := m.Called(ctx, r)
 	return args.Get(0).(*QueryResponse), args.Error(1)