Fix queries and storage block operations synchronisation (#699)

* Fix ingest/flush race

* Allow concurrent reads in cutRowGroup

* profileStore synchronization

* alloc dedup slice buffer

* Fix block eviction synchronisation

Author: kolesnikovae

pkg/ingester/ingester.go

@@ -193,7 +193,7 @@ func (i *Ingester) forInstance(ctx context.Context, f func(*instance) error) err
 	return f(instance)
 }
 
-func (i *Ingester) evictBlock(tenantID string, b ulid.ULID, fn func() error) error {
+func (i *Ingester) evictBlock(tenantID string, b ulid.ULID, fn func() error) (err error) {
 	// We lock instances map for writes to ensure that no new instances are
 	// created during the procedure. Otherwise, during initialization, the
 	// new PhlareDB instance may try to load a block that has already been
@@ -204,12 +204,18 @@ func (i *Ingester) evictBlock(tenantID string, b ulid.ULID, fn func() error) err
 	// the process start, therefore there is no guarantee that we will find the
 	// discovered candidate block there. If it is the case, we have to ensure that
 	// the block won't be accessed, before and during deleting it from the disk.
+	var evicted bool
 	if tenantInstance, ok := i.instances[tenantID]; ok {
-		if _, err := tenantInstance.Evict(b); err != nil {
+		if evicted, err = tenantInstance.Evict(b, fn); err != nil {
 			return fmt.Errorf("failed to evict block %s/%s: %w", tenantID, b, err)
 		}
 	}
-	return fn()
+	// If the instance is not found, or the querier is not aware of the block,
+	// and thus the callback has not been invoked, do it now.
+	if !evicted {
+		return fn()
+	}
+	return nil
 }
 
 func (i *Ingester) Push(ctx context.Context, req *connect.Request[pushv1.PushRequest]) (*connect.Response[pushv1.PushResponse], error) {
@@ -222,7 +228,7 @@ func (i *Ingester) Push(ctx context.Context, req *connect.Request[pushv1.PushReq
 					if err != nil {
 						return err
 					}
-					if err = instance.Head().Ingest(ctx, p, id, series.Labels...); err != nil {
+					if err = instance.Ingest(ctx, p, id, series.Labels...); err != nil {
 						reason := validation.ReasonOf(err)
 						if reason != validation.Unknown {
 							validation.DiscardedProfiles.WithLabelValues(string(reason), instance.tenantID).Add(float64(1))

pkg/ingester/query.go

@@ -12,28 +12,28 @@ import (
 // LabelValues returns the possible label values for a given label name.
 func (i *Ingester) LabelValues(ctx context.Context, req *connect.Request[typesv1.LabelValuesRequest]) (*connect.Response[typesv1.LabelValuesResponse], error) {
 	return forInstanceUnary(ctx, i, func(instance *instance) (*connect.Response[typesv1.LabelValuesResponse], error) {
-		return instance.Head().LabelValues(ctx, req)
+		return instance.LabelValues(ctx, req)
 	})
 }
 
 // LabelNames returns the possible label names.
 func (i *Ingester) LabelNames(ctx context.Context, req *connect.Request[typesv1.LabelNamesRequest]) (*connect.Response[typesv1.LabelNamesResponse], error) {
 	return forInstanceUnary(ctx, i, func(instance *instance) (*connect.Response[typesv1.LabelNamesResponse], error) {
-		return instance.Head().LabelNames(ctx, req)
+		return instance.LabelNames(ctx, req)
 	})
 }
 
 // ProfileTypes returns the possible profile types.
 func (i *Ingester) ProfileTypes(ctx context.Context, req *connect.Request[ingestv1.ProfileTypesRequest]) (*connect.Response[ingestv1.ProfileTypesResponse], error) {
 	return forInstanceUnary(ctx, i, func(instance *instance) (*connect.Response[ingestv1.ProfileTypesResponse], error) {
-		return instance.Head().ProfileTypes(ctx, req)
+		return instance.ProfileTypes(ctx, req)
 	})
 }
 
 // Series returns labels series for the given set of matchers.
 func (i *Ingester) Series(ctx context.Context, req *connect.Request[ingestv1.SeriesRequest]) (*connect.Response[ingestv1.SeriesResponse], error) {
 	return forInstanceUnary(ctx, i, func(instance *instance) (*connect.Response[ingestv1.SeriesResponse], error) {
-		return instance.Head().Series(ctx, req)
+		return instance.Series(ctx, req)
 	})
 }
 

pkg/ingester/retention.go

@@ -76,9 +76,10 @@ func (*realFileSystem) RemoveAll(path string) error                { return os.R
 
 // blockEvicter unloads blocks from tenant instance.
 type blockEvicter interface {
-	// evictBlock evicts the block by its ID for the given tenant and invokes
-	// fn callback, if the tenant is found. The call is thread-safe: tenant
-	// can't be added or removed during the execution.
+	// evictBlock evicts the block by its ID from the memory and
+	// invokes fn callback, regardless of if the tenant is found.
+	// The call is thread-safe: tenant can't be added or removed
+	// during the execution.
 	evictBlock(tenant string, b ulid.ULID, fn func() error) error
 }
 

pkg/phlaredb/block_querier.go

@@ -225,10 +225,25 @@ func (b *BlockQuerier) Sync(ctx context.Context) error {
 	return nil
 }
 
+func (b *BlockQuerier) AddBlockQuerierByMeta(m *block.Meta) {
+	q := newSingleBlockQuerierFromMeta(b.phlarectx, b.bucketReader, m)
+	b.queriersLock.Lock()
+	defer b.queriersLock.Unlock()
+	i := sort.Search(len(b.queriers), func(i int) bool {
+		return b.queriers[i].meta.MinTime >= m.MinTime
+	})
+	if i < len(b.queriers) && b.queriers[i].meta.ULID == m.ULID {
+		// Block with this meta is already present, skipping.
+		return
+	}
+	b.queriers = append(b.queriers, q) // Ensure we have enough capacity.
+	copy(b.queriers[i+1:], b.queriers[i:])
+	b.queriers[i] = q
+}
+
 // evict removes the block with the given ULID from the querier.
 func (b *BlockQuerier) evict(blockID ulid.ULID) (bool, error) {
 	b.queriersLock.Lock()
-	defer b.queriersLock.Unlock()
 	// N.B: queriers are sorted by meta.MinTime.
 	j := -1
 	for i, q := range b.queriers {
@@ -238,17 +253,16 @@ func (b *BlockQuerier) evict(blockID ulid.ULID) (bool, error) {
 		}
 	}
 	if j < 0 {
+		b.queriersLock.Unlock()
 		return false, nil
 	}
 	blockQuerier := b.queriers[j]
-	if err := blockQuerier.Close(); err != nil {
-		return true, err
-	}
 	// Delete the querier from the slice and make it eligible for GC.
 	copy(b.queriers[j:], b.queriers[j+1:])
 	b.queriers[len(b.queriers)-1] = nil
 	b.queriers = b.queriers[:len(b.queriers)-1]
-	return true, nil
+	b.queriersLock.Unlock()
+	return true, blockQuerier.Close()
 }
 
 func (b *BlockQuerier) Close() error {

pkg/phlaredb/deduplicating_slice.go

@@ -42,7 +42,6 @@ type deduplicatingSlice[M Models, K comparable, H Helper[M, K], P schemav1.Persi
 	metrics *headMetrics
 	writer  *parquet.GenericWriter[P]
 
-	buffer      *parquet.Buffer
 	rowsFlushed int
 }
 
@@ -110,17 +109,14 @@ func (s *deduplicatingSlice[M, K, H, P]) maxRowsPerRowGroup() int {
 }
 
 func (s *deduplicatingSlice[M, K, H, P]) Flush(ctx context.Context) (numRows uint64, numRowGroups uint64, err error) {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
-	// intialise buffer if not existing
-	if s.buffer == nil {
-		s.buffer = parquet.NewBuffer(
-			s.persister.Schema(),
-			parquet.SortingRowGroupConfig(s.persister.SortingColumns()),
-			parquet.ColumnBufferCapacity(s.cfg.MaxBufferRowCount),
-		)
-	}
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+
+	buffer := parquet.NewBuffer(
+		s.persister.Schema(),
+		parquet.SortingRowGroupConfig(s.persister.SortingColumns()),
+		parquet.ColumnBufferCapacity(s.cfg.MaxBufferRowCount),
+	)
 
 	var (
 		maxRows = s.maxRowsPerRowGroup()
@@ -153,14 +149,14 @@ func (s *deduplicatingSlice[M, K, H, P]) Flush(ctx context.Context) (numRows uin
 			rows[pos] = s.persister.Deconstruct(rows[pos], uint64(slicePos), s.slice[slicePos])
 		}
 
-		s.buffer.Reset()
-		if _, err := s.buffer.WriteRows(rows); err != nil {
+		buffer.Reset()
+		if _, err = buffer.WriteRows(rows); err != nil {
 			return 0, 0, err
 		}
 
-		sort.Sort(s.buffer)
+		sort.Sort(buffer)
 
-		if _, err := s.writer.WriteRowGroup(s.buffer); err != nil {
+		if _, err = s.writer.WriteRowGroup(buffer); err != nil {
 			return 0, 0, err
 		}
 

pkg/phlaredb/head.go

@@ -127,9 +127,9 @@ type Head struct {
 	headPath  string // path while block is actively appended to
 	localPath string // path once block has been cut
 
-	flushCh chan struct{} // this channel is closed once the Head should be flushed, should be used externally
-
-	flushForcedTimer *time.Timer // this timer will phlare after the maximum
+	inFlightProfiles sync.WaitGroup // ongoing ingestion requests.
+	flushCh          chan struct{}  // this channel is closed once the Head should be flushed, should be used externally
+	flushForcedTimer *time.Timer    // this timer will phlare after the maximum
 
 	metaLock sync.RWMutex
 	meta     *block.Meta
@@ -569,8 +569,8 @@ func (h *Head) InRange(start, end model.Time) bool {
 
 // Returns underlying queries, the queriers should be roughly ordered in TS increasing order
 func (h *Head) Queriers() Queriers {
-	h.profiles.lock.RLock()
-	defer h.profiles.lock.RUnlock()
+	h.profiles.rowsLock.RLock()
+	defer h.profiles.rowsLock.RUnlock()
 
 	queriers := make([]Querier, 0, len(h.profiles.rowGroups)+1)
 	for idx := range h.profiles.rowGroups {
@@ -865,7 +865,9 @@ func (h *Head) Close() error {
 	return merr.Err()
 }
 
-// Flush closes the head and writes data to disk
+// Flush closes the head and writes data to disk. No ingestion requests should
+// be made concurrently with the call, or after it returns.
+// The call is thread-safe for reads.
 func (h *Head) Flush(ctx context.Context) error {
 	start := time.Now()
 	defer func() {
@@ -880,7 +882,11 @@ func (h *Head) Flush(ctx context.Context) error {
 }
 
 func (h *Head) flush(ctx context.Context) error {
-	if h.profiles.empty() {
+	// Ensure all the in-flight ingestion requests have finished.
+	// It must be guaranteed that no new inserts will happen
+	// after the call start.
+	h.inFlightProfiles.Wait()
+	if len(h.profiles.slice) == 0 {
 		level.Info(h.logger).Log("msg", "head empty - no block written")
 		return os.RemoveAll(h.headPath)
 	}
@@ -941,6 +947,20 @@ func (h *Head) flush(ctx context.Context) error {
 		return err
 	}
 	h.metrics.blockDurationSeconds.Observe(h.meta.MaxTime.Sub(h.meta.MinTime).Seconds())
+	return nil
+}
+
+// Move moves the head directory to local blocks. The call is not thread-safe:
+// no concurrent reads and writes are allowed.
+//
+// After the call, head in-memory representation is not valid and should not
+// be accessed for querying.
+func (h *Head) Move() error {
+	// Remove intermediate row groups before the move as they are still
+	// referencing files on the disk.
+	if err := h.profiles.DeleteRowGroups(); err != nil {
+		return err
+	}
 
 	// move block to the local directory
 	if err := os.MkdirAll(filepath.Dir(h.localPath), defaultFolderMode); err != nil {

pkg/phlaredb/head_queriers.go

@@ -24,8 +24,8 @@ type headOnDiskQuerier struct {
 }
 
 func (q *headOnDiskQuerier) rowGroup() *rowGroupOnDisk {
-	q.head.profiles.lock.RLock()
-	defer q.head.profiles.lock.RUnlock()
+	q.head.profiles.rowsLock.RLock()
+	defer q.head.profiles.rowsLock.RUnlock()
 	return q.head.profiles.rowGroups[q.rowGroupIdx]
 }
 
@@ -119,7 +119,7 @@ func (q *headOnDiskQuerier) MergePprof(ctx context.Context, rows iter.Iterator[P
 
 	stacktraceSamples := profileSampleMap{}
 
-	if err := mergeByStacktraces(ctx, q.head.profiles.rowGroups[q.rowGroupIdx], rows, stacktraceSamples); err != nil {
+	if err := mergeByStacktraces(ctx, q.rowGroup(), rows, stacktraceSamples); err != nil {
 		return nil, err
 	}
 
@@ -132,7 +132,7 @@ func (q *headOnDiskQuerier) MergeByLabels(ctx context.Context, rows iter.Iterato
 
 	seriesByLabels := make(seriesByLabels)
 
-	if err := mergeByLabels(ctx, q.head.profiles.rowGroups[q.rowGroupIdx], rows, seriesByLabels, by...); err != nil {
+	if err := mergeByLabels(ctx, q.rowGroup(), rows, seriesByLabels, by...); err != nil {
 		return nil, err
 	}