syz-cluster: enable atomic updates for SessionTestRepo

A raw InsertOrUpdate method is not very reliable in case of concurrent
update requests. Add a callback inside which the modified fields would
be set.

Refactor the existing code that used to call the old method.

4 files changed, 12 insertions, 4 deletions

diff --git a/syz-cluster/pkg/db/report_repo_test.go b/syz-cluster/pkg/db/report_repo_test.go
index 9648a7bc2..ee05a2e18 100644
--- a/syz-cluster/pkg/db/report_repo_test.go
+++ b/syz-cluster/pkg/db/report_repo_test.go
@@ -74,7 +74,7 @@ func TestSessionsWithoutReports(t *testing.T) {
 				SessionID: session.ID,
 				TestName:  "test",
 				Result:    api.TestPassed,
-			})
+			}, nil)
 			assert.NoError(t, err)
 			err = findingRepo.Save(ctx, &Finding{
 				SessionID: session.ID,
diff --git a/syz-cluster/pkg/db/session_test_repo.go b/syz-cluster/pkg/db/session_test_repo.go
index 43ada137c..e5b923fd4 100644
--- a/syz-cluster/pkg/db/session_test_repo.go
+++ b/syz-cluster/pkg/db/session_test_repo.go
@@ -20,7 +20,9 @@ func NewSessionTestRepository(client *spanner.Client) *SessionTestRepository {
 	}
 }
 
-func (repo *SessionTestRepository) InsertOrUpdate(ctx context.Context, test *SessionTest) error {
+// If the beforeSave callback is specified, it will be called before saving the entity.
+func (repo *SessionTestRepository) InsertOrUpdate(ctx context.Context, test *SessionTest,
+	beforeSave func(*SessionTest)) error {
 	_, err := repo.client.ReadWriteTransaction(ctx,
 		func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
 			// Check if the test already exists.
@@ -38,6 +40,9 @@ func (repo *SessionTestRepository) InsertOrUpdate(ctx context.Context, test *Ses
 
 			_, iterErr := iter.Next()
 			if iterErr == nil {
+				if beforeSave != nil {
+					beforeSave(test)
+				}
 				m, err := spanner.UpdateStruct("SessionTests", test)
 				if err != nil {
 					return err
@@ -46,6 +51,9 @@ func (repo *SessionTestRepository) InsertOrUpdate(ctx context.Context, test *Ses
 			} else if iterErr != iterator.Done {
 				return iterErr
 			} else {
+				if beforeSave != nil {
+					beforeSave(test)
+				}
 				m, err := spanner.InsertStruct("SessionTests", test)
 				if err != nil {
 					return err
diff --git a/syz-cluster/pkg/db/session_test_repo_test.go b/syz-cluster/pkg/db/session_test_repo_test.go
index f18349339..49eb2f3b4 100644
--- a/syz-cluster/pkg/db/session_test_repo_test.go
+++ b/syz-cluster/pkg/db/session_test_repo_test.go
@@ -43,7 +43,7 @@ func TestSessionTestRepository(t *testing.T) {
 			PatchedBuildID: spanner.NullString{StringVal: build2.ID, Valid: true},
 			Result:         api.TestPassed,
 		}
-		err = testsRepo.InsertOrUpdate(ctx, test)
+		err = testsRepo.InsertOrUpdate(ctx, test, nil)
 		assert.NoError(t, err)
 	}
 
diff --git a/syz-cluster/pkg/db/util_test.go b/syz-cluster/pkg/db/util_test.go
index f33f1e221..64fedeab2 100644
--- a/syz-cluster/pkg/db/util_test.go
+++ b/syz-cluster/pkg/db/util_test.go
@@ -26,7 +26,7 @@ func (d *dummyTestData) addSessionTest(session *Session, names ...string) {
 			SessionID: session.ID,
 			TestName:  name,
 			Result:    api.TestPassed,
-		})
+		}, nil)
 		assert.NoError(d.t, err)
 	}
 }