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

2 files changed, 479 insertions, 0 deletions

diff --git a/pkg/coveragedb/coveragedb.go b/pkg/coveragedb/coveragedb.go
index 5c60958fa..4c51df515 100644
--- a/pkg/coveragedb/coveragedb.go
+++ b/pkg/coveragedb/coveragedb.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"maps"
 	"sync/atomic"
 	"time"
 
@@ -378,3 +379,230 @@ func goSpannerDelete(ctx context.Context, batch []spanner.Key, eg *errgroup.Grou
 		return err
 	})
 }
+
+type FileCoverageWithDetails struct {
+	Subsystem    string
+	Filepath     string
+	Instrumented int64
+	Covered      int64
+	TimePeriod   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)
+}
+
+func MakeCovMap(keys, vals []int64) map[int]int64 {
+	res := map[int]int64{}
+	for i, key := range keys {
+		res[int(key)] = vals[i]
+	}
+	return res
+}
+
+type SelectScope struct {
+	Ns        string
+	Subsystem string
+	Manager   string
+	Periods   []TimePeriod
+}
+
+// FilesCoverageWithDetails fetches the data directly from DB. No caching.
+// Flag onlyUnique is quite expensive.
+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
+}
+
+func filesCoverageWithDetailsStmt(ns, subsystem, manager string, timePeriod 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 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 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
+}
diff --git a/pkg/coveragedb/coveragedb_test.go b/pkg/coveragedb/coveragedb_test.go
new file mode 100644
index 000000000..406d7d7a2
--- /dev/null
+++ b/pkg/coveragedb/coveragedb_test.go
@@ -0,0 +1,251 @@
+// Copyright 2025 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package coveragedb
+
+import (
+	"context"
+	"testing"
+
+	"cloud.google.com/go/civil"
+	"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, _ := 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: []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: []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: []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: []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: []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
+}