github.com/grafana/pyroscope@v1.18.0/pkg/util/delayhandler/delay.go (about)

     1  package delayhandler
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"strconv"
     8  	"time"
     9  
    10  	"github.com/grafana/dskit/tenant"
    11  
    12  	"github.com/grafana/pyroscope/pkg/util"
    13  )
    14  
    15  var (
    16  	// have local variables to allow for mocking in tests
    17  	timeNow   = time.Now
    18  	timeAfter = time.After
    19  )
    20  
    21  type Limits interface {
    22  	IngestionArtificialDelay(tenantID string) time.Duration
    23  }
    24  
    25  type delayCancelCtxKey struct{}
    26  
    27  func CancelDelay(ctx context.Context) {
    28  	if cancel, ok := ctx.Value(delayCancelCtxKey{}).(context.CancelFunc); ok && cancel != nil {
    29  		cancel()
    30  	}
    31  }
    32  
    33  func WithDelayCancel(ctx context.Context, cancel context.CancelFunc) context.Context {
    34  	return context.WithValue(ctx, delayCancelCtxKey{}, cancel)
    35  }
    36  
    37  func addDelayHeader(h http.Header, delay time.Duration) {
    38  	durationInMs := strconv.FormatFloat(float64(delay)/float64(time.Millisecond), 'f', -1, 64)
    39  	h.Add("Server-Timing", fmt.Sprintf("artificial_delay;dur=%s", durationInMs))
    40  }
    41  
    42  func getDelay(ctx context.Context, limits Limits) time.Duration {
    43  	tenantID, err := tenant.TenantID(ctx)
    44  	if err != nil {
    45  		return 0
    46  	}
    47  
    48  	delay := limits.IngestionArtificialDelay(tenantID)
    49  	if delay > 0 {
    50  		return util.DurationWithJitter(delay, 0.10)
    51  	}
    52  	return 0
    53  }