feat(v2): cancel artificial delay on async ingest

pkg/distributor/distributor.go

@@ -140,9 +140,9 @@ type Limits interface {
 	EnforceLabelsOrder(tenantID string) bool
 	IngestionRelabelingRules(tenantID string) []*relabel.Config
 	DistributorUsageGroups(tenantID string) *validation.UsageGroupConfig
+	WritePathOverrides(tenantID string) writepath.Config
 	validation.ProfileValidationLimits
 	aggregator.Limits
-	writepath.Overrides
 }
 
 func New(
@@ -184,11 +184,7 @@ func New(
 
 	ingesterRoute := writepath.IngesterFunc(d.sendRequestsToIngester)
 	segmentWriterRoute := writepath.IngesterFunc(d.sendRequestsToSegmentWriter)
-	d.router = writepath.NewRouter(
-		logger, reg, limits,
-		ingesterRoute,
-		segmentWriterRoute,
-	)
+	d.router = writepath.NewRouter(logger, reg, ingesterRoute, segmentWriterRoute)
 
 	var err error
 	subservices := []services.Service(nil)
@@ -433,7 +429,8 @@ func (d *Distributor) PushParsed(ctx context.Context, req *distributormodel.Push
 	// functions to send the request to the appropriate service; these are
 	// called independently, and may be called concurrently: the request is
 	// cloned in this case – the callee may modify the request safely.
-	if err = d.router.Send(ctx, req); err != nil {
+	config := d.limits.WritePathOverrides(req.TenantID)
+	if err = d.router.Send(ctx, req, config); err != nil {
 		return nil, err
 	}
 
@@ -518,7 +515,8 @@ func (d *Distributor) aggregate(ctx context.Context, req *distributormodel.PushR
 					Annotations: annotations,
 				}},
 			}
-			return d.router.Send(localCtx, aggregated)
+			config := d.limits.WritePathOverrides(req.TenantID)
+			return d.router.Send(localCtx, aggregated, config)
 		})()
 		if sendErr != nil {
 			_ = level.Error(d.logger).Log("msg", "failed to handle aggregation", "tenant", req.TenantID, "err", err)

pkg/distributor/write_path/router.go

@@ -20,6 +20,7 @@ import (
 	"github.com/grafana/pyroscope/pkg/tenant"
 	"github.com/grafana/pyroscope/pkg/util"
 	"github.com/grafana/pyroscope/pkg/util/connectgrpc"
+	"github.com/grafana/pyroscope/pkg/util/delayhandler"
 	httputil "github.com/grafana/pyroscope/pkg/util/http"
 )
 
@@ -43,17 +44,12 @@ func (f IngesterFunc) Push(
 	return f(ctx, req)
 }
 
-type Overrides interface {
-	WritePathOverrides(tenantID string) Config
-}
-
 type Router struct {
 	service  services.Service
 	inflight sync.WaitGroup
 
-	logger    log.Logger
-	overrides Overrides
-	metrics   *metrics
+	logger  log.Logger
+	metrics *metrics
 
 	ingester  IngesterClient
 	segwriter IngesterClient
@@ -62,13 +58,11 @@ type Router struct {
 func NewRouter(
 	logger log.Logger,
 	registerer prometheus.Registerer,
-	overrides Overrides,
 	ingester IngesterClient,
 	segwriter IngesterClient,
 ) *Router {
 	r := &Router{
 		logger:    logger,
-		overrides: overrides,
 		metrics:   newMetrics(registerer),
 		ingester:  ingester,
 		segwriter: segwriter,
@@ -92,15 +86,17 @@ func (m *Router) running(ctx context.Context) error {
 	return nil
 }
 
-func (m *Router) Send(ctx context.Context, req *distributormodel.PushRequest) error {
-	config := m.overrides.WritePathOverrides(req.TenantID)
+func (m *Router) Send(ctx context.Context, req *distributormodel.PushRequest, config Config) error {
+	if config.AsyncIngest {
+		delayhandler.CancelDelay(ctx)
+	}
 	switch config.WritePath {
 	case SegmentWriterPath:
-		return m.send(m.segwriterRoute(true))(ctx, req)
+		return m.sendToSegmentWriterOnly(ctx, req, &config)
 	case CombinedPath:
 		return m.sendToBoth(ctx, req, &config)
 	default:
-		return m.send(m.ingesterRoute())(ctx, req)
+		return m.sendToIngesterOnly(ctx, req)
 	}
 }
 
@@ -160,7 +156,7 @@ func (m *Router) sendToBoth(ctx context.Context, req *distributormodel.PushReque
 		// The request is to be sent to both asynchronously, therefore we're
 		// cloning it. We do not wait for the secondary request to complete.
 		// On shutdown, however, we will wait for all inflight requests.
-		segwriter.client = m.sendClone(ctx, req.Clone(), segwriter.client, config)
+		segwriter.client = m.detachedClient(ctx, req.Clone(), segwriter.client, config)
 	}
 
 	if segwriter != nil {
@@ -179,6 +175,22 @@ func (m *Router) sendToBoth(ctx context.Context, req *distributormodel.PushReque
 	return nil
 }
 
+func (m *Router) sendToSegmentWriterOnly(ctx context.Context, req *distributormodel.PushRequest, config *Config) error {
+	r := m.segwriterRoute(true)
+	if !config.AsyncIngest {
+		return m.send(r)(ctx, req)
+	}
+	r.client = m.detachedClient(ctx, req, r.client, config)
+	m.sendAsync(ctx, req, r)
+	return nil
+}
+
+func (m *Router) sendToIngesterOnly(ctx context.Context, req *distributormodel.PushRequest) error {
+	// NOTE(kolesnikovae): If we also want to support async requests to ingesters,
+	// we should implement it here and in sendToBoth.
+	return m.send(m.ingesterRoute())(ctx, req)
+}
+
 type sendFunc func(context.Context, *distributormodel.PushRequest) error
 
 type route struct {
@@ -187,7 +199,9 @@ type route struct {
 	primary bool
 }
 
-func (m *Router) sendClone(ctx context.Context, req *distributormodel.PushRequest, client IngesterClient, config *Config) IngesterFunc {
+// detachedClient creates a new IngesterFunc that wraps the call with a local context
+// that has a timeout and tenant ID injected so it can be used for asynchronous requests.
+func (m *Router) detachedClient(ctx context.Context, req *distributormodel.PushRequest, client IngesterClient, config *Config) IngesterFunc {
 	return func(context.Context, *distributormodel.PushRequest) (*connect.Response[pushv1.PushResponse], error) {
 		localCtx, cancel := context.WithTimeout(context.Background(), config.SegmentWriterTimeout)
 		localCtx = tenant.InjectTenantID(localCtx, req.TenantID)