resultdb/internal/services/testmetadataupdator/updator.go - infra/luci/luci-go - Git at Google

 // Copyright 2023 The LUCI Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 // Package testmetadataupdator implements a task which ingest test metadata from invocations.
 package testmetadataupdator

 import (
 	"context"
 	"math"
 	"strings"
 	"time"

 	"cloud.google.com/go/spanner"
 	"go.chromium.org/luci/common/errors"
 	"go.chromium.org/luci/common/proto/mask"
 	"go.chromium.org/luci/resultdb/internal/invocations"
 	"go.chromium.org/luci/resultdb/internal/spanutil"
 	"go.chromium.org/luci/resultdb/internal/testmetadata"
 	"go.chromium.org/luci/resultdb/internal/testresults"
 	"go.chromium.org/luci/resultdb/pbutil"
 	pb "go.chromium.org/luci/resultdb/proto/v1"
 	"go.chromium.org/luci/server/auth/realms"
 	"go.chromium.org/luci/server/span"
 	"golang.org/x/sync/errgroup"
 	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
 )

 // Expiration time of a test metadata.
 // After expiry, update to less metadata information is allowed.
 var metadataExpiry = 48 * time.Hour

 // Maximum number of incoming test metadata we hold in memory.
 var batchSize = 1000

 type testID string

 // testMetadataUpdator updates TestMetadata rows with test results in a set of invocations.
 type testMetadataUpdator struct {
 	sources *pb.Sources
 	invIDs  invocations.IDSet
 	start   time.Time // Used to determine if a test metadata is expired.
 }

 func newUpdator(sources *pb.Sources, invIDs invocations.IDSet, start time.Time) *testMetadataUpdator {
 	return &testMetadataUpdator{sources: sources, invIDs: invIDs, start: start}
 }

 func (u *testMetadataUpdator) run(ctx context.Context) error {
 	for invID := range u.invIDs {
 		// Query test results and use them to update test metadata.
 		batchC := make(chan map[testID]*pb.TestMetadata)
 		// Batch update test metadata.
 		eg, ctx := errgroup.WithContext(ctx)

 		eg.Go(func() error {
 			defer close(batchC)
 			if err := u.queryTestResultMetadata(ctx, invID, batchC); err != nil {
 				return errors.Annotate(err, "query test result metadata, invocation %s", invID).Err()
 			}
 			return nil
 		})

 		eg.Go(func() error {
 			var realm string
 			if err := invocations.ReadColumns(span.Single(ctx), invID, map[string]any{"Realm": &realm}); err != nil {
 				return errors.Annotate(err, "read realm, invocation %s", invID).Err()
 			}
 			if err := u.batchCreateOrUpdateTestMetadata(ctx, realm, batchC); err != nil {
 				return errors.Annotate(err, "create or update test metadata, invocation %s", invID).Err()
 			}
 			return nil
 		})
 		if err := eg.Wait(); err != nil {
 			return err
 		}
 	}
 	return nil
 }

 // queryTestResultMetadata visits all test results in the given invocations.
 // and returns one test metadata with the most metadata information for each test.
 func (u *testMetadataUpdator) queryTestResultMetadata(ctx context.Context, invID invocations.ID, batchC chan<- map[testID]*pb.TestMetadata) error {
 	ctx, cancel := span.ReadOnlyTransaction(ctx)
 	defer cancel()

 	q := testresults.Query{
 		InvocationIDs: invocations.NewIDSet(invID),
 		Mask: mask.MustFromReadMask(&pb.TestResult{},
 			"test_id",
 			"test_metadata",
 		),
 	}
 	deduplicatedMetadata := make(map[testID]*pb.TestMetadata, batchSize)
 	err := q.Run(ctx, func(tr *pb.TestResult) error {
 		existing, ok := deduplicatedMetadata[testID(tr.TestId)]
 		if !ok {
 			if len(deduplicatedMetadata) >= batchSize {
 				select {
 				case <-ctx.Done():
 					return ctx.Err()
 				case batchC <- deduplicatedMetadata:
 				}
 				deduplicatedMetadata = make(map[testID]*pb.TestMetadata, batchSize)
 			}
 			deduplicatedMetadata[testID(tr.TestId)] = tr.TestMetadata
 			return nil
 		}
 		if fieldExistenceBitField(existing) < fieldExistenceBitField(tr.TestMetadata) {
 			deduplicatedMetadata[testID(tr.TestId)] = tr.TestMetadata
 		}
 		return nil
 	})
 	if err != nil {
 		return err
 	}
 	if len(deduplicatedMetadata) > 0 {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		case batchC <- deduplicatedMetadata:
 		}
 	}
 	return nil
 }

 func (u *testMetadataUpdator) batchCreateOrUpdateTestMetadata(ctx context.Context, realm string, batchC <-chan map[testID]*pb.TestMetadata) error {
 	for testResults := range batchC {
 		if err := u.updateOrCreateRows(ctx, realm, testResults); err != nil {
 			return err
 		}
 	}
 	return nil
 }

 func (u *testMetadataUpdator) updateOrCreateRows(ctx context.Context, realm string, testMetadata map[testID]*pb.TestMetadata) error {
 	project, subRealm := realms.Split(realm)
 	testIDs := make([]string, 0, len(testMetadata))
 	for k := range testMetadata {
 		testIDs = append(testIDs, string(k))
 	}
 	f := func(ctx context.Context) error {
 		seen := make(map[string]bool)
 		ms := []*spanner.Mutation{}
 		opts := testmetadata.ReadTestMetadataOptions{
 			Project:   project,
 			TestIDs:   testIDs,
 			SourceRef: pbutil.SourceRefFromSources(u.sources),
 			SubRealm:  subRealm,
 		}
 		// Update existing metadata entries.
 		err := testmetadata.ReadTestMetadata(ctx, opts, func(tmd *testmetadata.TestMetadataRow) error {
 			seen[tmd.TestID] = true
 			// No update if test result from a lower commit position than existing test metadata.
 			if u.sources.GitilesCommit.Position < tmd.Position {
 				return nil
 			}
 			newMetadata := testMetadata[testID(tmd.TestID)]
 			newBitField := fieldExistenceBitField(newMetadata)
 			existingBitField := fieldExistenceBitField(tmd.TestMetadata)
 			// If test result from the same commit position,
 			// then only update when more metadata fields were found in the test results.
 			// If test result from a higher commit position,
 			// then update when more or equal number of metadata fields were found in the test results
 			// OR existing test metadata expired.
 			if (u.sources.GitilesCommit.Position == tmd.Position && newBitField > existingBitField) ||
 				(u.sources.GitilesCommit.Position > tmd.Position &&
 					(tmd.LastUpdated.Add(metadataExpiry).Before(u.start) || newBitField >= existingBitField)) {
 				mutation := u.testMetadataMutation(project, tmd.TestID, subRealm, testMetadata[testID(tmd.TestID)])
 				ms = append(ms, mutation)
 			}
 			return nil
 		})
 		if err != nil {
 			return err
 		}
 		// Create new metadata entries which aren't exist yet.
 		for testID, tm := range testMetadata {
 			if !seen[string(testID)] {
 				mutation := u.testMetadataMutation(project, string(testID), subRealm, tm)
 				ms = append(ms, mutation)
 			}
 		}
 		span.BufferWrite(ctx, ms...)
 		return nil
 	}
 	_, err := span.ReadWriteTransaction(ctx, f)
 	if err != nil {
 		return err
 	}
 	return nil
 }

 var saveCols = []string{
 	"Project",
 	"TestId",
 	"RefHash",
 	"SubRealm",
 	"LastUpdated",
 	"TestMetadata",
 	"SourceRef",
 	"Position",
 }

 func (u *testMetadataUpdator) testMetadataMutation(project, testID, subRealm string, tm *pb.TestMetadata) *spanner.Mutation {
 	ref := pbutil.SourceRefFromSources(u.sources)
 	vals := []any{
 		project,
 		testID,
 		pbutil.SourceRefHash(ref),
 		subRealm,
 		spanner.CommitTimestamp,
 		spanutil.Compressed(pbutil.MustMarshal(tm)).ToSpanner(),
 		spanutil.Compressed(pbutil.MustMarshal(ref)).ToSpanner(),
 		u.sources.GitilesCommit.Position,
 	}
 	return spanner.InsertOrUpdate("TestMetadata", saveCols, vals)
 }

 // fieldExistenceBitField returns bit fields contains 0 or 1 in each bit to indicate
 // the existence of each field in the test metadata.
 // Fields from the lowest bit to highest bit:
 // TestMetadata.name
 // TestMetadata.location.repo
 // TestMetadata.location.file_name
 // TestMetadata.location.line
 // TestMetadata.bug_component
 // TestMetadata.properties
 func fieldExistenceBitField(metadata *pb.TestMetadata) uint8 {
 	bitField := uint8(0)
 	bitFieldOrder := []string{
 		"name",
 		"location.repo",
 		"location.file_name",
 		"location.line",
 		"bug_component",
 		"properties",
 	}
 	for i, k := range bitFieldOrder {
 		if exist(strings.Split(k, "."), metadata.ProtoReflect()) {
 			bitField += uint8(math.Pow(2, float64(i)))
 		}
 	}
 	return bitField
 }

 func exist(fieldNameTokens []string, message protoreflect.Message) bool {
 	if len(fieldNameTokens) == 0 {
 		return true
 	}
 	curName := fieldNameTokens[0]
 	fd := message.Descriptor().Fields().ByTextName(curName)
 	if message.Has(fd) {
 		if fd.Kind() == protoreflect.MessageKind {
 			return exist(fieldNameTokens[1:], message.Get(fd).Message())
 		}
 		return len(fieldNameTokens) == 1
 	}
 	return false
 }
	// Copyright 2023 The LUCI Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	// Package testmetadataupdator implements a task which ingest test metadata from invocations.
	package testmetadataupdator

	import (
	"context"
	"math"
	"strings"
	"time"

	"cloud.google.com/go/spanner"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/proto/mask"
	"go.chromium.org/luci/resultdb/internal/invocations"
	"go.chromium.org/luci/resultdb/internal/spanutil"
	"go.chromium.org/luci/resultdb/internal/testmetadata"
	"go.chromium.org/luci/resultdb/internal/testresults"
	"go.chromium.org/luci/resultdb/pbutil"
	pb "go.chromium.org/luci/resultdb/proto/v1"
	"go.chromium.org/luci/server/auth/realms"
	"go.chromium.org/luci/server/span"
	"golang.org/x/sync/errgroup"
	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
	)

	// Expiration time of a test metadata.
	// After expiry, update to less metadata information is allowed.
	var metadataExpiry = 48 * time.Hour

	// Maximum number of incoming test metadata we hold in memory.
	var batchSize = 1000

	type testID string

	// testMetadataUpdator updates TestMetadata rows with test results in a set of invocations.
	type testMetadataUpdator struct {
	sources *pb.Sources
	invIDs invocations.IDSet
	start time.Time // Used to determine if a test metadata is expired.
	}

	func newUpdator(sources pb.Sources, invIDs invocations.IDSet, start time.Time) testMetadataUpdator {
	return &testMetadataUpdator{sources: sources, invIDs: invIDs, start: start}
	}

	func (u *testMetadataUpdator) run(ctx context.Context) error {
	for invID := range u.invIDs {
	// Query test results and use them to update test metadata.
	batchC := make(chan map[testID]*pb.TestMetadata)
	// Batch update test metadata.
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
	defer close(batchC)
	if err := u.queryTestResultMetadata(ctx, invID, batchC); err != nil {
	return errors.Annotate(err, "query test result metadata, invocation %s", invID).Err()
	}
	return nil
	})

	eg.Go(func() error {
	var realm string
	if err := invocations.ReadColumns(span.Single(ctx), invID, map[string]any{"Realm": &realm}); err != nil {
	return errors.Annotate(err, "read realm, invocation %s", invID).Err()
	}
	if err := u.batchCreateOrUpdateTestMetadata(ctx, realm, batchC); err != nil {
	return errors.Annotate(err, "create or update test metadata, invocation %s", invID).Err()
	}
	return nil
	})
	if err := eg.Wait(); err != nil {
	return err
	}
	}
	return nil
	}

	// queryTestResultMetadata visits all test results in the given invocations.
	// and returns one test metadata with the most metadata information for each test.
	func (u testMetadataUpdator) queryTestResultMetadata(ctx context.Context, invID invocations.ID, batchC chan<- map[testID]pb.TestMetadata) error {
	ctx, cancel := span.ReadOnlyTransaction(ctx)
	defer cancel()

	q := testresults.Query{
	InvocationIDs: invocations.NewIDSet(invID),
	Mask: mask.MustFromReadMask(&pb.TestResult{},
	"test_id",
	"test_metadata",
	),
	}
	deduplicatedMetadata := make(map[testID]*pb.TestMetadata, batchSize)
	err := q.Run(ctx, func(tr *pb.TestResult) error {
	existing, ok := deduplicatedMetadata[testID(tr.TestId)]
	if !ok {
	if len(deduplicatedMetadata) >= batchSize {
	select {
	case <-ctx.Done():
	return ctx.Err()
	case batchC <- deduplicatedMetadata:
	}
	deduplicatedMetadata = make(map[testID]*pb.TestMetadata, batchSize)
	}
	deduplicatedMetadata[testID(tr.TestId)] = tr.TestMetadata
	return nil
	}
	if fieldExistenceBitField(existing) < fieldExistenceBitField(tr.TestMetadata) {
	deduplicatedMetadata[testID(tr.TestId)] = tr.TestMetadata
	}
	return nil
	})
	if err != nil {
	return err
	}
	if len(deduplicatedMetadata) > 0 {
	select {
	case <-ctx.Done():
	return ctx.Err()
	case batchC <- deduplicatedMetadata:
	}
	}
	return nil
	}

	func (u testMetadataUpdator) batchCreateOrUpdateTestMetadata(ctx context.Context, realm string, batchC <-chan map[testID]pb.TestMetadata) error {
	for testResults := range batchC {
	if err := u.updateOrCreateRows(ctx, realm, testResults); err != nil {
	return err
	}
	}
	return nil
	}

	func (u testMetadataUpdator) updateOrCreateRows(ctx context.Context, realm string, testMetadata map[testID]pb.TestMetadata) error {
	project, subRealm := realms.Split(realm)
	testIDs := make([]string, 0, len(testMetadata))
	for k := range testMetadata {
	testIDs = append(testIDs, string(k))
	}
	f := func(ctx context.Context) error {
	seen := make(map[string]bool)
	ms := []*spanner.Mutation{}
	opts := testmetadata.ReadTestMetadataOptions{
	Project: project,
	TestIDs: testIDs,
	SourceRef: pbutil.SourceRefFromSources(u.sources),
	SubRealm: subRealm,
	}
	// Update existing metadata entries.
	err := testmetadata.ReadTestMetadata(ctx, opts, func(tmd *testmetadata.TestMetadataRow) error {
	seen[tmd.TestID] = true
	// No update if test result from a lower commit position than existing test metadata.
	if u.sources.GitilesCommit.Position < tmd.Position {
	return nil
	}
	newMetadata := testMetadata[testID(tmd.TestID)]
	newBitField := fieldExistenceBitField(newMetadata)
	existingBitField := fieldExistenceBitField(tmd.TestMetadata)
	// If test result from the same commit position,
	// then only update when more metadata fields were found in the test results.
	// If test result from a higher commit position,
	// then update when more or equal number of metadata fields were found in the test results
	// OR existing test metadata expired.
	if (u.sources.GitilesCommit.Position == tmd.Position && newBitField > existingBitField) \|\|
	(u.sources.GitilesCommit.Position > tmd.Position &&
	(tmd.LastUpdated.Add(metadataExpiry).Before(u.start) \|\| newBitField >= existingBitField)) {
	mutation := u.testMetadataMutation(project, tmd.TestID, subRealm, testMetadata[testID(tmd.TestID)])
	ms = append(ms, mutation)
	}
	return nil
	})
	if err != nil {
	return err
	}
	// Create new metadata entries which aren't exist yet.
	for testID, tm := range testMetadata {
	if !seen[string(testID)] {
	mutation := u.testMetadataMutation(project, string(testID), subRealm, tm)
	ms = append(ms, mutation)
	}
	}
	span.BufferWrite(ctx, ms...)
	return nil
	}
	_, err := span.ReadWriteTransaction(ctx, f)
	if err != nil {
	return err
	}
	return nil
	}

	var saveCols = []string{
	"Project",
	"TestId",
	"RefHash",
	"SubRealm",
	"LastUpdated",
	"TestMetadata",
	"SourceRef",
	"Position",
	}

	func (u testMetadataUpdator) testMetadataMutation(project, testID, subRealm string, tm pb.TestMetadata) *spanner.Mutation {
	ref := pbutil.SourceRefFromSources(u.sources)
	vals := []any{
	project,
	testID,
	pbutil.SourceRefHash(ref),
	subRealm,
	spanner.CommitTimestamp,
	spanutil.Compressed(pbutil.MustMarshal(tm)).ToSpanner(),
	spanutil.Compressed(pbutil.MustMarshal(ref)).ToSpanner(),
	u.sources.GitilesCommit.Position,
	}
	return spanner.InsertOrUpdate("TestMetadata", saveCols, vals)
	}

	// fieldExistenceBitField returns bit fields contains 0 or 1 in each bit to indicate
	// the existence of each field in the test metadata.
	// Fields from the lowest bit to highest bit:
	// TestMetadata.name
	// TestMetadata.location.repo
	// TestMetadata.location.file_name
	// TestMetadata.location.line
	// TestMetadata.bug_component
	// TestMetadata.properties
	func fieldExistenceBitField(metadata *pb.TestMetadata) uint8 {
	bitField := uint8(0)
	bitFieldOrder := []string{
	"name",
	"location.repo",
	"location.file_name",
	"location.line",
	"bug_component",
	"properties",
	}
	for i, k := range bitFieldOrder {
	if exist(strings.Split(k, "."), metadata.ProtoReflect()) {
	bitField += uint8(math.Pow(2, float64(i)))
	}
	}
	return bitField
	}

	func exist(fieldNameTokens []string, message protoreflect.Message) bool {
	if len(fieldNameTokens) == 0 {
	return true
	}
	curName := fieldNameTokens[0]
	fd := message.Descriptor().Fields().ByTextName(curName)
	if message.Has(fd) {
	if fd.Kind() == protoreflect.MessageKind {
	return exist(fieldNameTokens[1:], message.Get(fd).Message())
	}
	return len(fieldNameTokens) == 1
	}
	return false
	}