feat(query-backend): hedge GetRange calls to reduce tail latency

Symbol table fetches (locations, mappings, functions, strings,
stacktraces) call GetRange against object storage and are a primary
source of tail latency on the read path. This adds opt-in speculative
hedging: after a configurable delay, a second parallel GetRange is
issued; whichever response arrives first is used and the other is
cancelled.

The hedge wraps only the GetRange call, not the decode — the winning
response body is decoded once. A new Cleanup field on retry.Hedged
ensures the losing response body is always closed, preventing connection
leaks when both calls succeed.

Config flag (default 0 = disabled):
  --query-backend.block-read-hedge-after=<duration>

Author: simonswine

pkg/block/object.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/grafana/dskit/multierror"
 	"github.com/oklog/ulid/v2"
@@ -39,6 +40,7 @@ type Object struct {
 
 	memSize     int
 	downloadDir string
+	hedgeAfter  time.Duration
 }
 
 type ObjectOption func(*Object)
@@ -61,6 +63,12 @@ func WithObjectDownload(dir string) ObjectOption {
 	}
 }
 
+func WithObjectHedgeAfter(d time.Duration) ObjectOption {
+	return func(obj *Object) {
+		obj.hedgeAfter = d
+	}
+}
+
 func NewObjectFromPath(ctx context.Context, storage objstore.Bucket, path string, opts ...ObjectOption) (*Object, error) {
 	attrs, err := storage.Attributes(ctx, path)
 	if err != nil {

pkg/block/section_symbols.go

@@ -14,8 +14,11 @@ func openSymbols(ctx context.Context, s *Dataset) (err error) {
 		offset -= int64(s.offset())
 		s.symbols, err = symdb.OpenObject(ctx, s.inMemoryBucket(buf), s.obj.path, offset, size)
 	} else {
-		s.symbols, err = symdb.OpenObject(ctx, s.obj.storage, s.obj.path, offset, size,
-			symdb.WithPrefetchSize(symbolsPrefetchSize))
+		opts := []symdb.Option{symdb.WithPrefetchSize(symbolsPrefetchSize)}
+		if s.obj.hedgeAfter > 0 {
+			opts = append(opts, symdb.WithHedgeAfter(s.obj.hedgeAfter))
+		}
+		s.symbols, err = symdb.OpenObject(ctx, s.obj.storage, s.obj.path, offset, size, opts...)
 	}
 	if err != nil {
 		return fmt.Errorf("opening symbols: %w", err)

pkg/phlaredb/symdb/block_reader.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"sync"
+	"time"
 
 	"github.com/grafana/dskit/multierror"
 	"github.com/grafana/dskit/tracing"
@@ -21,6 +22,7 @@ import (
 	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
 	"github.com/grafana/pyroscope/pkg/util/bufferpool"
 	"github.com/grafana/pyroscope/pkg/util/refctr"
+	"github.com/grafana/pyroscope/pkg/util/retry"
 )
 
 type Reader struct {
@@ -38,6 +40,7 @@ type Reader struct {
 	parquetFiles *parquetFiles
 
 	prefetchSize uint64
+	hedgeAfter   time.Duration
 }
 
 type Option func(*Reader)
@@ -48,6 +51,12 @@ func WithPrefetchSize(size uint64) Option {
 	}
 }
 
+func WithHedgeAfter(d time.Duration) Option {
+	return func(r *Reader) {
+		r.hedgeAfter = d
+	}
+}
+
 func OpenObject(ctx context.Context, b objstore.BucketReader, name string, offset, size int64, options ...Option) (*Reader, error) {
 	f := block.File{
 		RelPath:   name,
@@ -560,7 +569,7 @@ func (c *stacktraceBlock) fetch(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		rc, err := c.reader.bucket.GetRange(ctx, path, c.header.Offset, c.header.Size)
+		rc, err := c.getRange(ctx, path)
 		if err != nil {
 			return err
 		}
@@ -573,15 +582,19 @@ func (c *stacktraceBlock) fetch(ctx context.Context) error {
 	})
 }
 
-func (c *stacktraceBlock) stacktracesFile() (string, error) {
-	f := c.reader.file
-	if c.reader.index.Header.Version < 3 {
-		var err error
-		if f, err = c.reader.lookupFile(StacktracesFileName); err != nil {
-			return "", err
-		}
+// getRange issues a GetRange call, hedging with a second call if hedgeAfter is set.
+func (c *stacktraceBlock) getRange(ctx context.Context, path string) (io.ReadCloser, error) {
+	if c.reader.hedgeAfter <= 0 {
+		return c.reader.bucket.GetRange(ctx, path, c.header.Offset, c.header.Size)
 	}
-	return f.RelPath, nil
+	return retry.Hedged[io.ReadCloser]{
+		Trigger:  time.After(c.reader.hedgeAfter),
+		FailFast: true,
+		Cleanup:  func(rc io.ReadCloser) { rc.Close() },
+		Call: func(ctx context.Context, _ bool) (io.ReadCloser, error) {
+			return c.reader.bucket.GetRange(ctx, path, c.header.Offset, c.header.Size)
+		},
+	}.Do(ctx)
 }
 
 func (c *stacktraceBlock) readFrom(r *bufio.Reader) error {
@@ -607,6 +620,17 @@ func (c *stacktraceBlock) readFrom(r *bufio.Reader) error {
 	return nil
 }
 
+func (c *stacktraceBlock) stacktracesFile() (string, error) {
+	f := c.reader.file
+	if c.reader.index.Header.Version < 3 {
+		var err error
+		if f, err = c.reader.lookupFile(StacktracesFileName); err != nil {
+			return "", err
+		}
+	}
+	return f.RelPath, nil
+}
+
 func (c *stacktraceBlock) release() {
 	c.r.Dec(func() {
 		c.t = nil
@@ -627,10 +651,7 @@ func (t *rawTable[T]) fetch(ctx context.Context) error {
 	span.SetTag("length", t.header.Length)
 	defer span.Finish()
 	return t.r.Inc(func() error {
-		rc, err := t.reader.bucket.GetRange(ctx,
-			t.reader.file.RelPath,
-			int64(t.header.Offset),
-			int64(t.header.Size))
+		rc, err := t.getRange(ctx)
 		if err != nil {
 			return err
 		}
@@ -643,6 +664,27 @@ func (t *rawTable[T]) fetch(ctx context.Context) error {
 	})
 }
 
+// getRange issues a GetRange call, hedging with a second call if hedgeAfter is set.
+func (t *rawTable[T]) getRange(ctx context.Context) (io.ReadCloser, error) {
+	if t.reader.hedgeAfter <= 0 {
+		return t.reader.bucket.GetRange(ctx,
+			t.reader.file.RelPath,
+			int64(t.header.Offset),
+			int64(t.header.Size))
+	}
+	return retry.Hedged[io.ReadCloser]{
+		Trigger:  time.After(t.reader.hedgeAfter),
+		FailFast: true,
+		Cleanup:  func(rc io.ReadCloser) { rc.Close() },
+		Call: func(ctx context.Context, _ bool) (io.ReadCloser, error) {
+			return t.reader.bucket.GetRange(ctx,
+				t.reader.file.RelPath,
+				int64(t.header.Offset),
+				int64(t.header.Size))
+		},
+	}.Do(ctx)
+}
+
 func (t *rawTable[T]) readFrom(r *bufio.Reader) error {
 	crc := crc32.New(castagnoli)
 	tee := io.TeeReader(r, crc)

pkg/pyroscope/modules_experimental.go

@@ -370,7 +370,8 @@ func (f *Pyroscope) initQueryBackend() (services.Service, error) {
 		logger,
 		f.reg,
 		f.queryBackendClient,
-		querybackend.NewBlockReader(f.logger, f.storageBucket, f.reg),
+		querybackend.NewBlockReader(f.logger, f.storageBucket, f.reg,
+			querybackend.WithBlockReaderHedgeAfter(f.Cfg.QueryBackend.BlockReadHedgeAfter)),
 	)
 	if err != nil {
 		return nil, err

pkg/querybackend/backend.go

@@ -21,14 +21,16 @@ import (
 )
 
 type Config struct {
-	Address          string            `yaml:"address"`
-	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config" doc:"description=Configures the gRPC client used to communicate between the query-frontends and the query-schedulers."`
-	ClientTimeout    time.Duration     `yaml:"client_timeout"`
+	Address             string            `yaml:"address"`
+	GRPCClientConfig    grpcclient.Config `yaml:"grpc_client_config" doc:"description=Configures the gRPC client used to communicate between the query-frontends and the query-schedulers."`
+	ClientTimeout       time.Duration     `yaml:"client_timeout"`
+	BlockReadHedgeAfter time.Duration     `yaml:"block_read_hedge_after" category:"advanced"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Address, "query-backend.address", "localhost:9095", "")
 	f.DurationVar(&cfg.ClientTimeout, "query-backend.client-timeout", 30*time.Second, "Timeout for query-backend client requests.")
+	f.DurationVar(&cfg.BlockReadHedgeAfter, "query-backend.block-read-hedge-after", 0, "If non-zero, issue a speculative second GetRange request for symbol tables after this duration. 0 disables hedging.")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-backend.grpc-client-config", f)
 }
 

pkg/querybackend/block_reader.go

@@ -50,8 +50,9 @@ type BlockReader struct {
 	log     log.Logger
 	storage objstore.Bucket
 
-	metrics  *metrics
-	hostname string
+	metrics    *metrics
+	hostname   string
+	hedgeAfter time.Duration
 
 	// TODO:
 	//  - Use a worker pool instead of the errgroup.
@@ -61,14 +62,26 @@ type BlockReader struct {
 	//    Instead, they should share the processing pipeline, if possible.
 }
 
-func NewBlockReader(logger log.Logger, storage objstore.Bucket, reg prometheus.Registerer) *BlockReader {
+func NewBlockReader(logger log.Logger, storage objstore.Bucket, reg prometheus.Registerer, opts ...BlockReaderOption) *BlockReader {
 	hostname, _ := os.Hostname()
-	return &BlockReader{
+	br := &BlockReader{
 		log:      logger,
 		storage:  storage,
 		metrics:  newMetrics(reg),
 		hostname: hostname,
 	}
+	for _, opt := range opts {
+		opt(br)
+	}
+	return br
+}
+
+type BlockReaderOption func(*BlockReader)
+
+func WithBlockReaderHedgeAfter(d time.Duration) BlockReaderOption {
+	return func(br *BlockReader) {
+		br.hedgeAfter = d
+	}
 }
 
 func (b *BlockReader) Invoke(
@@ -115,7 +128,7 @@ func (b *BlockReader) Invoke(
 		}
 		blocksCount++
 		datasetsCount += int64(len(md.Datasets))
-		obj := block.NewObject(b.storage, md)
+		obj := block.NewObject(b.storage, md, block.WithObjectHedgeAfter(b.hedgeAfter))
 		g.Go(util.RecoverPanic((&blockContext{
 			ctx:             ctx,
 			log:             b.log,

pkg/util/retry/hedged.go

@@ -26,23 +26,36 @@ type Hedged[T any] struct {
 	//  - the result received first is returned, regardless of anything.
 	//  - if Call fails before the trigger fires, it won't be retried.
 	FailFast bool
+
+	// Cleanup is called on the result of a losing attempt when it succeeded
+	// but another attempt already won. Use this to release resources (e.g.,
+	// close an io.ReadCloser) that would otherwise be abandoned.
+	Cleanup func(T)
 }
 
 type Call[T any] func(ctx context.Context, isRetry bool) (T, error)
 
 func (s Hedged[T]) Do(ctx context.Context) (T, error) {
-	attemptCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
+	// Each attempt gets its own independent cancellable context derived from
+	// the parent. The winner cancels only the loser — never itself — so that
+	// values tied to the winning context (e.g. an io.ReadCloser backed by an
+	// HTTP connection) remain usable after Do returns. The winner's context
+	// lives until the parent ctx is cancelled.
+	ctx1, cancel1 := context.WithCancel(ctx)
+	ctx2, cancel2 := context.WithCancel(ctx)
+
 	var (
 		ret    T
 		err    error
 		failed uint64
 
 		wg sync.WaitGroup
 		do sync.Once
+		// decided is closed by the goroutine that stores the winning result.
+		decided = make(chan struct{})
 	)
 
-	attempt := func(isRetry bool) {
+	attempt := func(attemptCtx context.Context, cancelSelf, cancelOther context.CancelFunc, isRetry bool) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
@@ -55,21 +68,43 @@ func (s Hedged[T]) Do(ctx context.Context) (T, error) {
 				// ongoing attempt.
 				return
 			}
-			// If there is an ongoing attempt, it will be cancelled,
-			// because we already got the result.
-			cancel()
+			stored := false
 			do.Do(func() {
 				ret, err = attemptRet, attemptErr
+				stored = true
+				close(decided)
 			})
+			if stored {
+				// We won: cancel the other attempt.
+				// Do NOT cancel our own context — the caller may still be
+				// reading from our result (e.g. io.ReadCloser).
+				cancelOther()
+			} else {
+				// We lost: cancel our own context and release our result.
+				cancelSelf()
+				if attemptErr == nil && s.Cleanup != nil {
+					s.Cleanup(attemptRet)
+				}
+			}
 		}()
 	}
 
-	attempt(false)
+	attempt(ctx1, cancel1, cancel2, false)
 	select {
-	case <-attemptCtx.Done():
-		// Call has returned, or caller cancelled the request.
+	case <-decided:
+		// A winner was found before the trigger fired.
+	case <-ctx.Done():
+		// Caller cancelled: abort both attempts.
+		cancel1()
+		cancel2()
 	case <-s.Trigger:
-		attempt(true)
+		attempt(ctx2, cancel2, cancel1, true)
+		select {
+		case <-decided:
+		case <-ctx.Done():
+			cancel1()
+			cancel2()
+		}
 	}
 
 	wg.Wait()