pkg/cover: split logic, move some data processing to pkg/coveragedb

1 files changed, 10 insertions, 238 deletions

diff --git a/pkg/cover/heatmap.go b/pkg/cover/heatmap.go
index f02ee132f..375817b7d 100644
--- a/pkg/cover/heatmap.go
+++ b/pkg/cover/heatmap.go
@@ -12,13 +12,10 @@ import (
 	"sort"
 	"strings"
 
-	"cloud.google.com/go/spanner"
 	"github.com/google/syzkaller/pkg/coveragedb"
 	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	_ "github.com/google/syzkaller/pkg/subsystem/lists"
 	"golang.org/x/exp/maps"
-	"golang.org/x/sync/errgroup"
-	"google.golang.org/api/iterator"
 )
 
 type templateHeatmapRow struct {
@@ -117,32 +114,12 @@ func (thm *templateHeatmapRow) prepareDataFor(pageColumns []pageColumnTarget, sk
 	}
 }
 
-type fileCoverageWithDetails struct {
-	Subsystem    string
-	Filepath     string
-	Instrumented int64
-	Covered      int64
-	TimePeriod   coveragedb.TimePeriod `spanner:"-"`
-	Commit       string
-	Subsystems   []string
-}
-
-type FileCoverageWithLineInfo struct {
-	fileCoverageWithDetails
-	LinesInstrumented []int64
-	HitCounts         []int64
-}
-
-func (fc *FileCoverageWithLineInfo) CovMap() map[int]int64 {
-	return MakeCovMap(fc.LinesInstrumented, fc.HitCounts)
-}
-
 type pageColumnTarget struct {
 	TimePeriod coveragedb.TimePeriod
 	Commit     string
 }
 
