plan.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 convert

import (
	"cmp"
	"slices"
	"time"

	"github.com/thanos-io/thanos-parquet-gateway/internal/util"
	"github.com/thanos-io/thanos-parquet-gateway/schema"
	"github.com/thanos-io/thanos/pkg/block/metadata"
)

type Step struct {
	Date           util.Date       // Date for which we are building a parquet block.
	Sources        []metadata.Meta // Source TSDB blocks we will use to generate a parquet block.
	ExternalLabels schema.ExternalLabels
}

// isFullyCovered returns true if blocks for this date cover the whole day.
// We return true even if there are gaps in coverage, we only care that there
// are blocks that cover both min and max timestamps.
func (s Step) isFullyCovered() bool {
	mint := s.Date.MinT()
	maxt := s.Date.MaxT()
	var gotMin, gotMax bool
	for _, s := range s.Sources {
		if s.MinTime <= mint {
			gotMin = true
		}
		if s.MaxTime >= maxt {
			gotMax = true
		}
	}
	return gotMin && gotMax
}

type Plan struct {
	Steps []Step
}

type Planner struct {
	// Do not create parquet blocks that are younger then this.
	notAfter time.Time
	// Maximum number of days to plan conversions for. Planner might still produce a plan with more days
	// if it would be converting a TSDB block that spans over more days, to avoid re-downloading that block
	// for the next plan.
	maxDays int
}

func NewPlanner(notAfter time.Time, maxDays int) Planner {
	return Planner{notAfter: notAfter, maxDays: maxDays}
}

func (p Planner) planStream(tsdb schema.TSDBBlocksStream, parquet schema.ParquetBlocksStream, minTimeDate string, maxTimeDate string) Plan {

	var (
		windowActive bool
		winStartDay  time.Time
		winEndDay    time.Time
	)

	if minTimeDate != "" || maxTimeDate != "" {
		windowActive = true

		// Parse minDate and round to start of day (00:00:00 UTC, inclusive)
		if minTimeDate != "" {
			t, err := time.Parse("2006-01-02", minTimeDate)
			if err != nil {
				// Invalid date format, disable window filtering
				windowActive = false
			} else {
				winStartDay = time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
			}
		}

		// Parse maxDate and round to start of next day (00:00:00 UTC, exclusive)
		if maxTimeDate != "" {
			t, err := time.Parse("2006-01-02", maxTimeDate)
			if err != nil {
				// Invalid date format, disable window filtering
				windowActive = false
			} else {
				// Add 1 day to make the boundary exclusive
				winEndDay = time.Date(t.Year(), t.Month(), t.Day()+1, 0, 0, 0, 0, time.UTC)
			}
		}
	}
	// Make a list of days covered by TSDB blocks.
	tsdbDates := map[util.Date][]metadata.Meta{}
	for _, tsdb := range tsdb.Metas {
		for _, partialDate := range util.SplitIntoDates(tsdb.MinTime, tsdb.MaxTime) {
			if windowActive {
				pt := partialDate.ToTime()
				if minTimeDate != "" && pt.Before(winStartDay) {
					continue
				}
				if maxTimeDate != "" && !pt.Before(winEndDay) {
					continue
				}
			}
			tsdbDates[partialDate] = append(tsdbDates[partialDate], tsdb)
		}
	}

	pqDates := map[util.Date]struct{}{}
	for _, pq := range parquet.Metas {
		for _, partialDate := range util.SplitIntoDates(pq.Mint, pq.Maxt) {

			if windowActive {
				pt := partialDate.ToTime()
				if minTimeDate != "" && pt.Before(winStartDay) {
					continue
				}
				if maxTimeDate != "" && !pt.Before(winEndDay) {
					continue
				}
			}
			pqDates[partialDate] = struct{}{}
		}
	}

	// Find TSDB dates not covered by parquet dates.
	steps := make([]Step, 0, len(tsdbDates))
	for date, metas := range tsdbDates {
		if !date.ToTime().Before(p.notAfter) {
			// Ignore TSDB blocks that are for dates excluded from conversions.
			continue
		}

		if _, ok := pqDates[date]; ok {
			// This date is already covered by a parquet block.
			continue
		}

		// Sort our tsdb block metas from oldest to the newest.
		slices.SortFunc(metas, func(a, b metadata.Meta) int {
			return cmp.Compare(a.MinTime, b.MinTime)
		})

		steps = append(steps, Step{
			Date:    date,
			Sources: metas,
		})
	}

	// Sort our days, most recent first.
	slices.SortFunc(steps, func(a, b Step) int {
		return cmp.Compare(b.Date.MinT(), a.Date.MinT())
	})

	// Remove the most recent day if it's not fully covered by TSDB blocks.
	// We do this because we might get some delayed blocks for it.
	// Any gaps in days older than the most recent one are ignored.
	steps = truncateLastPartialDay(steps)

	// Restrict our plan to have only up to maxDays number of steps.
	// But allow more steps if they come from TSDB blocks that are only on the plan.
	steps = limitSteps(steps, p.maxDays)

	return Plan{Steps: steps}
}

