syncer.go

// Copyright (c) The Thanos Authors.
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
//     https://opensource.org/licenses/Apache-2.0

package locate

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"sync"

	"github.com/alecthomas/units"
	"github.com/efficientgo/core/errcapture"
	"github.com/parquet-go/parquet-go"
	"github.com/prometheus/prometheus/tsdb/fileutil"
	"github.com/thanos-io/objstore"

	"github.com/thanos-io/thanos-parquet-gateway/db"
	"github.com/thanos-io/thanos-parquet-gateway/internal/util"
	"github.com/thanos-io/thanos-parquet-gateway/schema"
)

type Syncer struct {
	bkt objstore.Bucket

	blockOpts        []BlockOption
	metaFilter       MetaFilter
	concurrency      int
	shardConcurrency int

	mu     sync.Mutex
	blocks map[schema.ExternalLabelsHash]map[util.Date]*db.Block

	cached []*db.Block
}

type SyncerOption func(*syncerConfig)

type syncerConfig struct {
	blockOpts        []BlockOption
	metaFilter       MetaFilter
	concurrency      int
	shardConcurrency int
}

func BlockOptions(opts ...BlockOption) SyncerOption {
	return func(cfg *syncerConfig) {
		cfg.blockOpts = opts
	}
}

func FilterMetas(f MetaFilter) SyncerOption {
	return func(cfg *syncerConfig) {
		cfg.metaFilter = f
	}
}

func BlockConcurrency(c int) SyncerOption {
	return func(cfg *syncerConfig) {
		cfg.concurrency = c
	}
}
func ShardConcurrency(c int) SyncerOption {
	return func(cfg *syncerConfig) {
		cfg.shardConcurrency = c
	}
}

type BlockOption func(*blockConfig)

type blockConfig struct {
	readBufferSize units.Base2Bytes
	labelFilesDir  string
}

func ReadBufferSize(sz units.Base2Bytes) BlockOption {
	return func(cfg *blockConfig) {
		cfg.readBufferSize = sz
	}
}

func LabelFilesDir(d string) BlockOption {
	return func(cfg *blockConfig) {
		cfg.labelFilesDir = d
	}
}

func NewSyncer(bkt objstore.Bucket, opts ...SyncerOption) *Syncer {
	cfg := syncerConfig{
		metaFilter:       AllMetasMetaFilter,
		concurrency:      1,
		shardConcurrency: 0,
	}

	for _, o := range opts {
		o(&cfg)
	}

	return &Syncer{
		bkt:              bkt,
		blocks:           make(map[schema.ExternalLabelsHash]map[util.Date]*db.Block),
		blockOpts:        cfg.blockOpts,
		metaFilter:       cfg.metaFilter,
		concurrency:      cfg.concurrency,
		shardConcurrency: cfg.shardConcurrency,
	}
}

func (s *Syncer) Blocks() []*db.Block {
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.filterBlocks(s.cached)
}

