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,6 +26,11 @@ 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)
@@ -58,9 +63,14 @@ func (s Hedged[T]) Do(ctx context.Context) (T, error) {
 			// 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
 			})
+			if !stored && attemptErr == nil && s.Cleanup != nil {
+				s.Cleanup(attemptRet)
+			}
 		}()
 	}
 

pkg/util/retry/hedged_test.go

@@ -3,6 +3,7 @@ package retry
 import (
 	"context"
 	"errors"
+	"sync/atomic"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -135,3 +136,105 @@ func Test_Hedging(t *testing.T) {
 		})
 	}
 }
+
+func Test_Hedging_Cleanup(t *testing.T) {
+	t.Run("cleanup called on loser when both succeed", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			var cleaned int64
+			// First call is slow; hedge fires and wins. First call then
+			// succeeds but loses — Cleanup must be called on its result.
+			const hedgeDelay = time.Second
+			a := Hedged[*int]{
+				Trigger:  time.After(hedgeDelay),
+				FailFast: true,
+				Cleanup:  func(v *int) { atomic.AddInt64(&cleaned, 1) },
+				Call: func(ctx context.Context, isRetry bool) (*int, error) {
+					if !isRetry {
+						// slow: block until cancelled by the winning hedge
+						<-ctx.Done()
+					}
+					v := 1
+					return &v, nil
+				},
+			}
+			done := make(chan struct{})
+			go func() {
+				defer close(done)
+				_, err := a.Do(context.Background())
+				if err != nil {
+					t.Errorf("unexpected error: %v", err)
+				}
+			}()
+			synctest.Wait()
+			time.Sleep(hedgeDelay) // hedge fires and wins; slow call gets cancelled
+			synctest.Wait()
+			<-done
+			if atomic.LoadInt64(&cleaned) != 1 {
+				t.Fatal("expected Cleanup to be called exactly once on the loser")
+			}
+		})
+	})
+
+	t.Run("cleanup not called when only one attempt runs", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			var cleaned int64
+			a := Hedged[*int]{
+				Trigger:  time.After(time.Hour),
+				FailFast: true,
+				Cleanup:  func(v *int) { atomic.AddInt64(&cleaned, 1) },
+				Call: func(ctx context.Context, _ bool) (*int, error) {
+					v := 1
+					return &v, nil
+				},
+			}
+			_, err := a.Do(context.Background())
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if atomic.LoadInt64(&cleaned) != 0 {
+				t.Fatal("expected Cleanup not to be called")
+			}
+		})
+	})
+
+	t.Run("cleanup not called on loser that errored", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			var cleaned int64
+			e := errors.New("fail")
+			const hedgeDelay = time.Second
+			// Hedge fires; first call errors, second succeeds.
+			// Cleanup must NOT be called on the errored loser.
+			firstBlocked := make(chan struct{})
+			a := Hedged[*int]{
+				Trigger:  time.After(hedgeDelay),
+				FailFast: false,
+				Cleanup:  func(v *int) { atomic.AddInt64(&cleaned, 1) },
+				Call: func(ctx context.Context, isRetry bool) (*int, error) {
+					if !isRetry {
+						close(firstBlocked)
+						<-ctx.Done()
+						return nil, e
+					}
+					v := 1
+					return &v, nil
+				},
+			}
+			done := make(chan struct{})
+			go func() {
+				defer close(done)
+				_, err := a.Do(context.Background())
+				if err != nil {
+					t.Errorf("unexpected error: %v", err)
+				}
+			}()
+			<-firstBlocked
+			synctest.Wait()
+			time.Sleep(hedgeDelay)
+			synctest.Wait()
+			<-done
+			if atomic.LoadInt64(&cleaned) != 0 {
+				t.Fatal("expected Cleanup not to be called on errored attempt")
+			}
+		})
+	})
+}