func (p Planner) Plan(tsdbStreams map[schema.ExternalLabelsHash]schema.TSDBBlocksStream, parquetStreams map[schema.ExternalLabelsHash]schema.ParquetBlocksStream, minTimeDate string, maxTimeDate string) Plan {
	outPlan := Plan{Steps: []Step{}}

	for tsdbEH := range tsdbStreams {
		parquet, ok := parquetStreams[tsdbEH]
		if !ok {
			parquet = schema.ParquetBlocksStream{}
		}

		streamPlan := p.planStream(tsdbStreams[tsdbEH], parquet, minTimeDate, maxTimeDate)
		for i := range streamPlan.Steps {
			streamPlan.Steps[i].ExternalLabels = tsdbStreams[tsdbEH].ExternalLabels
		}

		outPlan.Steps = append(outPlan.Steps, streamPlan.Steps...)
	}

	return outPlan
}

// Get rid of the most recent day if we don't have TSDB blocks for the whole day.
func truncateLastPartialDay(steps []Step) []Step {
	// Empty source, return it as is.
	if len(steps) < 1 {
		return steps
	}
	// Most recent day is NOT fully covered, return all but most recent day.
	if !steps[0].isFullyCovered() {
		return steps[1:]
	}
	// Return as is.
	return steps
}

// Limit the plan to specified max number of days.
// This is a soft limit, final plan might have more days if it makes sense.
func limitSteps(steps []Step, limit int) []Step {
	if len(steps) <= limit {
		// Plan is <= the limit, return it as is.
		return steps
	}

	metas := MergeMetas(steps[:limit])
	for i := limit; i < len(steps); i++ {
		// Loop over all excess days and check if they would require a new block.
		// If yes then exclude them from the plan.
		// If no then keep them on the plan.
		var newBlock bool
		for _, meta := range steps[i].Sources {
			if !slices.ContainsFunc(metas, func(m metadata.Meta) bool {
				return meta.ULID == m.ULID
			}) {
				newBlock = true
			}
		}
		if newBlock {
			break
		}
		limit = i + 1
	}

	return steps[:limit]
}

// Merge dates from all steps into a single slice.
func MergeDates(steps []Step) []util.Date {
	dates := make([]util.Date, 0, len(steps))
	for _, step := range steps {
		dates = append(dates, step.Date)
	}
	return dates
}

// Merge TSDB block metas from all steps into a single slice.
func MergeMetas(steps []Step) (metas []metadata.Meta) {
	for _, step := range steps {
		for _, m := range step.Sources {
			if !slices.ContainsFunc(metas, func(meta metadata.Meta) bool {
				return meta.ULID == m.ULID
			}) {
				metas = append(metas, m)
			}
		}
	}
	return metas
}