pkg/spanner/coveragedb: move package to pkg/coveragedb

3 files changed, 560 insertions, 0 deletions

diff --git a/pkg/coveragedb/spanner.go b/pkg/coveragedb/spanner.go
new file mode 100644
index 000000000..e5f7fb2a4
--- /dev/null
+++ b/pkg/coveragedb/spanner.go
@@ -0,0 +1,184 @@
+// Copyright 2024 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"
+	"fmt"
+	"time"
+
+	"cloud.google.com/go/civil"
+	"cloud.google.com/go/spanner"
+	"github.com/google/syzkaller/pkg/subsystem"
+	_ "github.com/google/syzkaller/pkg/subsystem/lists"
+	"github.com/google/uuid"
+	"google.golang.org/api/iterator"
+)
+
+type FilesRecord struct {
+	Session      string
+	FilePath     string
+	Instrumented int64
+	Covered      int64
+}
+
+type FileSubsystems struct {
+	Namespace  string
+	FilePath   string
+	Subsystems []string
+}
+
+type HistoryRecord struct {
+	Session   string
+	Time      time.Time
+	Namespace string
+	Repo      string
+	Commit    string
+	Duration  int64
+	DateTo    civil.Date
+	TotalRows int64
+}
+
+func NewClient(ctx context.Context, projectID string) (*spanner.Client, error) {
+	database := "projects/" + projectID + "/instances/syzbot/databases/coverage"
+	return spanner.NewClient(ctx, database)
+}
+
+type Coverage struct {
+	Instrumented int64
+	Covered      int64
+}
+
+func SaveMergeResult(ctx context.Context, projectID string, covMap map[string]*Coverage,
+	template *HistoryRecord, totalRows int64, sss []*subsystem.Subsystem) error {
+	client, err := NewClient(ctx, projectID)
+	if err != nil {
+		return fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
+	}
+	defer client.Close()
+
+	ssMatcher := subsystem.MakePathMatcher(sss)
+	ssCache := make(map[string][]string)
+
+	session := uuid.New().String()
+	mutations := []*spanner.Mutation{}
+	for filePath, record := range covMap {
+		mutations = append(mutations, fileRecordMutation(session, filePath, record))
+		subsystems := fileSubsystems(filePath, ssMatcher, ssCache)
+		mutations = append(mutations, fileSubsystemsMutation(template.Namespace, filePath, subsystems))
+		// 80k mutations is a DB limit. 4 fields * 2k records is apx 8k mutations
+		// let keep this value 10x lower to have a room for indexes
+		// indexes update are also counted
+		if len(mutations) > 2000 {
+			if _, err = client.Apply(ctx, mutations); err != nil {
+				return fmt.Errorf("failed to spanner.Apply(inserts): %s", err.Error())
+			}
+			mutations = nil
+		}
+	}
+	mutations = append(mutations, historyMutation(session, template, totalRows))
+	if _, err = client.Apply(ctx, mutations); err != nil {
+		return fmt.Errorf("failed to spanner.Apply(inserts): %s", err.Error())
+	}
+	return nil
+}
+
+func historyMutation(session string, template *HistoryRecord, totalRows int64) *spanner.Mutation {
+	historyInsert, err := spanner.InsertOrUpdateStruct("merge_history", &HistoryRecord{
+		Session:   session,
+		Time:      time.Now(),
+		Namespace: template.Namespace,
+		Repo:      template.Repo,
+		Commit:    template.Commit,
+		Duration:  template.Duration,
+		DateTo:    template.DateTo,
+		TotalRows: totalRows,
+	})
+	if err != nil {
+		panic(fmt.Sprintf("failed to spanner.InsertStruct(): %s", err.Error()))
+	}
+	return historyInsert
+}
+
+func fileRecordMutation(session, filePath string, record *Coverage) *spanner.Mutation {
+	insert, err := spanner.InsertOrUpdateStruct("files", &FilesRecord{
+		Session:      session,
+		FilePath:     filePath,
+		Instrumented: record.Instrumented,
+		Covered:      record.Covered,
+	})
+	if err != nil {
+		panic(fmt.Sprintf("failed to fileRecordMutation(): %s", err.Error()))
+	}
+	return insert
+}
+
+func fileSubsystemsMutation(ns, filePath string, subsystems []string) *spanner.Mutation {
+	insert, err := spanner.InsertOrUpdateStruct("file_subsystems", &FileSubsystems{
+		Namespace:  ns,
+		FilePath:   filePath,
+		Subsystems: subsystems,
+	})
+	if err != nil {
+		panic(fmt.Sprintf("failed to fileSubsystemsMutation(): %s", err.Error()))
+	}
+	return insert
+}
+
+func fileSubsystems(filePath string, ssMatcher *subsystem.PathMatcher, ssCache map[string][]string) []string {
+	sss, cached := ssCache[filePath]
+	if !cached {
+		for _, match := range ssMatcher.Match(filePath) {
+			sss = append(sss, match.Name)
+		}
+		ssCache[filePath] = sss
+	}
+	return sss
+}
+
+func NsDataMerged(ctx context.Context, projectID, ns string) ([]TimePeriod, []int64, error) {
+	client, err := NewClient(ctx, projectID)
+	if err != nil {
+		return nil, nil, fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
+	}
+	defer client.Close()
+
+	stmt := spanner.Statement{
+		SQL: `
+			select
+				dateto,
+				duration as days,
+				totalrows
+			from merge_history
+			where
+				namespace=$1`,
+		Params: map[string]interface{}{
+			"p1": ns,
+		},
+	}
+	iter := client.Single().Query(ctx, stmt)
+	defer iter.Stop()
+	var periods []TimePeriod
+	var totalRows []int64
+	for {
+		row, err := iter.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to iter.Next() spanner DB: %w", err)
+		}
+		var r struct {
+			Days      int64
+			DateTo    civil.Date
+			TotalRows int64
+		}
+		if err = row.ToStruct(&r); err != nil {
+			return nil, nil, fmt.Errorf("failed to row.ToStruct() spanner DB: %w", err)
+		}
+		periods = append(periods, TimePeriod{DateTo: r.DateTo, Days: int(r.Days)})
+		totalRows = append(totalRows, r.TotalRows)
+	}
+	return periods, totalRows, nil
+}
diff --git a/pkg/coveragedb/time_period.go b/pkg/coveragedb/time_period.go
new file mode 100644
index 000000000..463ee6c41
--- /dev/null
+++ b/pkg/coveragedb/time_period.go
@@ -0,0 +1,113 @@
+// Copyright 2024 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 (
+	"sort"
+
+	"cloud.google.com/go/civil"
+)
+
+type TimePeriod struct {
+	DateTo civil.Date
+	Days   int
+}
+
+type periodOps interface {
+	isValidPeriod(p TimePeriod) bool
+	lastPeriodDate(d civil.Date) civil.Date
+	pointedPeriodDays(d civil.Date) int
+}
+
+type DayPeriodOps struct{}
+
+func (dpo *DayPeriodOps) lastPeriodDate(d civil.Date) civil.Date {
+	return d
+}
+
+func (dpo *DayPeriodOps) isValidPeriod(p TimePeriod) bool {
+	return p.Days == 1
+}
+
+func (dpo *DayPeriodOps) pointedPeriodDays(d civil.Date) int {
+	return 1
+}
+
+type MonthPeriodOps struct{}
+
+func (m *MonthPeriodOps) lastPeriodDate(d civil.Date) civil.Date {
+	d.Day = 1
+	d = d.AddDays(32)
+	d.Day = 1
+	return d.AddDays(-1)
+}
+
+func (m *MonthPeriodOps) isValidPeriod(p TimePeriod) bool {
+	lmd := m.lastPeriodDate(p.DateTo)
+	return lmd == p.DateTo && p.Days == lmd.Day
+}
+
+func (m *MonthPeriodOps) pointedPeriodDays(d civil.Date) int {
+	return m.lastPeriodDate(d).Day
+}
+
+type QuarterPeriodOps struct{}
+
+func (q *QuarterPeriodOps) isValidPeriod(p TimePeriod) bool {
+	lmd := q.lastPeriodDate(p.DateTo)
+	return lmd == p.DateTo && p.Days == q.pointedPeriodDays(lmd)
+}
+
+func (q *QuarterPeriodOps) lastPeriodDate(d civil.Date) civil.Date {
+	d.Month = ((d.Month-1)/3)*3 + 3
+	d.Day = 1
+	return (&MonthPeriodOps{}).lastPeriodDate(d)
+}
+
+func (q *QuarterPeriodOps) pointedPeriodDays(d civil.Date) int {
+	d = q.lastPeriodDate(d)
+	d.Day = 1
+	res := 0
+	for i := 0; i < 3; i++ {
+		res += (&MonthPeriodOps{}).pointedPeriodDays(d)
+		d.Month--
+	}
+	return res
+}
+
+func PeriodsToMerge(srcDates, mergedPeriods []TimePeriod, srcRows, mergedRows []int64, ops periodOps) []TimePeriod {
+	periodRows := map[civil.Date]int64{}
+	for i, srcDate := range srcDates {
+		periodID := ops.lastPeriodDate(srcDate.DateTo)
+		periodRows[periodID] += srcRows[i]
+	}
+	for i, period := range mergedPeriods {
+		if !ops.isValidPeriod(period) {
+			continue
+		}
+		mergerPeriodID := period.DateTo
+		if rowsAvailable, ok := periodRows[mergerPeriodID]; ok && rowsAvailable == mergedRows[i] {
+			delete(periodRows, mergerPeriodID)
+		}
+	}
+	periods := []TimePeriod{}
+	for periodEndDate := range periodRows {
+		periods = append(periods,
+			TimePeriod{DateTo: periodEndDate, Days: ops.pointedPeriodDays(periodEndDate)})
+	}
+	sort.Slice(periods, func(i, j int) bool {
+		return periods[i].DateTo.After(periods[j].DateTo)
+	})
+	return periods
+}
+
+func AtMostNLatestPeriods(periods []TimePeriod, n int) []TimePeriod {
+	sort.Slice(periods, func(i, j int) bool {
+		return periods[i].DateTo.After(periods[j].DateTo)
+	})
+	if len(periods) <= n {
+		return periods
+	}
+	return periods[:n]
+}
diff --git a/pkg/coveragedb/time_period_test.go b/pkg/coveragedb/time_period_test.go
new file mode 100644
index 000000000..85fb4b407
--- /dev/null
+++ b/pkg/coveragedb/time_period_test.go
@@ -0,0 +1,263 @@
+// Copyright 2024 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 (
+	"testing"
+	"time"
+
+	"cloud.google.com/go/civil"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDayPeriodOps(t *testing.T) {
+	ops := &DayPeriodOps{}
+	d := civil.Date{Year: 2024, Month: time.February, Day: 20}
+	goodPeriod := TimePeriod{DateTo: d, Days: 1}
+	badPeriod := TimePeriod{DateTo: d, Days: 2}
+
+	assert.Equal(t, "2024-02-20", ops.lastPeriodDate(d).String())
+
+	assert.True(t, ops.isValidPeriod(goodPeriod))
+	assert.False(t, ops.isValidPeriod(badPeriod))
+
+	assert.Equal(t, 1, ops.pointedPeriodDays(d))
+}
+
+func TestMonthPeriodOps(t *testing.T) {
+	ops := &MonthPeriodOps{}
+	midMonthDate := civil.Date{Year: 2024, Month: time.February, Day: 20}
+	goodPeriod := TimePeriod{DateTo: midMonthDate, Days: 29}
+	goodPeriod.DateTo.Day = goodPeriod.Days
+	badPeriod1 := goodPeriod
+	badPeriod1.DateTo.Day--
+	badPeriod2 := goodPeriod
+	badPeriod2.Days--
+
+	assert.Equal(t, "2024-02-29", ops.lastPeriodDate(midMonthDate).String())
+
+	assert.True(t, ops.isValidPeriod(goodPeriod))
+	assert.False(t, ops.isValidPeriod(badPeriod1))
+	assert.False(t, ops.isValidPeriod(badPeriod2))
+
+	assert.Equal(t, 29, ops.pointedPeriodDays(midMonthDate))
+}
+
+func TestQuarterPeriodOps(t *testing.T) {
+	ops := &QuarterPeriodOps{}
+	midQuarterDate := civil.Date{Year: 2024, Month: time.February, Day: 20}
+	goodPeriod := TimePeriod{DateTo: midQuarterDate, Days: 31 + 29 + 31}
+	goodPeriod.DateTo.Month = time.March
+	goodPeriod.DateTo.Day = 31
+	badPeriod1 := goodPeriod
+	badPeriod1.DateTo.Day--
+	badPeriod2 := goodPeriod
+	badPeriod2.Days--
+
+	assert.Equal(t, "2024-03-31", ops.lastPeriodDate(midQuarterDate).String())
+
+	assert.True(t, ops.isValidPeriod(goodPeriod))
+	assert.False(t, ops.isValidPeriod(badPeriod1))
+	assert.False(t, ops.isValidPeriod(badPeriod2))
+
+	assert.Equal(t, 31+29+31, ops.pointedPeriodDays(midQuarterDate))
+}
+
+func TestPeriodsToMerge(t *testing.T) {
+	sampleDays := []TimePeriod{
+		makeTimePeriod("2024-04-01", 1),
+		makeTimePeriod("2024-04-02", 1),
+		makeTimePeriod("2024-05-03", 1),
+		makeTimePeriod("2024-05-04", 1),
+		makeTimePeriod("2024-06-05", 1),
+		makeTimePeriod("2024-06-06", 1),
+	}
+	sampleRows := []int64{1, 2, 4, 8, 16, 32}
+
+	tests := []struct {
+		name          string
+		srcDates      []TimePeriod
+		srcRows       []int64
+		mergedPeriods []TimePeriod
+		mergedRows    []int64
+		ops           periodOps
+		expected      []TimePeriod
+	}{
+		{
+			name:     "days/all_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-04-01", 1),
+				makeTimePeriod("2024-04-02", 1),
+				makeTimePeriod("2024-05-03", 1),
+				makeTimePeriod("2024-05-04", 1),
+				makeTimePeriod("2024-06-05", 1),
+				makeTimePeriod("2024-06-06", 1),
+			},
+			mergedRows: []int64{1, 2, 4, 8, 16, 32},
+			ops:        &DayPeriodOps{},
+			expected:   []TimePeriod{},
+		},
+		{
+			name:     "days/some_not_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-04-01", 1),
+				makeTimePeriod("2024-05-03", 1),
+				makeTimePeriod("2024-05-04", 1),
+				makeTimePeriod("2024-06-06", 1),
+			},
+			mergedRows: []int64{1, 4, 8, 32},
+			ops:        &DayPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-05", 1),
+				makeTimePeriod("2024-04-02", 1),
+			},
+		},
+		{
+			name:     "days/some_partially_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-04-01", 1),
+				makeTimePeriod("2024-04-02", 1),
+				makeTimePeriod("2024-05-03", 1),
+				makeTimePeriod("2024-05-04", 1),
+				makeTimePeriod("2024-06-05", 1),
+				makeTimePeriod("2024-06-06", 1),
+			},
+			mergedRows: []int64{1, 2, 1, 8, 16, 1},
+			ops:        &DayPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-06", 1),
+				makeTimePeriod("2024-05-03", 1),
+			},
+		},
+		{
+			name:     "months/all_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-04-30", 30),
+				makeTimePeriod("2024-05-31", 31),
+				makeTimePeriod("2024-06-30", 30),
+			},
+			mergedRows: []int64{3, 12, 48},
+			ops:        &MonthPeriodOps{},
+			expected:   []TimePeriod{},
+		},
+		{
+			name:     "months/some_not_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-04-30", 30),
+				makeTimePeriod("2024-05-31", 31),
+			},
+			mergedRows: []int64{3, 12},
+			ops:        &MonthPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30),
+			},
+		},
+		{
+			name:     "months/some_partially_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-04-30", 30),
+				makeTimePeriod("2024-05-31", 31),
+				makeTimePeriod("2024-06-30", 30),
+			},
+			mergedRows: []int64{1, 12, 1},
+			ops:        &MonthPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30),
+				makeTimePeriod("2024-04-30", 30),
+			},
+		},
+
+		{
+			name:     "quarter/all_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30+31+30),
+			},
+			mergedRows: []int64{63},
+			ops:        &QuarterPeriodOps{},
+			expected:   []TimePeriod{},
+		},
+		{
+			name:          "quarter/not_merged",
+			srcDates:      sampleDays,
+			srcRows:       sampleRows,
+			mergedPeriods: []TimePeriod{},
+			mergedRows:    []int64{},
+			ops:           &QuarterPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30+31+30),
+			},
+		},
+		{
+			name:     "quarter/partially_merged",
+			srcDates: sampleDays,
+			srcRows:  sampleRows,
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30+31+30),
+			},
+			mergedRows: []int64{60},
+			ops:        &QuarterPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30+31+30),
+			},
+		},
+		{
+			name:     "quarters/not_all_merged_with_invalid_periods",
+			srcDates: append(sampleDays, makeTimePeriod("2024-01-01", 1)),
+			srcRows:  append(sampleRows, 128),
+			mergedPeriods: []TimePeriod{
+				makeTimePeriod("2024-03-31", 31+29+31),
+				makeTimePeriod("2024-06-30", 30+31+30),
+				makeTimePeriod("2024-01-10", 30),
+				makeTimePeriod("2024-01-20", 1),
+			},
+			mergedRows: []int64{128, 60, 1, 1},
+			ops:        &QuarterPeriodOps{},
+			expected: []TimePeriod{
+				makeTimePeriod("2024-06-30", 30+31+30),
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			assert.Equal(t, len(test.srcRows), len(test.srcDates))
+			actual := PeriodsToMerge(test.srcDates, test.mergedPeriods, test.srcRows, test.mergedRows, test.ops)
+			assert.Equal(t, test.expected, actual)
+		})
+	}
+}
+
+func makeTimePeriod(s string, days int) TimePeriod {
+	d, err := civil.ParseDate(s)
+	if err != nil {
+		panic(err.Error())
+	}
+	return TimePeriod{DateTo: d, Days: days}
+}
+
+func TestAtMostNLatestPeriods(t *testing.T) {
+	sampleDays := []TimePeriod{
+		makeTimePeriod("2024-04-01", 1),
+		makeTimePeriod("2024-04-02", 1),
+		makeTimePeriod("2024-05-03", 1),
+		makeTimePeriod("2024-05-04", 1),
+		makeTimePeriod("2024-06-05", 1),
+		makeTimePeriod("2024-06-06", 1),
+	}
+	assert.Equal(t, []TimePeriod{makeTimePeriod("2024-06-06", 1)}, AtMostNLatestPeriods(sampleDays, 1))
+	assert.Equal(t, sampleDays, AtMostNLatestPeriods(sampleDays, 100))
+}