Implement the merge API in the query path

pkg/phlaredb/block_querier.go

@@ -372,6 +372,7 @@ func (b *singleBlockQuerier) Close() error {
 			errs.Add(err)
 		}
 	}
+	b.opened = false
 	return errs.Err()
 }
 
@@ -483,16 +484,32 @@ type Querier interface {
 	MergeByStacktraces(ctx context.Context, rows iter.Iterator[Profile]) (*ingestv1.MergeProfilesStacktracesResult, error)
 	MergeByLabels(ctx context.Context, rows iter.Iterator[Profile], by ...string) ([]*typesv1.Series, error)
 	MergePprof(ctx context.Context, rows iter.Iterator[Profile]) (*profile.Profile, error)
-
+	Open(ctx context.Context) error
 	// Sorts profiles for retrieval.
 	Sort([]Profile) []Profile
 }
 
 type Queriers []Querier
 
+func (queriers Queriers) Open(ctx context.Context) error {
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(128)
+	for _, q := range queriers {
+		q := q
+		g.Go(func() error {
+			if err := q.Open(ctx); err != nil {
+				return err
+			}
+			return nil
+		})
+	}
+	return g.Wait()
+}
+
 func (queriers Queriers) SelectMatchingProfiles(ctx context.Context, params *ingestv1.SelectProfilesRequest) (iter.Iterator[Profile], error) {
 	iters := make([]iter.Iterator[Profile], 0, len(queriers))
-
+	// todo: if blocks overlap in time we should be deduping
+	// otherwise we can just append
 	for _, q := range queriers {
 		it, err := q.SelectMatchingProfiles(ctx, params)
 		if err != nil {
@@ -504,17 +521,19 @@ func (queriers Queriers) SelectMatchingProfiles(ctx context.Context, params *ing
 	return iter.NewSortProfileIterator(iters), nil
 }
 
-func (queriers Queriers) ForTimeRange(start, end model.Time) Queriers {
+func (queriers Queriers) ForTimeRange(_ context.Context, start, end model.Time) (Queriers, error) {
 	result := make(Queriers, 0, len(queriers))
 	for _, q := range queriers {
 		if q.InRange(start, end) {
 			result = append(result, q)
 		}
 	}
-	return result
+	return result, nil
 }
 
-func (q Queriers) MergeProfilesStacktraces(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesStacktracesRequest, ingestv1.MergeProfilesStacktracesResponse]) error {
+type BlockGetter func(ctx context.Context, start, end model.Time) (Queriers, error)
+
+func MergeProfilesStacktraces(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesStacktracesRequest, ingestv1.MergeProfilesStacktracesResponse], blockGetter BlockGetter) error {
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "MergeProfilesStacktraces")
 	defer sp.Finish()
 
@@ -537,7 +556,10 @@ func (q Queriers) MergeProfilesStacktraces(ctx context.Context, stream *connect.
 		otlog.String("profile_id", request.Type.ID),
 	)
 
-	queriers := q.ForTimeRange(model.Time(request.Start), model.Time(request.End))
+	queriers, err := blockGetter(ctx, model.Time(request.Start), model.Time(request.End))
+	if err != nil {
+		return err
+	}
 
 	result := make([]*ingestv1.MergeProfilesStacktracesResult, 0, len(queriers))
 	var lock sync.Mutex
@@ -600,7 +622,7 @@ func (q Queriers) MergeProfilesStacktraces(ctx context.Context, stream *connect.
 	return nil
 }
 
-func (q Queriers) MergeProfilesLabels(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesLabelsRequest, ingestv1.MergeProfilesLabelsResponse]) error {
+func MergeProfilesLabels(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesLabelsRequest, ingestv1.MergeProfilesLabelsResponse], blockGetter BlockGetter) error {
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "MergeProfilesLabels")
 	defer sp.Finish()
 
@@ -626,7 +648,10 @@ func (q Queriers) MergeProfilesLabels(ctx context.Context, stream *connect.BidiS
 		otlog.String("by", strings.Join(by, ",")),
 	)
 
-	queriers := q.ForTimeRange(model.Time(request.Start), model.Time(request.End))
+	queriers, err := blockGetter(ctx, model.Time(request.Start), model.Time(request.End))
+	if err != nil {
+		return err
+	}
 	result := make([][]*typesv1.Series, 0, len(queriers))
 	g, ctx := errgroup.WithContext(ctx)
 	s := lo.Synchronize()
@@ -686,7 +711,7 @@ func (q Queriers) MergeProfilesLabels(ctx context.Context, stream *connect.BidiS
 	return nil
 }
 
-func (q Queriers) MergeProfilesPprof(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesPprofRequest, ingestv1.MergeProfilesPprofResponse]) error {
+func MergeProfilesPprof(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesPprofRequest, ingestv1.MergeProfilesPprofResponse], blockGetter BlockGetter) error {
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "MergeProfilesPprof")
 	defer sp.Finish()
 
@@ -709,7 +734,10 @@ func (q Queriers) MergeProfilesPprof(ctx context.Context, stream *connect.BidiSt
 		otlog.String("profile_id", request.Type.ID),
 	)
 
-	queriers := q.ForTimeRange(model.Time(request.Start), model.Time(request.End))
+	queriers, err := blockGetter(ctx, model.Time(request.Start), model.Time(request.End))
+	if err != nil {
+		return err
+	}
 
 	result := make([]*profile.Profile, 0, len(queriers))
 	var lock sync.Mutex
@@ -1079,6 +1107,8 @@ func (r *parquetReader[M, P]) Close() error {
 	if r.reader != nil {
 		return r.reader.Close()
 	}
+	r.reader = nil
+	r.file = nil
 	return nil
 }
 
@@ -1193,6 +1223,9 @@ func (r *inMemoryparquetReader[M, P]) Close() error {
 	if r.reader != nil {
 		return r.reader.Close()
 	}
+	r.reader = nil
+	r.file = nil
+	r.cache = nil
 	return nil
 }
 

pkg/phlaredb/head_queriers.go

@@ -29,6 +29,10 @@ func (q *headOnDiskQuerier) rowGroup() *rowGroupOnDisk {
 	return q.head.profiles.rowGroups[q.rowGroupIdx]
 }
 
+func (q *headOnDiskQuerier) Open(_ context.Context) error {
+	return nil
+}
+
 func (q *headOnDiskQuerier) SelectMatchingProfiles(ctx context.Context, params *ingestv1.SelectProfilesRequest) (iter.Iterator[Profile], error) {
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "SelectMatchingProfiles - HeadOnDisk")
 	defer sp.Finish()
@@ -157,6 +161,10 @@ type headInMemoryQuerier struct {
 	head *Head
 }
 
+func (q *headInMemoryQuerier) Open(_ context.Context) error {
+	return nil
+}
+
 func (q *headInMemoryQuerier) SelectMatchingProfiles(ctx context.Context, params *ingestv1.SelectProfilesRequest) (iter.Iterator[Profile], error) {
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "SelectMatchingProfiles - HeadInMemory")
 	defer sp.Finish()
@@ -184,7 +192,7 @@ func (q *headInMemoryQuerier) SelectMatchingProfiles(ctx context.Context, params
 			continue
 		}
 
-		var profiles = make([]*schemav1.Profile, len(profileSeries.profiles))
+		profiles := make([]*schemav1.Profile, len(profileSeries.profiles))
 		copy(profiles, profileSeries.profiles)
 
 		iters = append(iters,
@@ -259,7 +267,6 @@ func (q *headInMemoryQuerier) MergePprof(ctx context.Context, rows iter.Iterator
 	}
 
 	return q.head.resolvePprof(ctx, stacktraceSamples), nil
-
 }
 
 func (q *headInMemoryQuerier) MergeByLabels(ctx context.Context, rows iter.Iterator[Profile], by ...string) ([]*typesv1.Series, error) {

pkg/phlaredb/phlaredb.go

@@ -190,15 +190,15 @@ func (f *PhlareDB) Queriers() Queriers {
 }
 
 func (f *PhlareDB) MergeProfilesStacktraces(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesStacktracesRequest, ingestv1.MergeProfilesStacktracesResponse]) error {
-	return f.Queriers().MergeProfilesStacktraces(ctx, stream)
+	return MergeProfilesStacktraces(ctx, stream, f.blockQuerier.Queriers().ForTimeRange)
 }
 
 func (f *PhlareDB) MergeProfilesLabels(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesLabelsRequest, ingestv1.MergeProfilesLabelsResponse]) error {
-	return f.Queriers().MergeProfilesLabels(ctx, stream)
+	return MergeProfilesLabels(ctx, stream, f.blockQuerier.Queriers().ForTimeRange)
 }
 
 func (f *PhlareDB) MergeProfilesPprof(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesPprofRequest, ingestv1.MergeProfilesPprofResponse]) error {
-	return f.Queriers().MergeProfilesPprof(ctx, stream)
+	return MergeProfilesPprof(ctx, stream, f.blockQuerier.Queriers().ForTimeRange)
 }
 
 type BidiServerMerge[Res any, Req any] interface {

pkg/phlaredb/phlaredb_test.go

@@ -98,6 +98,18 @@ type ingesterHandlerPhlareDB struct {
 	// *PhlareDB
 }
 
+func (i *ingesterHandlerPhlareDB) MergeProfilesStacktraces(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesStacktracesRequest, ingestv1.MergeProfilesStacktracesResponse]) error {
+	return MergeProfilesStacktraces(ctx, stream, i.ForTimeRange)
+}
+
+func (i *ingesterHandlerPhlareDB) MergeProfilesLabels(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesLabelsRequest, ingestv1.MergeProfilesLabelsResponse]) error {
+	return MergeProfilesLabels(ctx, stream, i.ForTimeRange)
+}
+
+func (i *ingesterHandlerPhlareDB) MergeProfilesPprof(ctx context.Context, stream *connect.BidiStream[ingestv1.MergeProfilesPprofRequest, ingestv1.MergeProfilesPprofResponse]) error {
+	return MergeProfilesPprof(ctx, stream, i.ForTimeRange)
+}
+
 func (i *ingesterHandlerPhlareDB) Push(context.Context, *connect.Request[pushv1.PushRequest]) (*connect.Response[pushv1.PushResponse], error) {
 	return nil, errors.New("not implemented")
 }

