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

2 files changed, 15 insertions, 483 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,
diff --git a/pkg/cover/heatmap_test.go b/pkg/cover/heatmap_test.go
index 60660b0a0..6a88ffee2 100644
--- a/pkg/cover/heatmap_test.go
+++ b/pkg/cover/heatmap_test.go
@@ -4,264 +4,24 @@
 package cover
 
 import (
-	"context"
 	"testing"
 	"time"
 
 	"cloud.google.com/go/civil"
 	"github.com/google/syzkaller/pkg/coveragedb"
-	"github.com/google/syzkaller/pkg/coveragedb/mocks"
-	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"google.golang.org/api/iterator"
 )
 
-func TestFilesCoverageWithDetails(t *testing.T) {
-	period, _ := coveragedb.MakeTimePeriod(
-		civil.Date{Year: 2025, Month: 1, Day: 1},
-		"day")
-	tests := []struct {
-		name       string
-		scope      *SelectScope
-		client     func() spannerclient.SpannerClient
-		onlyUnique bool
-		want       []*fileCoverageWithDetails
-		wantErr    bool
-	}{
-		{
-			name:    "empty scope",
-			scope:   &SelectScope{},
-			want:    nil,
-			wantErr: false,
-		},
-		{
-			name: "single day, no filters, empty DB => no coverage",
-			scope: &SelectScope{
-				Ns:      "upstream",
-				Periods: []coveragedb.TimePeriod{period},
-			},
-			client:  func() spannerclient.SpannerClient { return emptyCoverageDBFixture(t, 1) },
-			want:    nil,
-			wantErr: false,
-		},
-		{
-			name: "single day, unique coverage, empty DB => no coverage",
-			scope: &SelectScope{
-				Ns:      "upstream",
-				Periods: []coveragedb.TimePeriod{period},
-			},
-			client:     func() spannerclient.SpannerClient { return emptyCoverageDBFixture(t, 2) },
-			onlyUnique: true,
-			want:       nil,
-			wantErr:    false,
-		},
-		{
-			name: "single day, unique coverage, empty partial result => 0/3 covered",
-			scope: &SelectScope{
-				Ns:      "upstream",
-				Periods: []coveragedb.TimePeriod{period},
-			},
-			client: func() spannerclient.SpannerClient {
-				return fullCoverageDBFixture(
-					t,
-					[]*FileCoverageWithLineInfo{
-						{
-							fileCoverageWithDetails: fileCoverageWithDetails{
-								Filepath:     "file1",
-								Instrumented: 3,
-								Covered:      3,
-							},
-							LinesInstrumented: []int64{1, 2, 3},
-							HitCounts:         []int64{1, 1, 1},
-						},
-					},
-					nil)
-			},
-			onlyUnique: true,
-			want: []*fileCoverageWithDetails{
-				{
-					Filepath:     "file1",
-					Instrumented: 3,
-					Covered:      0,
-					TimePeriod:   period,
-				},
-			},
-			wantErr: false,
-		},
-		{
-			name: "single day, unique coverage, full result match => 3/3 covered",
-			scope: &SelectScope{
-				Ns:      "upstream",
-				Periods: []coveragedb.TimePeriod{period},
-			},
-			client: func() spannerclient.SpannerClient {
-				return fullCoverageDBFixture(
-					t,
-					[]*FileCoverageWithLineInfo{
-						{
-							fileCoverageWithDetails: fileCoverageWithDetails{
-								Filepath:     "file1",
-								Instrumented: 3,
-								Covered:      3,
-							},
-							LinesInstrumented: []int64{1, 2, 3},
-							HitCounts:         []int64{1, 1, 1},
-						},
-					},
-					[]*FileCoverageWithLineInfo{
-						{
-							fileCoverageWithDetails: fileCoverageWithDetails{
-								Filepath:     "file1",
-								Instrumented: 3,
-								Covered:      3,
-							},
-							LinesInstrumented: []int64{1, 2, 3},
-							HitCounts:         []int64{1, 1, 1},
-						},
-					})
-			},
-			onlyUnique: true,
-			want: []*fileCoverageWithDetails{
-				{
-					Filepath:     "file1",
-					Instrumented: 3,
-					Covered:      3,
-					TimePeriod:   period,
-				},
-			},
-			wantErr: false,
-		},
-		{
-			name: "single day, unique coverage, partial result match => 3/5 covered",
-			scope: &SelectScope{
-				Ns:      "upstream",
-				Periods: []coveragedb.TimePeriod{period},
-			},
-			client: func() spannerclient.SpannerClient {
-				return fullCoverageDBFixture(
-					t,
-					[]*FileCoverageWithLineInfo{
-						{
-							fileCoverageWithDetails: fileCoverageWithDetails{
-								Filepath:     "file1",
-								Instrumented: 5,
-								Covered:      5,
-							},
-							LinesInstrumented: []int64{1, 2, 3, 4, 5},
-							HitCounts:         []int64{3, 4, 5, 6, 7},
-						},
-					},
-					[]*FileCoverageWithLineInfo{
-						{
-							fileCoverageWithDetails: fileCoverageWithDetails{
-								Filepath:     "file1",
-								Instrumented: 4,
-								Covered:      3,
-							},
-							LinesInstrumented: []int64{1, 2, 3, 5},
-							HitCounts:         []int64{3, 0, 5, 7},
-						},
-					})
-			},
-			onlyUnique: true,
-			want: []*fileCoverageWithDetails{
-				{
-					Filepath:     "file1",
-					Instrumented: 5,
-					Covered:      3,
-					TimePeriod:   period,
-				},
-			},
-			wantErr: false,
-		},
-	}
-	for _, test := range tests {
-		t.Run(test.name, func(t *testing.T) {
-			var testClient spannerclient.SpannerClient
-			if test.client != nil {
-				testClient = test.client()
-			}
-			got, gotErr := filesCoverageWithDetails(
-				context.Background(),
-				testClient, test.scope, test.onlyUnique)
-			if test.wantErr {
-				assert.Error(t, gotErr)
-			} else {
-				assert.NoError(t, gotErr)
-			}
-			assert.Equal(t, test.want, got)
-		})
-	}
-}
-
-func emptyCoverageDBFixture(t *testing.T, times int) spannerclient.SpannerClient {
-	mRowIterator := mocks.NewRowIterator(t)
-	mRowIterator.On("Stop").Return().Times(times)
-	mRowIterator.On("Next").
-		Return(nil, iterator.Done).Times(times)
-
-	mTran := mocks.NewReadOnlyTransaction(t)
-	mTran.On("Query", mock.Anything, mock.Anything).
-		Return(mRowIterator).Times(times)
-
-	m := mocks.NewSpannerClient(t)
-	m.On("Single").
-		Return(mTran).Times(times)
-	return m
-}
-
-func fullCoverageDBFixture(
-	t *testing.T, full, partial []*FileCoverageWithLineInfo,
-) spannerclient.SpannerClient {
-	mPartialTran := mocks.NewReadOnlyTransaction(t)
-	mPartialTran.On("Query", mock.Anything, mock.Anything).
-		Return(newRowIteratorMock(t, partial)).Once()
-
-	mFullTran := mocks.NewReadOnlyTransaction(t)
-	mFullTran.On("Query", mock.Anything, mock.Anything).
-		Return(newRowIteratorMock(t, full)).Once()
-
-	m := mocks.NewSpannerClient(t)
-	m.On("Single").
-		Return(mPartialTran).Once()
-	m.On("Single").
-		Return(mFullTran).Once()
-	return m
-}
-
-func newRowIteratorMock(t *testing.T, events []*FileCoverageWithLineInfo,
-) *mocks.RowIterator {
-	m := mocks.NewRowIterator(t)
-	m.On("Stop").Once().Return()
-	for _, item := range events {
-		mRow := mocks.NewRow(t)
-		mRow.On("ToStruct", mock.Anything).
-			Run(func(args mock.Arguments) {
-				arg := args.Get(0).(*FileCoverageWithLineInfo)
-				*arg = *item
-			}).
-			Return(nil).Once()
-
-		m.On("Next").
-			Return(mRow, nil).Once()
-	}
-
-	m.On("Next").
-		Return(nil, iterator.Done).Once()
-	return m
-}
-
 func TestFilesCoverageToTemplateData(t *testing.T) {
 	tests := []struct {
 		name      string
-		input     []*fileCoverageWithDetails
+		input     []*coveragedb.FileCoverageWithDetails
 		hideEmpty bool
 		want      *templateHeatmap
 	}{
 		{
 			name:  "empty input",
-			input: []*fileCoverageWithDetails{},
+			input: []*coveragedb.FileCoverageWithDetails{},
 			want: &templateHeatmap{
 				Root: &templateHeatmapRow{
 					IsDir: true,
@@ -270,7 +30,7 @@ func TestFilesCoverageToTemplateData(t *testing.T) {
 		},
 		{
 			name: "single file",
-			input: []*fileCoverageWithDetails{
+			input: []*coveragedb.FileCoverageWithDetails{
 				{
 					Filepath:     "file1",
 					Instrumented: 1,
@@ -309,7 +69,7 @@ func TestFilesCoverageToTemplateData(t *testing.T) {
 		},
 		{
 			name: "tree data",
-			input: []*fileCoverageWithDetails{
+			input: []*coveragedb.FileCoverageWithDetails{
 				{
 					Filepath:     "dir/file2",
 					Instrumented: 1,
@@ -385,7 +145,7 @@ func TestFilesCoverageToTemplateData(t *testing.T) {
 		{
 			name:      "hide empty",
 			hideEmpty: true,
-			input: []*fileCoverageWithDetails{
+			input: []*coveragedb.FileCoverageWithDetails{
 				{
 					Filepath:     "file1",
 					Instrumented: 1,