perf/go/dfbuilder/dfbuilder.go - buildbot - Git at Google

 package dfbuilder

 import (
 	"context"
 	"fmt"
 	"net/url"
 	"sync"
 	"time"

 	"go.opencensus.io/trace"
 	"go.skia.org/infra/go/git/gitinfo"
 	"go.skia.org/infra/go/paramtools"
 	"go.skia.org/infra/go/query"
 	"go.skia.org/infra/go/sklog"
 	"go.skia.org/infra/go/timer"
 	"go.skia.org/infra/go/vcsinfo"
 	"go.skia.org/infra/go/vec32"
 	"go.skia.org/infra/perf/go/btts"
 	"go.skia.org/infra/perf/go/cid"
 	"go.skia.org/infra/perf/go/dataframe"
 	"go.skia.org/infra/perf/go/tracesetbuilder"
 	"go.skia.org/infra/perf/go/types"
 	"golang.org/x/sync/errgroup"
 )

 const (
 	// NEW_N_FROM_KEY_STEP is the length of time to do a search for each step
 	// when constructing a NewN* query.
 	NEW_N_FROM_KEY_STEP = 12 * time.Hour

 	// NEW_N_MAX_SEARCH is the minimum number of queries to perform that returned
 	// no data before giving up.
 	NEW_N_MAX_SEARCH = 4
 )

 // builder implements DataFrameBuilder using btts.
 type builder struct {
 	vcs      vcsinfo.VCS
 	store    *btts.BigTableTraceStore
 	tileSize int32
 }

 func NewDataFrameBuilderFromBTTS(vcs vcsinfo.VCS, store *btts.BigTableTraceStore) dataframe.DataFrameBuilder {
 	return &builder{
 		vcs:      vcs,
 		store:    store,
 		tileSize: store.TileSize(),
 	}
 }

 // fromIndexCommit returns the slices of ColumnHeader and index.
 // The slices are populated from the given vcsinfo.IndexCommits.
 //
 // The value for 'skip', the number of commits skipped, is passed through to
 // the return values.
 func fromIndexCommit(resp []*vcsinfo.IndexCommit, skip int) ([]*dataframe.ColumnHeader, []int32, int) {
 	headers := []*dataframe.ColumnHeader{}
 	indices := []int32{}
 	for _, r := range resp {
 		headers = append(headers, &dataframe.ColumnHeader{
 			Source:    "master",
 			Offset:    int64(r.Index),
 			Timestamp: r.Timestamp.Unix(),
 		})
 		indices = append(indices, int32(r.Index))
 	}
 	return headers, indices, skip
 }

 // lastNCommits returns the slices of ColumnHeader and cid.CommitID that are
 // needed by DataFrame and ptracestore.PTraceStore, respectively. The slices
 // are for the last N commits in the repo.
 //
 // Returns 0 for 'skip', the number of commits skipped.
 func lastNCommits(vcs vcsinfo.VCS, n int) ([]*dataframe.ColumnHeader, []int32, int) {
 	return fromIndexCommit(vcs.LastNIndex(n), 0)
 }

 // fromIndexRange returns the headers and indices for all the commits
 // between beginIndex and endIndex inclusive.
 func fromIndexRange(ctx context.Context, vcs vcsinfo.VCS, beginIndex, endIndex int32) ([]*dataframe.ColumnHeader, []int32, int, error) {
 	ctx, span := trace.StartSpan(ctx, "dfbuilder fromIndexRange")
 	defer span.End()

 	g, ok := vcs.(*gitinfo.GitInfo)
 	headers := []*dataframe.ColumnHeader{}
 	indices := []int32{}
 	for i := beginIndex; i <= endIndex; i++ {
 		if ok {
 			// This is a temporary performance enhancement for Perf.
 			// It will be removed once Perf moves to gitstore.
 			ts, err := g.TimestampAtIndex(int(i))
 			if err != nil {
 				return nil, nil, 0, fmt.Errorf("Range of commits invalid: %s", err)
 			}
 			headers = append(headers, &dataframe.ColumnHeader{
 				Source:    "master",
 				Offset:    int64(i),
 				Timestamp: ts.Unix(),
 			})
 			indices = append(indices, i)
 		} else {
 			commit, err := vcs.ByIndex(ctx, int(i))
 			if err != nil {
 				return nil, nil, 0, fmt.Errorf("Range of commits invalid: %s", err)
 			}
 			headers = append(headers, &dataframe.ColumnHeader{
 				Source:    "master",
 				Offset:    int64(i),
 				Timestamp: commit.Timestamp.Unix(),
 			})
 			indices = append(indices, i)
 		}
 	}
 	return headers, indices, 0, nil
 }

 // fromTimeRange returns the slices of ColumnHeader and int32. The slices
 // are for the commits that fall in the given time range [begin, end).
 //
 // If 'downsample' is true then the number of commits returned is limited
 // to MAX_SAMPLE_SIZE.
 //
 // The value for 'skip', the number of commits skipped, is also returned.
 func fromTimeRange(vcs vcsinfo.VCS, begin, end time.Time, downsample bool) ([]*dataframe.ColumnHeader, []int32, int) {
 	commits := vcs.Range(begin, end)
 	skip := 0
 	if downsample {
 		commits, skip = dataframe.DownSample(commits, dataframe.MAX_SAMPLE_SIZE)
 	}
 	return fromIndexCommit(commits, skip)
 }

 // tileMapOffsetToIndex maps the offset of each point in a tile to the index it
 // should appear in the resulting Trace.
 type tileMapOffsetToIndex map[btts.TileKey]map[int32]int32

 // buildTileMapOffsetToIndex returns a tileMapOffsetToIndex for the given indices and the given BigTableTraceStore.
 //
 // The returned map is used when loading traces out of tiles.
 func buildTileMapOffsetToIndex(indices []int32, store *btts.BigTableTraceStore) tileMapOffsetToIndex {
 	ret := tileMapOffsetToIndex{}
 	for targetIndex, sourceIndex := range indices {
 		tileKey := store.TileKey(sourceIndex)
 		if traceMap, ok := ret[tileKey]; !ok {
 			ret[tileKey] = map[int32]int32{
 				store.OffsetFromIndex(sourceIndex): int32(targetIndex),
 			}
 		} else {
 			traceMap[store.OffsetFromIndex(sourceIndex)] = int32(targetIndex)
 		}
 	}
 	return ret
 }

 // new builds a DataFrame for the given columns and populates it with traces that match the given query.
 //
 // The progress callback is triggered once for every tile.
 func (b *builder) new(ctx context.Context, colHeaders []*dataframe.ColumnHeader, indices []int32, q *query.Query, progress types.Progress, skip int) (*dataframe.DataFrame, error) {
 	ctx, span := trace.StartSpan(ctx, "dfbuilder.new")
 	defer span.End()

 	// TODO tickle progress as each Go routine completes.
 	defer timer.New("dfbuilder_new").Stop()
 	// Determine which tiles we are querying over, and how each tile maps into our results.
 	mapper := buildTileMapOffsetToIndex(indices, b.store)

 	traceSetBuilder := tracesetbuilder.New(len(indices))
 	defer traceSetBuilder.Close()

 	var mutex sync.Mutex // mutex protects stepsCompleted.
 	stepsCompleted := 0
 	triggerProgress := func() {
 		mutex.Lock()
 		defer mutex.Unlock()
 		stepsCompleted += 1
 		if progress != nil {
 			progress(stepsCompleted, len(mapper))
 		}
 	}

 	var g errgroup.Group
 	// For each tile.
 	for tileKey, traceMap := range mapper {
 		tileKey := tileKey
 		traceMap := traceMap
 		// TODO(jcgregorio) If we query across a large number of tiles N then this will spawn N*8 Go routines
 		// all hitting the backend at the same time. Maybe we need a worker pool if this becomes a problem.
 		g.Go(func() error {
 			defer timer.New("dfbuilder_by_tile").Stop()

 			// Query for matching traces in the given tile.
 			traces, err := b.store.QueryTracesByIndex(ctx, tileKey, q)
 			if err != nil {
 				return err
 			}
 			sklog.Debugf("found %d traces for %s", len(traces), tileKey.OpsRowName())

 			traceSetBuilder.Add(traceMap, traces)
 			triggerProgress()
 			return nil
 		})
 	}
 	if err := g.Wait(); err != nil {
 		return nil, fmt.Errorf("Failed while querying: %s", err)
 	}
 	traceSet, paramSet := traceSetBuilder.Build(ctx)
 	paramSet.Normalize()
 	d := &dataframe.DataFrame{
 		TraceSet: traceSet,
 		Header:   colHeaders,
 		ParamSet: paramSet,
 		Skip:     skip,
 	}
 	return d, nil
 }

 // See DataFrameBuilder.
 func (b *builder) New(progress types.Progress) (*dataframe.DataFrame, error) {
 	return b.NewN(progress, dataframe.DEFAULT_NUM_COMMITS)
 }

 // See DataFrameBuilder.
 func (b *builder) NewN(progress types.Progress, n int) (*dataframe.DataFrame, error) {
 	colHeaders, indices, skip := lastNCommits(b.vcs, n)
 	q, err := query.New(url.Values{})
 	if err != nil {
 		return nil, err
 	}
 	return b.new(context.TODO(), colHeaders, indices, q, progress, skip)
 }

 // See DataFrameBuilder.
 func (b *builder) NewFromQueryAndRange(begin, end time.Time, q *query.Query, downsample bool, progress types.Progress) (*dataframe.DataFrame, error) {
 	colHeaders, indices, skip := fromTimeRange(b.vcs, begin, end, downsample)
 	return b.new(context.TODO(), colHeaders, indices, q, progress, skip)
 }

 // See DataFrameBuilder.
 func (b *builder) NewFromKeysAndRange(keys []string, begin, end time.Time, downsample bool, progress types.Progress) (*dataframe.DataFrame, error) {
 	// TODO tickle progress as each Go routine completes.
 	defer timer.New("NewFromKeysAndRange").Stop()
 	colHeaders, indices, skip := fromTimeRange(b.vcs, begin, end, downsample)

 	// Determine which tiles we are querying over, and how each tile maps into our results.
 	mapper := buildTileMapOffsetToIndex(indices, b.store)

 	var mutex sync.Mutex // mutex protects traceSet and paramSet.
 	traceSet := types.TraceSet{}
 	paramSet := paramtools.ParamSet{}
 	stepsCompleted := 0
 	// triggerProgress must only be called when the caller has mutex locked.
 	triggerProgress := func() {
 		stepsCompleted += 1
 		if progress != nil {
 			progress(stepsCompleted, len(mapper))
 		}
 	}

 	var g errgroup.Group
 	// For each tile.
 	for tileKey, traceMap := range mapper {
 		tileKey := tileKey
 		traceMap := traceMap
 		g.Go(func() error {
 			// Read the traces for the given keys.
 			traces, err := b.store.ReadTraces(tileKey, keys)
 			if err != nil {
 				return err
 			}
 			mutex.Lock()
 			defer mutex.Unlock()
 			// For each trace, convert the encodedKey to a structured key
 			// and copy the trace values into their final destination.
 			for key, tileTrace := range traces {
 				trace, ok := traceSet[key]
 				if !ok {
 					trace = types.NewTrace(len(indices))
 				}
 				for srcIndex, dstIndex := range traceMap {
 					trace[dstIndex] = tileTrace[srcIndex]
 				}
 				traceSet[key] = trace
 				p, err := query.ParseKey(key)
 				if err != nil {
 					continue
 				}
 				paramSet.AddParams(p)
 			}
 			return nil
 		})
 	}
 	if err := g.Wait(); err != nil {
 		return nil, fmt.Errorf("Failed while querying: %s", err)
 	}
 	d := &dataframe.DataFrame{
 		TraceSet: traceSet,
 		Header:   colHeaders,
 		ParamSet: paramSet,
 		Skip:     skip,
 	}
 	triggerProgress()
 	return d, nil
 }

 // See DataFrameBuilder.
 func (b *builder) NewFromCommitIDsAndQuery(ctx context.Context, cids []*cid.CommitID, cidl *cid.CommitIDLookup, q *query.Query, progress types.Progress) (*dataframe.DataFrame, error) {
 	details, err := cidl.Lookup(ctx, cids)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to look up CommitIDs: %s", err)
 	}
 	colHeaders := []*dataframe.ColumnHeader{}
 	indices := []int32{}
 	for _, d := range details {
 		colHeaders = append(colHeaders, &dataframe.ColumnHeader{
 			Source:    d.Source,
 			Offset:    int64(d.Offset),
 			Timestamp: d.Timestamp,
 		})
 		indices = append(indices, int32(d.Offset))
 	}
 	return b.new(ctx, colHeaders, indices, q, progress, 0)
 }

 // findIndexForTime finds the index of the closest commit <= 'end'.
 //
 // Pass in zero time, i.e. time.Time{} to indicate to just get the most recent commit.
 func (b *builder) findIndexForTime(ctx context.Context, end time.Time) (int32, error) {
 	ctx, span := trace.StartSpan(ctx, "dfbuilder.findIndexForTime")
 	defer span.End()

 	var err error
 	endIndex := 0

 	if end.IsZero() {
 		commits := b.vcs.LastNIndex(1)
 		if len(commits) == 0 {
 			return 0, fmt.Errorf("Failed to find an end commit.")
 		}
 		return int32(commits[0].Index), nil
 	}

 	hashes := b.vcs.From(end)
 	if len(hashes) > 0 {
 		endIndex, err = b.vcs.IndexOf(ctx, hashes[0])
 		if err != nil {
 			return 0, fmt.Errorf("Failed loading end commit: %s", err)
 		}
 	} else {
 		commits := b.vcs.LastNIndex(1)
 		if len(commits) == 0 {
 			return 0, fmt.Errorf("Failed to find an end commit.")
 		}
 		endIndex = commits[0].Index
 	}
 	return int32(endIndex), nil
 }

 // See DataFrameBuilder.
 func (b *builder) NewNFromQuery(ctx context.Context, end time.Time, q *query.Query, n int32, progress types.Progress) (*dataframe.DataFrame, error) {
 	ctx, span := trace.StartSpan(ctx, "dfbuilder.NewNFromQuery")
 	defer span.End()

 	sklog.Infof("Querying to: %v", end)

 	ret := dataframe.NewEmpty()
 	var total int32 // total number of commits we've added to ret so far.
 	steps := 1      // Total number of times we've gone through the loop below, used in the progress() callback.
 	numStepsNoData := 0

 	endIndex, err := b.findIndexForTime(ctx, end)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to find end index: %s", err)
 	}
 	// beginIndex is the index of the first commit in the tile that endIndex is
 	// in. We are OK if beginIndex == endIndex because fromIndexRange returns
 	// headers from begin to end *inclusive*.
 	beginIndex := b.store.IndexOfTileStart(endIndex)

 	sklog.Infof("BeginIndex: %d  EndIndex: %d", beginIndex, endIndex)
 	for total < n {

 		// Query for traces.
 		headers, indices, skip, err := fromIndexRange(ctx, b.vcs, beginIndex, endIndex)
 		if err != nil {
 			return nil, fmt.Errorf("Failed building index range: %s", err)
 		}
 		df, err := b.new(ctx, headers, indices, q, nil, skip)
 		if err != nil {
 			return nil, fmt.Errorf("Failed while querying: %s", err)
 		}

 		nonMissing := 0
 		// Total up the number of data points we have for each commit.
 		counts := make([]int, len(df.Header))
 		for _, tr := range df.TraceSet {
 			for i, x := range tr {
 				if x != vec32.MISSING_DATA_SENTINEL {
 					counts[i] += 1
 					nonMissing += 1
 				}
 			}
 		}
 		// If there are no matches then we might be done.
 		if nonMissing == 0 {
 			numStepsNoData += 1
 		}
 		if numStepsNoData > NEW_N_MAX_SEARCH {
 			sklog.Infof("Failed querying: %s", q)
 			break
 		}

 		ret.ParamSet.AddParamSet(df.ParamSet)

 		// For each commit that has data, copy the data from df into ret.
 		// Move backwards down the trace since we are building the result from 'end' backwards.
 		for i := len(counts) - 1; i >= 0; i-- {
 			if counts[i] > 0 {
 				ret.Header = append([]*dataframe.ColumnHeader{df.Header[i]}, ret.Header...)
 				for key, sourceTrace := range df.TraceSet {
 					if _, ok := ret.TraceSet[key]; !ok {
 						ret.TraceSet[key] = vec32.New(int(n))
 					}
 					ret.TraceSet[key][n-1-total] = sourceTrace[i]
 				}
 				total += 1

 				// If we've added enough commits to ret then we are done.
 				if total == n {
 					break
 				}
 			}
 		}
 		sklog.Infof("Total: %d Steps: %d NumStepsNoData: %d", total, steps, numStepsNoData)

 		if total == n {
 			break
 		}

 		if progress != nil {
 			progress(steps, steps+1)
 		}
 		steps += 1

 		// Now step back a full tile.

 		// At this point we know beginIndex points to the 0th column in a tile,
 		// so endIndex is easy to calculate.
 		endIndex = beginIndex - 1
 		if endIndex < 0 {
 			break
 		}
 		beginIndex = b.store.IndexOfTileStart(endIndex)
 		if beginIndex < 0 {
 			beginIndex = 0
 		}
 	}

 	if total < n {
 		// Trim down the traces so they are the same length as ret.Header.
 		for key, tr := range ret.TraceSet {
 			ret.TraceSet[key] = tr[n-total:]
 		}
 	}

 	return ret, nil
 }

 // See DataFrameBuilder.
 func (b *builder) NewNFromKeys(ctx context.Context, end time.Time, keys []string, n int32, progress types.Progress) (*dataframe.DataFrame, error) {
 	defer timer.New("NewNFromKeys").Stop()

 	endIndex, err := b.findIndexForTime(ctx, end)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to find end index: %s", err)
 	}
 	beginIndex := endIndex - (b.tileSize - 1)
 	if beginIndex < 0 {
 		beginIndex = 0
 	}

 	ret := dataframe.NewEmpty()
 	var total int32 // total number of commits we've added to ret so far.
 	steps := 1      // Total number of times we've gone through the loop below, used in the progress() callback.
 	numStepsNoData := 0

 	for total < n {
 		headers, indices, skip, err := fromIndexRange(ctx, b.vcs, beginIndex, endIndex)
 		if err != nil {
 			return nil, fmt.Errorf("Failed building index range: %s", err)
 		}

 		// Determine which tiles we are querying over, and how each tile maps into our results.
 		mapper := buildTileMapOffsetToIndex(indices, b.store)

 		traceSet := types.TraceSet{}
 		for tileKey, traceMap := range mapper {
 			// Read the traces for the given keys.
 			traces, err := b.store.ReadTraces(tileKey, keys)
 			if err != nil {
 				return nil, err
 			}
 			// For each trace, convert the encodedKey to a structured key
 			// and copy the trace values into their final destination.
 			for key, tileTrace := range traces {
 				trace, ok := traceSet[key]
 				if !ok {
 					trace = types.NewTrace(len(indices))
 				}
 				for srcIndex, dstIndex := range traceMap {
 					trace[dstIndex] = tileTrace[srcIndex]
 				}
 				traceSet[key] = trace
 			}
 		}
 		df := &dataframe.DataFrame{
 			TraceSet: traceSet,
 			Header:   headers,
 			ParamSet: paramtools.ParamSet{},
 			Skip:     skip,
 		}
 		df.BuildParamSet()

 		nonMissing := 0
 		// Total up the number of data points we have for each commit.
 		counts := make([]int, len(df.Header))
 		for _, tr := range df.TraceSet {
 			for i, x := range tr {
 				if x != vec32.MISSING_DATA_SENTINEL {
 					counts[i] += 1
 					nonMissing += 1
 				}
 			}
 		}
 		// If there are no matches then we might be done.
 		if nonMissing == 0 {
 			numStepsNoData += 1
 		}
 		if numStepsNoData > NEW_N_MAX_SEARCH {
 			break
 		}

 		ret.ParamSet.AddParamSet(df.ParamSet)

 		// For each commit that has data, copy the data from df into ret.
 		// Move backwards down the trace since we are building the result from 'end' backwards.
 		for i := len(counts) - 1; i >= 0; i-- {
 			if counts[i] > 0 {
 				ret.Header = append([]*dataframe.ColumnHeader{df.Header[i]}, ret.Header...)
 				for key, sourceTrace := range df.TraceSet {
 					if _, ok := ret.TraceSet[key]; !ok {
 						ret.TraceSet[key] = vec32.New(int(n))
 					}
 					ret.TraceSet[key][n-1-total] = sourceTrace[i]
 				}
 				total += 1

 				// If we've added enough commits to ret then we are done.
 				if total == n {
 					break
 				}
 			}
 		}

 		sklog.Infof("Total: %d Steps: %d NumStepsNoData: %d", total, steps, numStepsNoData)

 		if total == n {
 			break
 		}

 		if progress != nil {
 			progress(steps, steps+1)
 		}
 		steps += 1

 		endIndex -= b.tileSize
 		beginIndex -= b.tileSize
 		if endIndex < 0 {
 			break
 		}
 		if beginIndex < 0 {
 			beginIndex = 0
 		}

 	}

 	if total < n {
 		// Trim down the traces so they are the same length as ret.Header.
 		for key, tr := range ret.TraceSet {
 			ret.TraceSet[key] = tr[n-total:]
 		}
 	}

 	return ret, nil
 }

 // See DataFrameBuilder.
 func (b *builder) PreflightQuery(ctx context.Context, end time.Time, q *query.Query) (int64, paramtools.ParamSet, error) {
 	var count int64
 	ps := paramtools.ParamSet{}

 	tileKey, err := b.store.GetLatestTile()
 	if err != nil {
 		return -1, nil, err
 	}

 	if q.Empty() {
 		// If the query is empty then we have a shortcut for building the
 		// ParamSet by just using the OPS. In that case we only need to count
 		// encodedKeys to get the count.
 		for i := 0; i < 2; i++ {
 			ops, err := b.store.GetOrderedParamSet(ctx, tileKey)
 			if err != nil {
 				return -1, nil, err
 			}
 			ps.AddParamSet(ops.ParamSet)
 			tileKey = tileKey.PrevTile()
 		}

 		count, err = b.store.TraceCount(ctx, tileKey)
 		if err != nil {
 			return -1, nil, err
 		}

 	} else {
 		// Since the query isn't empty we'll have to run a partial query
 		// to build the ParamSet. Do so over the two most recent tiles.

 		// Record the OPS for the first tile.
 		opsOne, err := b.store.GetOrderedParamSet(ctx, tileKey)
 		if err != nil {
 			return -1, nil, err
 		}
 		// Count the matches and sum the params in the first tile.
 		out, err := b.store.QueryTracesIDOnlyByIndex(ctx, tileKey, q)
 		if err != nil {
 			return -1, nil, fmt.Errorf("Failed to query traces: %s", err)
 		}
 		var tileOneCount int64
 		for p := range out {
 			tileOneCount++
 			ps.AddParams(p)
 		}

 		// Now move to the previous tile.
 		tileKey = tileKey.PrevTile()

 		// Record the OPS for the second tile.
 		opsTwo, err := b.store.GetOrderedParamSet(ctx, tileKey)
 		if err != nil {
 			return -1, nil, err
 		}

 		// Count the matches and sum the params in the second tile.
 		out, err = b.store.QueryTracesIDOnlyByIndex(ctx, tileKey, q)
 		if err != nil {
 			return -1, nil, fmt.Errorf("Failed to query traces: %s", err)
 		}
 		var tileTwoCount int64
 		for p := range out {
 			tileTwoCount++
 			ps.AddParams(p)
 		}

 		// Use the larger of the two counts as our result.
 		if tileOneCount > tileTwoCount {
 			count = tileOneCount
 		} else {
 			count = tileTwoCount
 		}

 		// Use the larger of the two OPSs to work with.
 		var ops *paramtools.OrderedParamSet
 		if opsOne.ParamSet.Size() > opsTwo.ParamSet.Size() {
 			ops = opsOne
 		} else {
 			ops = opsTwo
 		}

 		// Now we have the ParamSet that corresponds to the query, but for each
 		// key in the query we need to go back and put in all the values that
 		// appear for that key since the user can make more selections in that
 		// key.
 		queryPlan, err := q.QueryPlan(ops)
 		if err != nil {
 			return -1, nil, err
 		}
 		for key := range queryPlan {
 			ps[key] = ops.ParamSet[key]
 		}
 	}

 	ps.Normalize()
 	return count, ps, nil
 }

 // Validate that the concrete bttsDataFrameBuilder faithfully implements the DataFrameBuidler interface.
 var _ dataframe.DataFrameBuilder = (*builder)(nil)
	package dfbuilder

	import (
	"context"
	"fmt"
	"net/url"
	"sync"
	"time"

	"go.opencensus.io/trace"
	"go.skia.org/infra/go/git/gitinfo"
	"go.skia.org/infra/go/paramtools"
	"go.skia.org/infra/go/query"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/timer"
	"go.skia.org/infra/go/vcsinfo"
	"go.skia.org/infra/go/vec32"
	"go.skia.org/infra/perf/go/btts"
	"go.skia.org/infra/perf/go/cid"
	"go.skia.org/infra/perf/go/dataframe"
	"go.skia.org/infra/perf/go/tracesetbuilder"
	"go.skia.org/infra/perf/go/types"
	"golang.org/x/sync/errgroup"
	)

	const (
	// NEW_N_FROM_KEY_STEP is the length of time to do a search for each step
	// when constructing a NewN* query.
	NEW_N_FROM_KEY_STEP = 12 * time.Hour

	// NEW_N_MAX_SEARCH is the minimum number of queries to perform that returned
	// no data before giving up.
	NEW_N_MAX_SEARCH = 4
	)

	// builder implements DataFrameBuilder using btts.
	type builder struct {
	vcs vcsinfo.VCS
	store *btts.BigTableTraceStore
	tileSize int32
	}

	func NewDataFrameBuilderFromBTTS(vcs vcsinfo.VCS, store *btts.BigTableTraceStore) dataframe.DataFrameBuilder {
	return &builder{
	vcs: vcs,
	store: store,
	tileSize: store.TileSize(),
	}
	}

	// fromIndexCommit returns the slices of ColumnHeader and index.
	// The slices are populated from the given vcsinfo.IndexCommits.
	//
	// The value for 'skip', the number of commits skipped, is passed through to
	// the return values.
	func fromIndexCommit(resp []vcsinfo.IndexCommit, skip int) ([]dataframe.ColumnHeader, []int32, int) {
	headers := []*dataframe.ColumnHeader{}
	indices := []int32{}
	for _, r := range resp {
	headers = append(headers, &dataframe.ColumnHeader{
	Source: "master",
	Offset: int64(r.Index),
	Timestamp: r.Timestamp.Unix(),
	})
	indices = append(indices, int32(r.Index))
	}
	return headers, indices, skip
	}

	// lastNCommits returns the slices of ColumnHeader and cid.CommitID that are
	// needed by DataFrame and ptracestore.PTraceStore, respectively. The slices
	// are for the last N commits in the repo.
	//
	// Returns 0 for 'skip', the number of commits skipped.
	func lastNCommits(vcs vcsinfo.VCS, n int) ([]*dataframe.ColumnHeader, []int32, int) {
	return fromIndexCommit(vcs.LastNIndex(n), 0)
	}

	// fromIndexRange returns the headers and indices for all the commits
	// between beginIndex and endIndex inclusive.
	func fromIndexRange(ctx context.Context, vcs vcsinfo.VCS, beginIndex, endIndex int32) ([]*dataframe.ColumnHeader, []int32, int, error) {
	ctx, span := trace.StartSpan(ctx, "dfbuilder fromIndexRange")
	defer span.End()

	g, ok := vcs.(*gitinfo.GitInfo)
	headers := []*dataframe.ColumnHeader{}
	indices := []int32{}
	for i := beginIndex; i <= endIndex; i++ {
	if ok {
	// This is a temporary performance enhancement for Perf.
	// It will be removed once Perf moves to gitstore.
	ts, err := g.TimestampAtIndex(int(i))
	if err != nil {
	return nil, nil, 0, fmt.Errorf("Range of commits invalid: %s", err)
	}
	headers = append(headers, &dataframe.ColumnHeader{
	Source: "master",
	Offset: int64(i),
	Timestamp: ts.Unix(),
	})
	indices = append(indices, i)
	} else {
	commit, err := vcs.ByIndex(ctx, int(i))
	if err != nil {
	return nil, nil, 0, fmt.Errorf("Range of commits invalid: %s", err)
	}
	headers = append(headers, &dataframe.ColumnHeader{
	Source: "master",
	Offset: int64(i),
	Timestamp: commit.Timestamp.Unix(),
	})
	indices = append(indices, i)
	}
	}
	return headers, indices, 0, nil
	}

	// fromTimeRange returns the slices of ColumnHeader and int32. The slices
	// are for the commits that fall in the given time range [begin, end).
	//
	// If 'downsample' is true then the number of commits returned is limited
	// to MAX_SAMPLE_SIZE.
	//
	// The value for 'skip', the number of commits skipped, is also returned.
	func fromTimeRange(vcs vcsinfo.VCS, begin, end time.Time, downsample bool) ([]*dataframe.ColumnHeader, []int32, int) {
	commits := vcs.Range(begin, end)
	skip := 0
	if downsample {
	commits, skip = dataframe.DownSample(commits, dataframe.MAX_SAMPLE_SIZE)
	}
	return fromIndexCommit(commits, skip)
	}

	// tileMapOffsetToIndex maps the offset of each point in a tile to the index it
	// should appear in the resulting Trace.
	type tileMapOffsetToIndex map[btts.TileKey]map[int32]int32

	// buildTileMapOffsetToIndex returns a tileMapOffsetToIndex for the given indices and the given BigTableTraceStore.
	//
	// The returned map is used when loading traces out of tiles.
	func buildTileMapOffsetToIndex(indices []int32, store *btts.BigTableTraceStore) tileMapOffsetToIndex {
	ret := tileMapOffsetToIndex{}
	for targetIndex, sourceIndex := range indices {
	tileKey := store.TileKey(sourceIndex)
	if traceMap, ok := ret[tileKey]; !ok {
	ret[tileKey] = map[int32]int32{
	store.OffsetFromIndex(sourceIndex): int32(targetIndex),
	}
	} else {
	traceMap[store.OffsetFromIndex(sourceIndex)] = int32(targetIndex)
	}
	}
	return ret
	}

	// new builds a DataFrame for the given columns and populates it with traces that match the given query.
	//
	// The progress callback is triggered once for every tile.
	func (b builder) new(ctx context.Context, colHeaders []dataframe.ColumnHeader, indices []int32, q query.Query, progress types.Progress, skip int) (dataframe.DataFrame, error) {
	ctx, span := trace.StartSpan(ctx, "dfbuilder.new")
	defer span.End()

	// TODO tickle progress as each Go routine completes.
	defer timer.New("dfbuilder_new").Stop()
	// Determine which tiles we are querying over, and how each tile maps into our results.
	mapper := buildTileMapOffsetToIndex(indices, b.store)

	traceSetBuilder := tracesetbuilder.New(len(indices))
	defer traceSetBuilder.Close()

	var mutex sync.Mutex // mutex protects stepsCompleted.
	stepsCompleted := 0
	triggerProgress := func() {
	mutex.Lock()
	defer mutex.Unlock()
	stepsCompleted += 1
	if progress != nil {
	progress(stepsCompleted, len(mapper))
	}
	}

	var g errgroup.Group
	// For each tile.
	for tileKey, traceMap := range mapper {
	tileKey := tileKey
	traceMap := traceMap
	// TODO(jcgregorio) If we query across a large number of tiles N then this will spawn N*8 Go routines
	// all hitting the backend at the same time. Maybe we need a worker pool if this becomes a problem.
	g.Go(func() error {
	defer timer.New("dfbuilder_by_tile").Stop()

	// Query for matching traces in the given tile.
	traces, err := b.store.QueryTracesByIndex(ctx, tileKey, q)
	if err != nil {
	return err
	}
	sklog.Debugf("found %d traces for %s", len(traces), tileKey.OpsRowName())

	traceSetBuilder.Add(traceMap, traces)
	triggerProgress()
	return nil
	})
	}
	if err := g.Wait(); err != nil {
	return nil, fmt.Errorf("Failed while querying: %s", err)
	}
	traceSet, paramSet := traceSetBuilder.Build(ctx)
	paramSet.Normalize()
	d := &dataframe.DataFrame{
	TraceSet: traceSet,
	Header: colHeaders,
	ParamSet: paramSet,
	Skip: skip,
	}
	return d, nil
	}

	// See DataFrameBuilder.
	func (b builder) New(progress types.Progress) (dataframe.DataFrame, error) {
	return b.NewN(progress, dataframe.DEFAULT_NUM_COMMITS)
	}

	// See DataFrameBuilder.
	func (b builder) NewN(progress types.Progress, n int) (dataframe.DataFrame, error) {
	colHeaders, indices, skip := lastNCommits(b.vcs, n)
	q, err := query.New(url.Values{})
	if err != nil {
	return nil, err
	}
	return b.new(context.TODO(), colHeaders, indices, q, progress, skip)
	}

	// See DataFrameBuilder.
	func (b builder) NewFromQueryAndRange(begin, end time.Time, q query.Query, downsample bool, progress types.Progress) (*dataframe.DataFrame, error) {
	colHeaders, indices, skip := fromTimeRange(b.vcs, begin, end, downsample)
	return b.new(context.TODO(), colHeaders, indices, q, progress, skip)
	}

	// See DataFrameBuilder.
	func (b builder) NewFromKeysAndRange(keys []string, begin, end time.Time, downsample bool, progress types.Progress) (dataframe.DataFrame, error) {
	// TODO tickle progress as each Go routine completes.
	defer timer.New("NewFromKeysAndRange").Stop()
	colHeaders, indices, skip := fromTimeRange(b.vcs, begin, end, downsample)

	// Determine which tiles we are querying over, and how each tile maps into our results.
	mapper := buildTileMapOffsetToIndex(indices, b.store)

	var mutex sync.Mutex // mutex protects traceSet and paramSet.
	traceSet := types.TraceSet{}
	paramSet := paramtools.ParamSet{}
	stepsCompleted := 0
	// triggerProgress must only be called when the caller has mutex locked.
	triggerProgress := func() {
	stepsCompleted += 1
	if progress != nil {
	progress(stepsCompleted, len(mapper))
	}
	}

	var g errgroup.Group
	// For each tile.
	for tileKey, traceMap := range mapper {
	tileKey := tileKey
	traceMap := traceMap
	g.Go(func() error {
	// Read the traces for the given keys.
	traces, err := b.store.ReadTraces(tileKey, keys)
	if err != nil {
	return err
	}
	mutex.Lock()
	defer mutex.Unlock()
	// For each trace, convert the encodedKey to a structured key
	// and copy the trace values into their final destination.
	for key, tileTrace := range traces {
	trace, ok := traceSet[key]
	if !ok {
	trace = types.NewTrace(len(indices))
	}
	for srcIndex, dstIndex := range traceMap {
	trace[dstIndex] = tileTrace[srcIndex]
	}
	traceSet[key] = trace
	p, err := query.ParseKey(key)
	if err != nil {
	continue
	}
	paramSet.AddParams(p)
	}
	return nil
	})
	}
	if err := g.Wait(); err != nil {
	return nil, fmt.Errorf("Failed while querying: %s", err)
	}
	d := &dataframe.DataFrame{
	TraceSet: traceSet,
	Header: colHeaders,
	ParamSet: paramSet,
	Skip: skip,
	}
	triggerProgress()
	return d, nil
	}

	// See DataFrameBuilder.
	func (b builder) NewFromCommitIDsAndQuery(ctx context.Context, cids []cid.CommitID, cidl cid.CommitIDLookup, q query.Query, progress types.Progress) (*dataframe.DataFrame, error) {
	details, err := cidl.Lookup(ctx, cids)
	if err != nil {
	return nil, fmt.Errorf("Failed to look up CommitIDs: %s", err)
	}
	colHeaders := []*dataframe.ColumnHeader{}
	indices := []int32{}
	for _, d := range details {
	colHeaders = append(colHeaders, &dataframe.ColumnHeader{
	Source: d.Source,
	Offset: int64(d.Offset),
	Timestamp: d.Timestamp,
	})
	indices = append(indices, int32(d.Offset))
	}
	return b.new(ctx, colHeaders, indices, q, progress, 0)
	}

	// findIndexForTime finds the index of the closest commit <= 'end'.
	//
	// Pass in zero time, i.e. time.Time{} to indicate to just get the most recent commit.
	func (b *builder) findIndexForTime(ctx context.Context, end time.Time) (int32, error) {
	ctx, span := trace.StartSpan(ctx, "dfbuilder.findIndexForTime")
	defer span.End()

	var err error
	endIndex := 0

	if end.IsZero() {
	commits := b.vcs.LastNIndex(1)
	if len(commits) == 0 {
	return 0, fmt.Errorf("Failed to find an end commit.")
	}
	return int32(commits[0].Index), nil
	}

	hashes := b.vcs.From(end)
	if len(hashes) > 0 {
	endIndex, err = b.vcs.IndexOf(ctx, hashes[0])
	if err != nil {
	return 0, fmt.Errorf("Failed loading end commit: %s", err)
	}
	} else {
	commits := b.vcs.LastNIndex(1)
	if len(commits) == 0 {
	return 0, fmt.Errorf("Failed to find an end commit.")
	}
	endIndex = commits[0].Index
	}
	return int32(endIndex), nil
	}

	// See DataFrameBuilder.
	func (b builder) NewNFromQuery(ctx context.Context, end time.Time, q query.Query, n int32, progress types.Progress) (*dataframe.DataFrame, error) {
	ctx, span := trace.StartSpan(ctx, "dfbuilder.NewNFromQuery")
	defer span.End()

	sklog.Infof("Querying to: %v", end)

	ret := dataframe.NewEmpty()
	var total int32 // total number of commits we've added to ret so far.
	steps := 1 // Total number of times we've gone through the loop below, used in the progress() callback.
	numStepsNoData := 0

	endIndex, err := b.findIndexForTime(ctx, end)
	if err != nil {
	return nil, fmt.Errorf("Failed to find end index: %s", err)
	}
	// beginIndex is the index of the first commit in the tile that endIndex is
	// in. We are OK if beginIndex == endIndex because fromIndexRange returns
	// headers from begin to end inclusive.
	beginIndex := b.store.IndexOfTileStart(endIndex)

	sklog.Infof("BeginIndex: %d EndIndex: %d", beginIndex, endIndex)
	for total < n {

	// Query for traces.
	headers, indices, skip, err := fromIndexRange(ctx, b.vcs, beginIndex, endIndex)
	if err != nil {
	return nil, fmt.Errorf("Failed building index range: %s", err)
	}
	df, err := b.new(ctx, headers, indices, q, nil, skip)
	if err != nil {
	return nil, fmt.Errorf("Failed while querying: %s", err)
	}

	nonMissing := 0
	// Total up the number of data points we have for each commit.
	counts := make([]int, len(df.Header))
	for _, tr := range df.TraceSet {
	for i, x := range tr {
	if x != vec32.MISSING_DATA_SENTINEL {
	counts[i] += 1
	nonMissing += 1
	}
	}
	}
	// If there are no matches then we might be done.
	if nonMissing == 0 {
	numStepsNoData += 1
	}
	if numStepsNoData > NEW_N_MAX_SEARCH {
	sklog.Infof("Failed querying: %s", q)
	break
	}

	ret.ParamSet.AddParamSet(df.ParamSet)

	// For each commit that has data, copy the data from df into ret.
	// Move backwards down the trace since we are building the result from 'end' backwards.
	for i := len(counts) - 1; i >= 0; i-- {
	if counts[i] > 0 {
	ret.Header = append([]*dataframe.ColumnHeader{df.Header[i]}, ret.Header...)
	for key, sourceTrace := range df.TraceSet {
	if _, ok := ret.TraceSet[key]; !ok {
	ret.TraceSet[key] = vec32.New(int(n))
	}
	ret.TraceSet[key][n-1-total] = sourceTrace[i]
	}
	total += 1

	// If we've added enough commits to ret then we are done.
	if total == n {
	break
	}
	}
	}
	sklog.Infof("Total: %d Steps: %d NumStepsNoData: %d", total, steps, numStepsNoData)

	if total == n {
	break
	}

	if progress != nil {
	progress(steps, steps+1)
	}
	steps += 1

	// Now step back a full tile.

	// At this point we know beginIndex points to the 0th column in a tile,
	// so endIndex is easy to calculate.
	endIndex = beginIndex - 1
	if endIndex < 0 {
	break
	}
	beginIndex = b.store.IndexOfTileStart(endIndex)
	if beginIndex < 0 {
	beginIndex = 0
	}
	}

	if total < n {
	// Trim down the traces so they are the same length as ret.Header.
	for key, tr := range ret.TraceSet {
	ret.TraceSet[key] = tr[n-total:]
	}
	}

	return ret, nil
	}

	// See DataFrameBuilder.
	func (b builder) NewNFromKeys(ctx context.Context, end time.Time, keys []string, n int32, progress types.Progress) (dataframe.DataFrame, error) {
	defer timer.New("NewNFromKeys").Stop()

	endIndex, err := b.findIndexForTime(ctx, end)
	if err != nil {
	return nil, fmt.Errorf("Failed to find end index: %s", err)
	}
	beginIndex := endIndex - (b.tileSize - 1)
	if beginIndex < 0 {
	beginIndex = 0
	}

	ret := dataframe.NewEmpty()
	var total int32 // total number of commits we've added to ret so far.
	steps := 1 // Total number of times we've gone through the loop below, used in the progress() callback.
	numStepsNoData := 0

	for total < n {
	headers, indices, skip, err := fromIndexRange(ctx, b.vcs, beginIndex, endIndex)
	if err != nil {
	return nil, fmt.Errorf("Failed building index range: %s", err)
	}

	// Determine which tiles we are querying over, and how each tile maps into our results.
	mapper := buildTileMapOffsetToIndex(indices, b.store)

	traceSet := types.TraceSet{}
	for tileKey, traceMap := range mapper {
	// Read the traces for the given keys.
	traces, err := b.store.ReadTraces(tileKey, keys)
	if err != nil {
	return nil, err
	}
	// For each trace, convert the encodedKey to a structured key
	// and copy the trace values into their final destination.
	for key, tileTrace := range traces {
	trace, ok := traceSet[key]
	if !ok {
	trace = types.NewTrace(len(indices))
	}
	for srcIndex, dstIndex := range traceMap {
	trace[dstIndex] = tileTrace[srcIndex]
	}
	traceSet[key] = trace
	}
	}
	df := &dataframe.DataFrame{
	TraceSet: traceSet,
	Header: headers,
	ParamSet: paramtools.ParamSet{},
	Skip: skip,
	}
	df.BuildParamSet()

	nonMissing := 0
	// Total up the number of data points we have for each commit.
	counts := make([]int, len(df.Header))
	for _, tr := range df.TraceSet {
	for i, x := range tr {
	if x != vec32.MISSING_DATA_SENTINEL {
	counts[i] += 1
	nonMissing += 1
	}
	}
	}
	// If there are no matches then we might be done.
	if nonMissing == 0 {
	numStepsNoData += 1
	}
	if numStepsNoData > NEW_N_MAX_SEARCH {
	break
	}

	ret.ParamSet.AddParamSet(df.ParamSet)

	// For each commit that has data, copy the data from df into ret.
	// Move backwards down the trace since we are building the result from 'end' backwards.
	for i := len(counts) - 1; i >= 0; i-- {
	if counts[i] > 0 {
	ret.Header = append([]*dataframe.ColumnHeader{df.Header[i]}, ret.Header...)
	for key, sourceTrace := range df.TraceSet {
	if _, ok := ret.TraceSet[key]; !ok {
	ret.TraceSet[key] = vec32.New(int(n))
	}
	ret.TraceSet[key][n-1-total] = sourceTrace[i]
	}
	total += 1

	// If we've added enough commits to ret then we are done.
	if total == n {
	break
	}
	}
	}

	sklog.Infof("Total: %d Steps: %d NumStepsNoData: %d", total, steps, numStepsNoData)

	if total == n {
	break
	}

	if progress != nil {
	progress(steps, steps+1)
	}
	steps += 1

	endIndex -= b.tileSize
	beginIndex -= b.tileSize
	if endIndex < 0 {
	break
	}
	if beginIndex < 0 {
	beginIndex = 0
	}

	}

	if total < n {
	// Trim down the traces so they are the same length as ret.Header.
	for key, tr := range ret.TraceSet {
	ret.TraceSet[key] = tr[n-total:]
	}
	}

	return ret, nil
	}

	// See DataFrameBuilder.
	func (b builder) PreflightQuery(ctx context.Context, end time.Time, q query.Query) (int64, paramtools.ParamSet, error) {
	var count int64
	ps := paramtools.ParamSet{}

	tileKey, err := b.store.GetLatestTile()
	if err != nil {
	return -1, nil, err
	}

	if q.Empty() {
	// If the query is empty then we have a shortcut for building the
	// ParamSet by just using the OPS. In that case we only need to count
	// encodedKeys to get the count.
	for i := 0; i < 2; i++ {
	ops, err := b.store.GetOrderedParamSet(ctx, tileKey)
	if err != nil {
	return -1, nil, err
	}
	ps.AddParamSet(ops.ParamSet)
	tileKey = tileKey.PrevTile()
	}

	count, err = b.store.TraceCount(ctx, tileKey)
	if err != nil {
	return -1, nil, err
	}

	} else {
	// Since the query isn't empty we'll have to run a partial query
	// to build the ParamSet. Do so over the two most recent tiles.

	// Record the OPS for the first tile.
	opsOne, err := b.store.GetOrderedParamSet(ctx, tileKey)
	if err != nil {
	return -1, nil, err
	}
	// Count the matches and sum the params in the first tile.
	out, err := b.store.QueryTracesIDOnlyByIndex(ctx, tileKey, q)
	if err != nil {
	return -1, nil, fmt.Errorf("Failed to query traces: %s", err)
	}
	var tileOneCount int64
	for p := range out {
	tileOneCount++
	ps.AddParams(p)
	}

	// Now move to the previous tile.
	tileKey = tileKey.PrevTile()

	// Record the OPS for the second tile.
	opsTwo, err := b.store.GetOrderedParamSet(ctx, tileKey)
	if err != nil {
	return -1, nil, err
	}

	// Count the matches and sum the params in the second tile.
	out, err = b.store.QueryTracesIDOnlyByIndex(ctx, tileKey, q)
	if err != nil {
	return -1, nil, fmt.Errorf("Failed to query traces: %s", err)
	}
	var tileTwoCount int64
	for p := range out {
	tileTwoCount++
	ps.AddParams(p)
	}

	// Use the larger of the two counts as our result.
	if tileOneCount > tileTwoCount {
	count = tileOneCount
	} else {
	count = tileTwoCount
	}

	// Use the larger of the two OPSs to work with.
	var ops *paramtools.OrderedParamSet
	if opsOne.ParamSet.Size() > opsTwo.ParamSet.Size() {
	ops = opsOne
	} else {
	ops = opsTwo
	}

	// Now we have the ParamSet that corresponds to the query, but for each
	// key in the query we need to go back and put in all the values that
	// appear for that key since the user can make more selections in that
	// key.
	queryPlan, err := q.QueryPlan(ops)
	if err != nil {
	return -1, nil, err
	}
	for key := range queryPlan {
	ps[key] = ops.ParamSet[key]
	}
	}

	ps.Normalize()
	return count, ps, nil
	}

	// Validate that the concrete bttsDataFrameBuilder faithfully implements the DataFrameBuidler interface.
	var _ dataframe.DataFrameBuilder = (*builder)(nil)