func (s *Syncer) Sync(ctx context.Context, parquetStreams map[schema.ExternalLabelsHash]schema.ParquetBlocksStream) error {
	type blockOrError struct {
		blk       *db.Block
		extLabels schema.ExternalLabels
		err       error
		blockName string
	}

	type metaHashTuple struct {
		streamHash schema.ExternalLabelsHash
		m          schema.Meta
		extLabels  schema.ExternalLabels
	}

	blkC := make(chan blockOrError)
	go func() {
		defer close(blkC)

		workerC := make(chan metaHashTuple, s.concurrency)
		go func() {
			defer close(workerC)

			for streamHash, stream := range parquetStreams {
				if _, ok := s.blocks[streamHash]; !ok {
					s.blocks[streamHash] = make(map[util.Date]*db.Block)
				}
				for _, m := range s.filterMetas(stream.Metas) {
					if _, ok := s.blocks[streamHash][m.Date]; ok {
						continue
					}
					workerC <- metaHashTuple{
						m:         m,
						extLabels: stream.ExternalLabels,
					}
				}
			}
		}()

		var wg sync.WaitGroup
		defer wg.Wait()

		for range s.concurrency {
			wg.Add(1)
			go func() {
				defer wg.Done()
				for mh := range workerC {
					blk, err := newBlockForMeta(ctx, s.bkt, mh.extLabels, mh.m, s.shardConcurrency, s.blockOpts...)
					if err != nil {
						blkC <- blockOrError{blk: nil, err: fmt.Errorf("unable to read block %s %q: %w", mh.streamHash.String(), mh.m.Date, err), blockName: mh.streamHash.String()}
					} else {
						blkC <- blockOrError{blk: blk, extLabels: mh.extLabels}
					}
				}
			}()
		}
	}()

	blocks := make(map[schema.ExternalLabelsHash]map[util.Date]*db.Block, 0)
	for b := range blkC {
		if b.err != nil {
			return fmt.Errorf("unable to read block: %w", b.err)
		}
		if _, ok := blocks[b.extLabels.Hash()]; !ok {
			blocks[b.extLabels.Hash()] = make(map[util.Date]*db.Block)
		}
		blocks[b.extLabels.Hash()][b.blk.Meta().Date] = b.blk
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	var numBlocks int
	for stream := range s.blocks {
		if _, ok := parquetStreams[stream]; !ok {
			delete(s.blocks, stream)
		}
	}

	for stream := range s.blocks {
		if _, ok := blocks[stream]; !ok {
			continue
		}
		s.blocks[stream] = blocks[stream]
	}
	for _, v := range s.blocks {
		numBlocks += len(v)
	}
	allBlocks := make([]*db.Block, 0, numBlocks)
	for _, m := range s.blocks {
		for _, b := range m {
			allBlocks = append(allBlocks, b)
		}
	}

	s.cached = allBlocks
	sort.Slice(s.cached, func(i, j int) bool {
		ls, _ := s.cached[i].Timerange()
		rs, _ := s.cached[j].Timerange()
		return ls < rs
	})

	if len(s.cached) != 0 {
		syncMinTime.WithLabelValues(whatSyncer).Set(float64(s.cached[0].Meta().Mint))
		syncMaxTime.WithLabelValues(whatSyncer).Set(float64(s.cached[len(s.cached)-1].Meta().Maxt))
	}
	syncLastSuccessfulTime.WithLabelValues(whatSyncer).SetToCurrentTime()

	return nil
}

func (s *Syncer) filterMetas(metas []schema.Meta) []schema.Meta {
	return s.metaFilter.filterMetas(metas)
}

func (s *Syncer) filterBlocks(blks []*db.Block) []*db.Block {
	return s.metaFilter.filterBlocks(blks)
}

/*
TODO: the following functions should be abstracted into syncer somehow, the syncer should decide
where shards live (s3, disk, memory) and should also do regular maintenance on its on disk state,
i.e. downloading files that it wants on disk, deleting files that are out of retention, etc.
*/

func newBlockForMeta(ctx context.Context, bkt objstore.Bucket, extLabels schema.ExternalLabels, m schema.Meta, shardConcurrency int, opts ...BlockOption) (*db.Block, error) {
	cfg := blockConfig{
		readBufferSize: 8 * units.MiB,
		labelFilesDir:  os.TempDir(),
	}
	for _, o := range opts {
		o(&cfg)
	}

	var shards []*db.Shard
	var err error

	// Use parallel loading if shardConcurrency > 0, otherwise use sequential (original) loading
	if shardConcurrency > 0 {
		shards, err = readShardsParallel(ctx, bkt, m, cfg, extLabels.Hash(), shardConcurrency)
	} else {
		shards, err = readShards(ctx, bkt, m, extLabels.Hash(), cfg)
	}
	if err != nil {
		return nil, fmt.Errorf("unable to read shards: %w", err)
	}

	return db.NewBlock(m, extLabels, shards...), nil
}

func readShards(ctx context.Context, bkt objstore.Bucket, m schema.Meta, extLabelsHash schema.ExternalLabelsHash, cfg blockConfig) ([]*db.Shard, error) {
	shards := make([]*db.Shard, 0, m.Shards)
	for sh := range int(m.Shards) {
		shard, err := readShard(ctx, bkt, m, extLabelsHash, sh, cfg)
		if err != nil {
			return nil, fmt.Errorf("unable to read shard %d: %w", sh, err)
		}
		shards = append(shards, shard)
	}
	return shards, nil
}

func readShardsParallel(ctx context.Context, bkt objstore.Bucket, m schema.Meta, cfg blockConfig, extLabelsHash schema.ExternalLabelsHash, shardConcurrency int) ([]*db.Shard, error) {
	type shardOrError struct {
		shard *db.Shard
		index int
		err   error
	}

	shardC := make(chan shardOrError, int(m.Shards))
	var wg sync.WaitGroup
	sem := make(chan struct{}, shardConcurrency)

	for i := range int(m.Shards) {
		wg.Add(1)
		go func(shardIndex int) {
			defer wg.Done()

			sem <- struct{}{}
			defer func() { <-sem }()

			shard, err := readShard(ctx, bkt, m, extLabelsHash, shardIndex, cfg)
			shardC <- shardOrError{shard: shard, index: shardIndex, err: err}
		}(i)
	}

	go func() {
		wg.Wait()
		close(shardC)
	}()

	shards := make([]*db.Shard, 0, int(m.Shards))
	var failedShards []int

	for result := range shardC {
		if result.err != nil {
			failedShards = append(failedShards, result.index)
			continue
		}
		shards = append(shards, result.shard)
	}

	// Return error only if ALL shards failed
	if len(shards) == 0 {
		return nil, fmt.Errorf("all shards failed to load, failed_shards=%v", failedShards)
	}

	return shards, nil
}

func bucketReaderFromContext(bkt objstore.Bucket, name string) func(context.Context) io.ReaderAt {
	return func(ctx context.Context) io.ReaderAt {
		return newBucketReaderAt(ctx, bkt, name)
	}
}

func readShard(ctx context.Context, bkt objstore.Bucket, m schema.Meta, extLabelsHash schema.ExternalLabelsHash, shard int, cfg blockConfig) (s *db.Shard, err error) {
	chunkspfile := schema.ChunksPfileNameForShard(extLabelsHash, m.Date, shard)
	attrs, err := bkt.Attributes(ctx, chunkspfile)
	if err != nil {
		return nil, fmt.Errorf("unable to attr chunks parquet file %q: %w", chunkspfile, err)
	}

	bktRdrAtFromCtx := bucketReaderFromContext(bkt, chunkspfile)

	chunkspf, err := parquet.OpenFile(bktRdrAtFromCtx(ctx), attrs.Size,
		parquet.FileReadMode(parquet.ReadModeAsync),
		parquet.ReadBufferSize(int(cfg.readBufferSize)),
		parquet.SkipMagicBytes(true),
		parquet.SkipBloomFilters(true),
		parquet.OptimisticRead(true),
	)
	if err != nil {
		return nil, fmt.Errorf("unable to open chunks parquet file %q: %w", chunkspfile, err)
	}

	labelspfile := schema.LabelsPfileNameForShard(extLabelsHash, m.Date, shard)
	labelspfilePath := filepath.Join(cfg.labelFilesDir, url.PathEscape(labelspfile))

	// If we were not able to read the file for any reason we delete it and retry.
	// This is also executed on paths that have nothing to do with the file, i.e. loading
	// its content from object storage, but just makes sure that we dont forget cleanup.
	defer func() {
		if err != nil {
			if cerr := os.RemoveAll(labelspfilePath); cerr != nil {
				err = errors.Join(err, fmt.Errorf("unable to remove labels parquet file %q: %w", labelspfilePath, cerr))
			}
		}
	}()

	// if the file was corrupted on its way to disk we remove it and will retry downloading it next try
	if stat, err := os.Stat(labelspfilePath); err != nil {
		if !os.IsNotExist(err) {
			return nil, fmt.Errorf("unable to stat label parquet file %q from disk: %w", labelspfile, err)
			// file didn't exist - we need to download and save it to disk
		}
	} else {
		f, err := fileutil.OpenMmapFileWithSize(labelspfilePath, int(stat.Size()))
		if err != nil {
			return nil, fmt.Errorf("unable to mmap label parquet file %q: %w", labelspfile, err)
		}
		labelspf, err := parquet.OpenFile(bytes.NewReader(f.Bytes()), stat.Size())
		if err != nil {
			syncCorruptedLabelFile.Add(1)
			rerr := fmt.Errorf("unable to read label parquet file %q: %w", labelspfile, err)
			if cerr := f.Close(); cerr != nil {
				return nil, errors.Join(rerr, fmt.Errorf("unable to close memory mapped parquet file %q: %w", labelspfile, cerr))
			}
			return nil, rerr
		}
		return db.NewShard(m, chunkspf, labelspf, bktRdrAtFromCtx), nil
	}
	rdr, err := bkt.Get(ctx, labelspfile)
	if err != nil {
		return nil, fmt.Errorf("unable to get %q: %w", labelspfile, err)
	}
	defer errcapture.Do(&err, rdr.Close, "labels parquet file close")

	f, err := os.Create(labelspfilePath)
	if err != nil {
		return nil, fmt.Errorf("unable to create label parquet file %q on disk: %w", labelspfile, err)
	}
	defer errcapture.Do(&err, f.Close, "labels parquet file close")

	n, err := io.Copy(f, rdr)
	if err != nil {
		return nil, fmt.Errorf("unable to copy label parquet file %q to disk: %w", labelspfile, err)
	}
	if err := f.Sync(); err != nil {
		return nil, fmt.Errorf("unable to close label parquet file %q: %w", labelspfile, err)
	}

	mf, err := fileutil.OpenMmapFileWithSize(labelspfilePath, int(n))
	if err != nil {
		return nil, fmt.Errorf("unable to mmap label parquet file %q: %w", labelspfile, err)
	}
	labelspf, err := parquet.OpenFile(bytes.NewReader(mf.Bytes()), int64(n))
	if err != nil {
		syncCorruptedLabelFile.Add(1)
		rerr := fmt.Errorf("unable to read label parquet file %q: %w", labelspfile, err)
		if cerr := f.Close(); cerr != nil {
			return nil, errors.Join(rerr, fmt.Errorf("unable to close memory mapped parquet file %q: %w", labelspfile, cerr))
		}
		return nil, rerr
	}
	return db.NewShard(m, chunkspf, labelspf, bktRdrAtFromCtx), nil
}