Add readiness check for distributor (#5142)

* add readiness check for distributor

* Add more tests

* Switch v2 run configuration to be v2-only and not combined

Author: aleks-p

.idea/runConfigurations/v2.xml

@@ -61,6 +61,11 @@ type PushClient interface {
 	Push(context.Context, *connect.Request[pushv1.PushRequest]) (*connect.Response[pushv1.PushResponse], error)
 }
 
+type SegmentWriterClient interface {
+	Push(context.Context, *segmentwriterv1.PushRequest) (*segmentwriterv1.PushResponse, error)
+	CheckReady(context.Context) error
+}
+
 const (
 	// distributorRingKey is the key under which we store the distributors ring in the KVStore.
 	distributorRingKey = "distributor"
@@ -121,7 +126,7 @@ type Distributor struct {
 	profileSizeStats        *usagestats.MultiStatistics
 
 	router        *writepath.Router
-	segmentWriter writepath.SegmentWriterClient
+	segmentWriter SegmentWriterClient
 }
 
 type Limits interface {
@@ -157,7 +162,7 @@ func New(
 	limits Limits,
 	reg prometheus.Registerer,
 	logger log.Logger,
-	segmentWriter writepath.SegmentWriterClient,
+	segmentWriter SegmentWriterClient,
 	ingesterClientsOptions ...connect.ClientOption,
 ) (*Distributor, error) {
 	ingesterClientsOptions = append(
@@ -240,6 +245,54 @@ func (d *Distributor) stopping(_ error) error {
 	return services.StopManagerAndAwaitStopped(context.Background(), d.subservices)
 }
 
+// CheckReady reports whether the distributor is ready to serve requests.
+// It verifies the destinations selected by the deployment default write path,
+// so the distributor does not accept traffic before the relevant ring is
+// populated during rollouts.
+func (d *Distributor) CheckReady(ctx context.Context) error {
+	if s := d.State(); s != services.Running && s != services.Stopping {
+		return fmt.Errorf("distributor not ready: %v", s)
+	}
+
+	switch d.limits.WritePathOverrides("").WritePath {
+	case writepath.SegmentWriterPath:
+		return d.checkSegmentWriterReady(ctx)
+	case writepath.CombinedPath:
+		if err := d.checkIngesterRingReady(ctx); err != nil {
+			return fmt.Errorf("ingester write path: %w", err)
+		}
+		if err := d.checkSegmentWriterReady(ctx); err != nil {
+			return fmt.Errorf("segment-writer write path: %w", err)
+		}
+		return nil
+	default:
+		return d.checkIngesterRingReady(ctx)
+	}
+}
+
+// checkIngesterRingReady reports whether the ingester ring has at least one
+// healthy instance available for writes.
+func (d *Distributor) checkIngesterRingReady(context.Context) error {
+	if d.ingestersRing == nil {
+		return errors.New("ingester ring not configured")
+	}
+	rs, err := d.ingestersRing.GetAllHealthy(ring.Write)
+	if err != nil {
+		return fmt.Errorf("ingester ring: %w", err)
+	}
+	if len(rs.Instances) == 0 {
+		return errors.New("ingester ring has no healthy instances")
+	}
+	return nil
+}
+
+func (d *Distributor) checkSegmentWriterReady(ctx context.Context) error {
+	if d.segmentWriter == nil {
+		return errors.New("segment-writer client not configured")
+	}
+	return d.segmentWriter.CheckReady(ctx)
+}
+
 func isKnownConnectError(err error) bool {
 	ce := new(connect.Error)
 	if !errors.As(err, &ce) {

pkg/distributor/distributor.go

@@ -0,0 +1,157 @@
+package distributor
+
+import (
+	"context"
+	"flag"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/grafana/dskit/grpcclient"
+	"github.com/grafana/dskit/ring"
+	"github.com/grafana/dskit/ring/client"
+	"github.com/grafana/dskit/services"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/pyroscope/v2/pkg/clientpool"
+	"github.com/grafana/pyroscope/v2/pkg/distributor/writepath"
+	segmentwriterclient "github.com/grafana/pyroscope/v2/pkg/segmentwriter/client"
+	"github.com/grafana/pyroscope/v2/pkg/testhelper"
+	"github.com/grafana/pyroscope/v2/pkg/validation"
+)
+
+func newReadinessSegmentWriterClient(t *testing.T, ctx context.Context, logger log.Logger, r ring.ReadRing) *segmentwriterclient.Client {
+	t.Helper()
+	var grpcCfg grpcclient.Config
+	grpcCfg.RegisterFlags(flag.NewFlagSet("", flag.PanicOnError))
+
+	swClient, err := segmentwriterclient.NewSegmentWriterClient(
+		grpcCfg, logger, nil, r, nil,
+	)
+	require.NoError(t, err)
+	require.NoError(t, services.StartAndAwaitRunning(ctx, swClient.Service()))
+	t.Cleanup(func() {
+		_ = services.StopAndAwaitTerminated(context.Background(), swClient.Service())
+	})
+	return swClient
+}
+
+func newReadinessDistributor(
+	t *testing.T,
+	ctx context.Context,
+	logger log.Logger,
+	path writepath.WritePath,
+	ingesterRing ring.ReadRing,
+	swClient SegmentWriterClient,
+) *Distributor {
+	t.Helper()
+	overrides := validation.MockOverrides(func(defaults *validation.Limits, _ map[string]*validation.Limits) {
+		defaults.WritePathOverrides.WritePath = path
+	})
+
+	d, err := New(
+		Config{
+			DistributorRing: ringConfig,
+			PoolConfig:      clientpool.PoolConfig{ClientCleanupPeriod: time.Second},
+		},
+		ingesterRing,
+		&poolFactory{f: func(string) (client.PoolClient, error) {
+			return newFakeIngester(t, false), nil
+		}},
+		overrides,
+		nil,
+		logger,
+		swClient,
+	)
+	require.NoError(t, err)
+	require.NoError(t, services.StartAndAwaitRunning(ctx, d))
+	t.Cleanup(func() {
+		_ = services.StopAndAwaitTerminated(context.Background(), d)
+	})
+	require.Equal(t, services.Running, d.State())
+	return d
+}
+
+func TestDistributor_CheckReady_V2SegmentWriterPath(t *testing.T) {
+	logger := log.NewLogfmtLogger(os.Stdout)
+	ctx := context.Background()
+
+	tests := []struct {
+		name              string
+		segmentWriterRing ring.ReadRing
+		wantErr           bool
+	}{
+		{
+			name:              "ready when segment-writer ring has healthy instances",
+			segmentWriterRing: testhelper.NewMockRing([]ring.InstanceDesc{{Addr: "foo"}}, 1),
+		},
+		{
+			name:              "not ready when segment-writer ring is empty",
+			segmentWriterRing: testhelper.NewMockRing(nil, 1),
+			wantErr:           true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			swClient := newReadinessSegmentWriterClient(t, ctx, logger, tt.segmentWriterRing)
+			d := newReadinessDistributor(
+				t,
+				ctx,
+				logger,
+				writepath.SegmentWriterPath,
+				testhelper.NewMockRing([]ring.InstanceDesc{{Addr: "foo"}}, 3),
+				swClient,
+			)
+
+			err := d.CheckReady(ctx)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestDistributor_CheckReady_V1IngesterPath(t *testing.T) {
+	logger := log.NewLogfmtLogger(os.Stdout)
+	ctx := context.Background()
+
+	tests := []struct {
+		name         string
+		ingesterRing ring.ReadRing
+		wantErr      bool
+	}{
+		{
+			name:         "ready when ingester ring has healthy instances",
+			ingesterRing: testhelper.NewMockRing([]ring.InstanceDesc{{Addr: "foo"}}, 3),
+		},
+		{
+			name:         "not ready when ingester ring is empty",
+			ingesterRing: testhelper.NewMockRing(nil, 3),
+			wantErr:      true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			d := newReadinessDistributor(
+				t,
+				ctx,
+				logger,
+				writepath.IngesterPath,
+				tt.ingesterRing,
+				nil,
+			)
+
+			err := d.CheckReady(ctx)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}

pkg/distributor/distributor_readiness_test.go

@@ -46,7 +46,6 @@ import (
 	pprof2 "github.com/grafana/pyroscope/v2/pkg/pprof"
 	pproftesthelper "github.com/grafana/pyroscope/v2/pkg/pprof/testhelper"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
-	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockwritepath"
 	"github.com/grafana/pyroscope/v2/pkg/testhelper"
 	"github.com/grafana/pyroscope/v2/pkg/util"
 	"github.com/grafana/pyroscope/v2/pkg/validation"
@@ -1795,7 +1794,7 @@ func Test_SampleLabels_SegmentWriter(t *testing.T) {
 				{Addr: "foo"},
 			}, 3), &poolFactory{func(addr string) (client.PoolClient, error) {
 				return newFakeIngester(t, false), nil
-			}}, overrides, nil, log.NewLogfmtLogger(os.Stdout), new(mockwritepath.MockSegmentWriterClient))
+			}}, overrides, nil, log.NewLogfmtLogger(os.Stdout), nil)
 
 			require.NoError(t, err)
 			var series []*distributormodel.ProfileSeries

pkg/distributor/distributor_test.go

@@ -34,7 +34,6 @@ import (
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
 
-	statusv1 "github.com/grafana/pyroscope/api/gen/proto/go/status/v1"
 	"github.com/grafana/pyroscope/v2/pkg/adhocprofiles"
 	apiversion "github.com/grafana/pyroscope/v2/pkg/api/version"
 	"github.com/grafana/pyroscope/v2/pkg/compactor"
@@ -60,6 +59,8 @@ import (
 	httputil "github.com/grafana/pyroscope/v2/pkg/util/http"
 	"github.com/grafana/pyroscope/v2/pkg/validation"
 	"github.com/grafana/pyroscope/v2/pkg/validation/exporter"
+
+	statusv1 "github.com/grafana/pyroscope/api/gen/proto/go/status/v1"
 )
 
 // The various modules that make up Pyroscope.
@@ -322,10 +323,15 @@ func (f *Pyroscope) initGRPCGateway() (services.Service, error) {
 func (f *Pyroscope) initDistributor() (services.Service, error) {
 	f.Cfg.Distributor.DistributorRing.ListenPort = f.Cfg.Server.HTTPListenPort
 	logger := log.With(f.logger, "component", "distributor")
-	d, err := distributor.New(f.Cfg.Distributor, f.ingesterRing, nil, f.Overrides, f.reg, logger, f.segmentWriterClient, f.auth)
+	var swClient distributor.SegmentWriterClient
+	if f.segmentWriterClient != nil {
+		swClient = f.segmentWriterClient
+	}
+	d, err := distributor.New(f.Cfg.Distributor, f.ingesterRing, nil, f.Overrides, f.reg, logger, swClient, f.auth)
 	if err != nil {
 		return nil, err
 	}
+	f.distributor = d
 	f.API.RegisterDistributor(d, f.Overrides, f.Cfg.Server)
 
 	if store, err := debuginfo.NewStore(f.logger, f.storageBucket, f.Cfg.DebugInfo); err != nil {

pkg/pyroscope/modules.go

@@ -392,8 +392,9 @@ type Pyroscope struct {
 
 	grpcGatewayMux *grpcgw.ServeMux
 
-	auth     connect.Option
-	frontend *frontend.Frontend
+	auth        connect.Option
+	frontend    *frontend.Frontend
+	distributor *distributor.Distributor
 
 	segmentWriter         *segmentwriter.SegmentWriterService
 	segmentWriterClient   *segmentwriterclient.Client
@@ -788,6 +789,13 @@ func (f *Pyroscope) readyHandler(sm *services.Manager) http.HandlerFunc {
 			}
 		}
 
+		if f.distributor != nil {
+			if err := f.distributor.CheckReady(r.Context()); err != nil {
+				http.Error(w, "Distributor not ready: "+err.Error(), http.StatusServiceUnavailable)
+				return
+			}
+		}
+
 		util.WriteTextResponse(w, "ready")
 	}
 }

pkg/pyroscope/pyroscope.go

@@ -183,6 +183,24 @@ func NewSegmentWriterClient(
 
 func (c *Client) Service() services.Service { return c.service }
 
+// CheckReady reports whether the client can dispatch requests, i.e. its
+// service is Running and the segment-writer ring has at least one healthy
+// instance. Used by callers (e.g. the distributor) to gate readiness so
+// that traffic isn't accepted before the ring has been populated.
+func (c *Client) CheckReady(_ context.Context) error {
+	if state := c.service.State(); state != services.Running {
+		return fmt.Errorf("segment-writer client not running: %s", state)
+	}
+	rs, err := c.ring.GetAllHealthy(ring.Reporting)
+	if err != nil {
+		return fmt.Errorf("segment-writer ring: %w", err)
+	}
+	if len(rs.Instances) == 0 {
+		return errors.New("segment-writer ring has no healthy instances")
+	}
+	return nil
+}
+
 func (c *Client) starting(ctx context.Context) error {
 	// Warm up connections. The pool does not do this.
 	instances, err := c.ring.GetAllHealthy(ring.Reporting)