dashboard/app: show manager unique coverage

1. Make heatmap testable, move out the spanner client instantiation.
2. Generate spannerdb.ReadOnlyTransaction mocks.
2. Generate spannerdb.RowIterator mocks.
3. Prepare spannerdb fixture.

dashboard/app/coverage.go

@@ -14,11 +14,13 @@ import (
 	"cloud.google.com/go/civil"
 	"github.com/google/syzkaller/pkg/cover"
 	"github.com/google/syzkaller/pkg/coveragedb"
+	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	"github.com/google/syzkaller/pkg/covermerger"
 	"github.com/google/syzkaller/pkg/validator"
 )
 
-type funcStyleBodyJS func(ctx context.Context, projectID string, scope *cover.SelectScope, sss, managers []string,
+type funcStyleBodyJS func(
+	ctx context.Context, client spannerclient.SpannerClient, scope *cover.SelectScope, onlyUnique bool, sss, managers []string,
 ) (template.CSS, template.HTML, template.HTML, error)
 
 func handleCoverageHeatmap(c context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -71,16 +73,24 @@ func handleHeatmap(c context.Context, w http.ResponseWriter, r *http.Request, f
 	slices.Sort(managers)
 	slices.Sort(subsystems)
 
+	onlyUnique := r.FormValue("only-unique") == "1"
+
+	spannerClient, err := spannerclient.NewClient(c, "syzkaller")
+	if err != nil {
+		return fmt.Errorf("spanner.NewClient: %s", err.Error())
+	}
+	defer spannerClient.Close()
+
 	var style template.CSS
 	var body, js template.HTML
-	if style, body, js, err = f(c, "syzkaller",
+	if style, body, js, err = f(c, spannerClient,
 		&cover.SelectScope{
 			Ns:        hdr.Namespace,
 			Subsystem: ss,
 			Manager:   manager,
 			Periods:   periods,
 		},
-		subsystems, managers); err != nil {
+		onlyUnique, subsystems, managers); err != nil {
 		return fmt.Errorf("failed to generate heatmap: %w", err)
 	}
 	return serveTemplate(w, "custom_content.html", struct {

pkg/cover/heatmap.go

@@ -17,6 +17,7 @@ import (
 	"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"
 )
 
@@ -115,6 +116,12 @@ type fileCoverageWithDetails struct {
 	Subsystems   []string
 }
 
+type fileCoverageWithLineInfo struct {
+	fileCoverageWithDetails
+	LinesInstrumented []int64
+	HitCounts         []int64
+}
+
 type pageColumnTarget struct {
 	TimePeriod coveragedb.TimePeriod
 	Commit     string
@@ -157,25 +164,25 @@ func filesCoverageToTemplateData(fCov []*fileCoverageWithDetails) *templateHeatm
 	return &res
 }
 
-func filesCoverageWithDetailsStmt(ns, subsystem, manager string, timePeriod coveragedb.TimePeriod) spanner.Statement {
+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
-  commit,
-  instrumented,
-  covered,
-  files.filepath,
-  subsystems
+		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`,
+  merge_history.namespace=$1 and dateto=$2 and duration=$3 and manager=$4
+order by files.filepath`,
 		Params: map[string]interface{}{
 			"p1": ns,
 			"p2": timePeriod.DateTo,
@@ -190,34 +197,134 @@ where
 	return stmt
 }
 
-func filesCoverageWithDetails(ctx context.Context, projectID string, scope *SelectScope,
-) ([]*fileCoverageWithDetails, error) {
-	client, err := spannerclient.NewClient(ctx, projectID)
+func readCoverage(iterManager spannerclient.RowIterator) ([]*fileCoverageWithDetails, error) {
+	res := []*fileCoverageWithDetails{}
+	ch := make(chan *fileCoverageWithDetails)
+	var err error
+	go func() {
+		defer close(ch)
+		err = readIterToChan(iterManager, ch)
+	}()
+	for fc := range ch {
+		res = append(res, fc)
+	}
 	if err != nil {
-		return nil, fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
+		return nil, fmt.Errorf("readIterToChan: %w", err)
 	}
-	defer client.Close()
+	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 := errgroup.Group{}
+	fullCh := make(chan *fileCoverageWithLineInfo)
+	eg.Go(func() error {
+		defer close(fullCh)
+		return readIterToChan(full, fullCh)
+	})
+	partCh := make(chan *fileCoverageWithLineInfo)
+	eg.Go(func() error {
+		defer close(partCh)
+		return readIterToChan(mgr, partCh)
+	})
 	res := []*fileCoverageWithDetails{}
-	for _, timePeriod := range scope.Periods {
-		stmt := filesCoverageWithDetailsStmt(scope.Ns, scope.Subsystem, scope.Manager, timePeriod)
-		iter := client.Single().Query(ctx, stmt)
-		defer iter.Stop()
-		for {
-			row, err := iter.Next()
-			if err == iterator.Done {
-				break
+	eg.Go(func() error {
+		partCov, ok := <-partCh
+		if !ok {
+			// no unique coverage exists because manager coverage is empty
+			return nil
+		}
+		for fullCov := range fullCh {
+			for partCov.Filepath < fullCov.Filepath {
+				partCov, ok = <-partCh
+				if !ok {
+					// all the file pairs found, exit
+					return nil
+				}
+			}
+			if partCov.Filepath == fullCov.Filepath {
+				if len(partCov.LinesInstrumented) != len(fullCov.LinesInstrumented) ||
+					len(partCov.HitCounts) != len(fullCov.HitCounts) ||
+					partCov.Commit != fullCov.Commit {
+					return fmt.Errorf("db record for file %s don't match", fullCov.Filepath)
+				}
+				res = append(res, uniqCoverage(fullCov, partCov))
 			}
+		}
+		return nil
+	})
+	if err := eg.Wait(); err != nil {
+		return nil, fmt.Errorf("eg.Wait: %w", err)
+	}
+	return res, nil
+}
+
+func uniqCoverage(full, partial *fileCoverageWithLineInfo) *fileCoverageWithDetails {
+	res := &full.fileCoverageWithDetails // Use Instrumented count from full aggregation.
+	res.Covered = 0                      // We're recalculating only the covered lines.
+	var fullCov map[int64]int64
+	for i, ln := range full.LinesInstrumented {
+		fullCov[ln] = full.HitCounts[i]
+	}
+	for i, ln := range partial.LinesInstrumented {
+		if hitCount, exist := fullCov[ln]; exist && hitCount > 0 && hitCount == partial.HitCounts[i] {
+			res.Covered++
+		}
+	}
+	return res
+}
+
+func readIterToChan[K fileCoverageWithLineInfo | fileCoverageWithDetails](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)
+		}
+		ch <- &r
+	}
+	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
+		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("failed to iter.Next() spanner DB: %w", err)
+				return nil, fmt.Errorf("uniqueFilesCoverageWithDetails: %w", err)
 			}
-			var r fileCoverageWithDetails
-			if err = row.ToStruct(&r); err != nil {
-				return nil, fmt.Errorf("failed to row.ToStruct() spanner DB: %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, &r)
 		}
+		res = append(res, periodRes...)
 	}
 	return res, nil
 }
@@ -252,9 +359,10 @@ type SelectScope struct {
 	Periods   []coveragedb.TimePeriod
 }
 
-func DoHeatMapStyleBodyJS(ctx context.Context, projectID string, scope *SelectScope, sss, managers []string,
+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, projectID, scope)
+	covAndDates, err := filesCoverageWithDetails(ctx, client, scope, onlyUnique)
 	if err != nil {
 		return "", "", "", fmt.Errorf("failed to filesCoverageWithDetails: %w", err)
 	}
@@ -264,9 +372,10 @@ func DoHeatMapStyleBodyJS(ctx context.Context, projectID string, scope *SelectSc
 	return stylesBodyJSTemplate(templData)
 }
 
-func DoSubsystemsHeatMapStyleBodyJS(ctx context.Context, projectID string, scope *SelectScope, sss, managers []string,
+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, projectID, scope)
+	covWithDetails, err := filesCoverageWithDetails(ctx, client, scope, onlyUnique)
 	if err != nil {
 		panic(err)
 	}

pkg/cover/heatmap_test.go

@@ -4,14 +4,79 @@
 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     spannerclient.SpannerClient
+		onlyUnique bool
+		want       []*fileCoverageWithDetails
+		wantErr    bool
+	}{
+		{
+			name:    "empty scope",
+			scope:   &SelectScope{},
+			want:    nil,
+			wantErr: false,
+		},
+		{
+			name: "single day, no filters => no coverage",
+			scope: &SelectScope{
+				Ns:      "upstream",
+				Periods: []coveragedb.TimePeriod{period},
+			},
+			client:  emptyCoverageDBFixture(t),
+			want:    nil,
+			wantErr: false,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got, gotErr := filesCoverageWithDetails(
+				context.Background(),
+				test.client, 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) spannerclient.SpannerClient {
+	mRowIterator := mocks.NewRowIterator(t)
+	mRowIterator.On("Stop").Once()
+	mRowIterator.On("Next").
+		Once().Return(nil, iterator.Done)
+
+	mTran := mocks.NewReadOnlyTransaction(t)
+	mTran.On("Query", mock.Anything, mock.Anything).
+		Once().Return(mRowIterator)
+
+	m := mocks.NewSpannerClient(t)
+	m.On("Single").
+		Once().Return(mTran)
+	return m
+}
+
 func TestFilesCoverageToTemplateData(t *testing.T) {
 	tests := []struct {
 		name  string

pkg/cover/templates/heatmap.html

@@ -139,7 +139,7 @@
           {{ end }}
         </select>
         <br>
-        <input type="checkbox" id="only-unique" name="unique-only" disabled>
+        <input type="checkbox" id="only-unique" name="unique-only">
       </div>
       <br>
       <button id="updateButton">Update</button>

pkg/coveragedb/coveragedb_mock_test.go

@@ -19,6 +19,8 @@ import (
 )
 
 //go:generate ../../tools/mockery.sh --name SpannerClient -r
+//go:generate ../../tools/mockery.sh --name ReadOnlyTransaction -r
+//go:generate ../../tools/mockery.sh --name RowIterator -r
 
 type spannerMockTune func(*testing.T, *mocks.SpannerClient)