-func filesCoverageToTemplateData(fCov []*fileCoverageWithDetails, hideEmpty bool) *templateHeatmap {
+func filesCoverageToTemplateData(fCov []*coveragedb.FileCoverageWithDetails, hideEmpty bool) *templateHeatmap {
 	res := templateHeatmap{
 		Root: &templateHeatmapRow{
 			IsDir:        true,
@@ -179,204 +156,6 @@ func filesCoverageToTemplateData(fCov []*fileCoverageWithDetails, hideEmpty bool
 	return &res
 }
 
-func filesCoverageWithDetailsStmt(ns, subsystem, manager string, timePeriod coveragedb.TimePeriod, withLines bool,
-) spanner.Statement {
-	if manager == "" {
-		manager = "*"
-	}
-	selectColumns := "commit, instrumented, covered, files.filepath, subsystems"
-	if withLines {
-		selectColumns += ", linesinstrumented, hitcounts"
-	}
-	stmt := spanner.Statement{
-		SQL: "select " + selectColumns + `
-from merge_history
-  join files
-    on merge_history.session = files.session
-  join file_subsystems
-    on merge_history.namespace = file_subsystems.namespace and files.filepath = file_subsystems.filepath
-where
-  merge_history.namespace=$1 and dateto=$2 and duration=$3 and manager=$4`,
-		Params: map[string]interface{}{
-			"p1": ns,
-			"p2": timePeriod.DateTo,
-			"p3": timePeriod.Days,
-			"p4": manager,
-		},
-	}
-	if subsystem != "" {
-		stmt.SQL += " and $5=ANY(subsystems)"
-		stmt.Params["p5"] = subsystem
-	}
-	stmt.SQL += "\norder by files.filepath"
-	return stmt
-}
-
-func readCoverage(iterManager spannerclient.RowIterator) ([]*fileCoverageWithDetails, error) {
-	res := []*fileCoverageWithDetails{}
-	ch := make(chan *fileCoverageWithDetails)
-	var err error
-	go func() {
-		defer close(ch)
-		err = readIterToChan(context.Background(), iterManager, ch)
-	}()
-	for fc := range ch {
-		res = append(res, fc)
-	}
-	if err != nil {
-		return nil, fmt.Errorf("readIterToChan: %w", err)
-	}
-	return res, nil
-}
-
-// Unique coverage from specific manager is more expensive to get.
-// We get unique coverage comparing manager and total coverage on the AppEngine side.
-func readCoverageUniq(full, mgr spannerclient.RowIterator,
-) ([]*fileCoverageWithDetails, error) {
-	eg, ctx := errgroup.WithContext(context.Background())
-	fullCh := make(chan *FileCoverageWithLineInfo)
-	eg.Go(func() error {
-		defer close(fullCh)
-		return readIterToChan(ctx, full, fullCh)
-	})
-	partCh := make(chan *FileCoverageWithLineInfo)
-	eg.Go(func() error {
-		defer close(partCh)
-		return readIterToChan(ctx, mgr, partCh)
-	})
-	res := []*fileCoverageWithDetails{}
-	eg.Go(func() error {
-		partCov := <-partCh
-		for fullCov := range fullCh {
-			if partCov == nil || partCov.Filepath > fullCov.Filepath {
-				// No pair for the file in full aggregation is available.
-				cov := fullCov.fileCoverageWithDetails
-				cov.Covered = 0
-				res = append(res, &cov)
-				continue
-			}
-			if partCov.Filepath == fullCov.Filepath {
-				if partCov.Commit != fullCov.Commit ||
-					!IsComparable(
-						fullCov.LinesInstrumented, fullCov.HitCounts,
-						partCov.LinesInstrumented, partCov.HitCounts) {
-					return fmt.Errorf("db record for file %s doesn't match", fullCov.Filepath)
-				}
-				resItem := fullCov.fileCoverageWithDetails // Use Instrumented count from full aggregation.
-				resItem.Covered = 0
-				for _, hc := range UniqCoverage(fullCov.CovMap(), partCov.CovMap()) {
-					if hc > 0 {
-						resItem.Covered++
-					}
-				}
-				res = append(res, &resItem)
-				partCov = <-partCh
-				continue
-			}
-			// Partial coverage is a subset of full coverage.
-			// File can't exist only in partial set.
-			return fmt.Errorf("currupted db, file %s can't exist", partCov.Filepath)
-		}
-		return nil
-	})
-	if err := eg.Wait(); err != nil {
-		return nil, fmt.Errorf("eg.Wait: %w", err)
-	}
-	return res, nil
-}
-
-func MakeCovMap(keys, vals []int64) map[int]int64 {
-	res := map[int]int64{}
-	for i, key := range keys {
-		res[int(key)] = vals[i]
-	}
-	return res
-}
-
-func IsComparable(fullLines, fullHitCounts, partialLines, partialHitCounts []int64) bool {
-	if len(fullLines) != len(fullHitCounts) ||
-		len(partialLines) != len(partialHitCounts) ||
-		len(fullLines) < len(partialLines) {
-		return false
-	}
-	fullCov := MakeCovMap(fullLines, fullHitCounts)
-	for iPartial, ln := range partialLines {
-		partialHitCount := partialHitCounts[iPartial]
-		if fullHitCount, fullExist := fullCov[int(ln)]; !fullExist || fullHitCount < partialHitCount {
-			return false
-		}
-	}
-	return true
-}
-
-// Returns partial hitcounts that are the only source of the full hitcounts.
-func UniqCoverage(fullCov, partCov map[int]int64) map[int]int64 {
-	res := maps.Clone(partCov)
-	for ln := range partCov {
-		if partCov[ln] != fullCov[ln] {
-			res[ln] = 0
-		}
-	}
-	return res
-}
-
-func readIterToChan[K FileCoverageWithLineInfo | fileCoverageWithDetails](
-	ctx context.Context, iter spannerclient.RowIterator, ch chan<- *K) error {
-	for {
-		row, err := iter.Next()
-		if err == iterator.Done {
-			break
-		}
-		if err != nil {
-			return fmt.Errorf("iter.Next: %w", err)
-		}
-		var r K
-		if err = row.ToStruct(&r); err != nil {
-			return fmt.Errorf("row.ToStruct: %w", err)
-		}
-		select {
-		case ch <- &r:
-		case <-ctx.Done():
-			return nil
-		}
-	}
-	return nil
-}
-
-func filesCoverageWithDetails(
-	ctx context.Context, client spannerclient.SpannerClient, scope *SelectScope, onlyUnique bool,
-) ([]*fileCoverageWithDetails, error) {
-	var res []*fileCoverageWithDetails
-	for _, timePeriod := range scope.Periods {
-		needLinesDetails := onlyUnique
-		iterManager := client.Single().Query(ctx,
-			filesCoverageWithDetailsStmt(scope.Ns, scope.Subsystem, scope.Manager, timePeriod, needLinesDetails))
-		defer iterManager.Stop()
-
-		var err error
-		var periodRes []*fileCoverageWithDetails
-		if onlyUnique {
-			iterAll := client.Single().Query(ctx,
-				filesCoverageWithDetailsStmt(scope.Ns, scope.Subsystem, "", timePeriod, needLinesDetails))
-			defer iterAll.Stop()
-			periodRes, err = readCoverageUniq(iterAll, iterManager)
-			if err != nil {
-				return nil, fmt.Errorf("uniqueFilesCoverageWithDetails: %w", err)
-			}
-		} else {
-			periodRes, err = readCoverage(iterManager)
-			if err != nil {
-				return nil, fmt.Errorf("readCoverage: %w", err)
-			}
-		}
-		for _, r := range periodRes {
-			r.TimePeriod = timePeriod
-		}
-		res = append(res, periodRes...)
-	}
-	return res, nil
-}
-
 type StyleBodyJS struct {
 	Style template.CSS
 	Body  template.HTML
@@ -400,19 +179,12 @@ func stylesBodyJSTemplate(templData *templateHeatmap,
 		template.HTML(js.Bytes()), nil
 }
 
-type SelectScope struct {
-	Ns        string
-	Subsystem string
-	Manager   string
-	Periods   []coveragedb.TimePeriod
-}
-
 func DoHeatMapStyleBodyJS(
-	ctx context.Context, client spannerclient.SpannerClient, scope *SelectScope, onlyUnique bool, sss, managers []string,
-) (template.CSS, template.HTML, template.HTML, error) {
-	covAndDates, err := filesCoverageWithDetails(ctx, client, scope, onlyUnique)
+	ctx context.Context, client spannerclient.SpannerClient, scope *coveragedb.SelectScope, onlyUnique bool,
+	sss, managers []string) (template.CSS, template.HTML, template.HTML, error) {
+	covAndDates, err := coveragedb.FilesCoverageWithDetails(ctx, client, scope, onlyUnique)
 	if err != nil {
-		return "", "", "", fmt.Errorf("failed to filesCoverageWithDetails: %w", err)
+		return "", "", "", fmt.Errorf("failed to FilesCoverageWithDetails: %w", err)
 	}
 	templData := filesCoverageToTemplateData(covAndDates, onlyUnique)
 	templData.Subsystems = sss
@@ -421,16 +193,16 @@ func DoHeatMapStyleBodyJS(
 }
 
 func DoSubsystemsHeatMapStyleBodyJS(
-	ctx context.Context, client spannerclient.SpannerClient, scope *SelectScope, onlyUnique bool, sss, managers []string,
-) (template.CSS, template.HTML, template.HTML, error) {
-	covWithDetails, err := filesCoverageWithDetails(ctx, client, scope, onlyUnique)
+	ctx context.Context, client spannerclient.SpannerClient, scope *coveragedb.SelectScope, onlyUnique bool,
+	sss, managers []string) (template.CSS, template.HTML, template.HTML, error) {
+	covWithDetails, err := coveragedb.FilesCoverageWithDetails(ctx, client, scope, onlyUnique)
 	if err != nil {
 		panic(err)
 	}
-	var ssCovAndDates []*fileCoverageWithDetails
+	var ssCovAndDates []*coveragedb.FileCoverageWithDetails
 	for _, cwd := range covWithDetails {
 		for _, ssName := range cwd.Subsystems {
-			newRecord := fileCoverageWithDetails{
+			newRecord := coveragedb.FileCoverageWithDetails{
 				Filepath:     cwd.Filepath,
 				Subsystem:    ssName,
 				Instrumented: cwd.Instrumented,