pkg/storegateway/block.go

@@ -0,0 +1,69 @@
+package storegateway
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/pkg/errors"
+	"github.com/prometheus/common/model"
+
+	"github.com/grafana/phlare/pkg/phlaredb"
+	"github.com/grafana/phlare/pkg/phlaredb/block"
+	"github.com/grafana/phlare/pkg/util"
+)
+
+type BlockCloser interface {
+	phlaredb.Querier
+	Close() error
+}
+
+type Block struct {
+	BlockCloser
+	meta   *block.Meta
+	logger log.Logger
+}
+
+func (bs *BucketStore) createBlock(ctx context.Context, meta *block.Meta) (*Block, error) {
+	blockLocalPath := filepath.Join(bs.syncDir, meta.ULID.String())
+	// add the dir if it doesn't exist
+	if _, err := os.Stat(blockLocalPath); errors.Is(err, os.ErrNotExist) {
+		if err := os.MkdirAll(bs.syncDir, 0o750); err != nil {
+			return nil, errors.Wrap(err, "create dir")
+		}
+	}
+	metaPath := filepath.Join(bs.syncDir, block.MetaFilename)
+	if _, err := os.Stat(metaPath); errors.Is(err, os.ErrNotExist) {
+		// add meta.json if it does not exist
+		if _, err := meta.WriteToFile(bs.logger, bs.syncDir); err != nil {
+			return nil, errors.Wrap(err, "write meta.json")
+		}
+	} else {
+		// read meta.json if it exists and validate it
+		if diskMeta, _, err := block.MetaFromDir(bs.syncDir); err != nil {
+			if meta.String() != diskMeta.String() {
+				return nil, errors.Wrap(err, "meta.json does not match")
+			}
+			return nil, errors.Wrap(err, "read meta.json")
+		}
+	}
+
+	blk := phlaredb.NewSingleBlockQuerierFromMeta(ctx, bs.bucket, meta)
+	// Load the block into memory if it's within the last 24 hours.
+	// Todo make this configurable
+	if blk.InRange(model.Now().Add(-24*time.Hour), model.Now()) {
+		go func() {
+			if err := blk.Open(ctx); err != nil {
+				level.Error(util.Logger).Log("msg", "open block", "err", err)
+			}
+		}()
+	}
+	return &Block{
+		meta:        meta,
+		logger:      bs.logger,
+		BlockCloser: blk,
+	}, nil
+}