fix(worker): lock the risk while resolving its evidence links

internal/service/relational/risks/scores.go

@@ -149,7 +149,7 @@ func (s *RiskService) RecordRiskScoreSnapshot(tx *gorm.DB, riskID uuid.UUID, sou
 		status = RiskScoreStatusDeleted
 		openBaselineScore = 0
 		openResidualScore = 0
-	} else if isTerminalRiskStatus(risk.Status) {
+	} else if IsTerminalRiskStatus(risk.Status) {
 		openBaselineScore = 0
 		openResidualScore = 0
 	}
@@ -180,7 +180,8 @@ func isScoreValueEvent(eventType RiskEventType) bool {
 	}
 }
 
-func isTerminalRiskStatus(status string) bool {
+// IsTerminalRiskStatus reports whether status is closed or remediated.
+func IsTerminalRiskStatus(status string) bool {
 	return status == string(RiskStatusClosed) || status == string(RiskStatusRemediated)
 }
 

internal/service/worker/risk_evidence_integration_helpers_test.go

@@ -0,0 +1,31 @@
+//go:build integration
+
+package worker
+
+import (
+	"testing"
+	"time"
+
+	"github.com/compliance-framework/api/internal/service/relational"
+	oscalTypes_1_1_3 "github.com/defenseunicorns/go-oscal/src/types/oscal-1-1-3"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"gorm.io/datatypes"
+	"gorm.io/gorm"
+)
+
+// createStreamEvidence creates an evidence of the given stream with the given state and end.
+func createStreamEvidence(t *testing.T, db *gorm.DB, stream uuid.UUID, state string, end time.Time) *relational.Evidence {
+	t.Helper()
+	id := uuid.New()
+	evidence := &relational.Evidence{
+		UUIDModel: relational.UUIDModel{ID: &id},
+		UUID:      stream,
+		Title:     "Test Evidence",
+		Start:     end.Add(-time.Hour),
+		End:       end,
+		Status:    datatypes.NewJSONType(oscalTypes_1_1_3.ObjectiveStatus{State: state}),
+	}
+	require.NoError(t, db.Create(evidence).Error)
+	return evidence
+}

internal/service/worker/risk_evidence_worker.go

@@ -1089,7 +1089,7 @@ func (w *RiskEvidenceWorker) handleEvidenceResolution(ctx context.Context, evide
 			continue
 		}
 		// Skip risks that are already closed or remediated — nothing to resolve.
-		if risk.Status == string(risks.RiskStatusClosed) || risk.Status == string(risks.RiskStatusRemediated) {
+		if risks.IsTerminalRiskStatus(risk.Status) {
 			continue
 		}
 
@@ -1098,7 +1098,7 @@ func (w *RiskEvidenceWorker) handleEvidenceResolution(ctx context.Context, evide
 			continue
 		}
 
-		if err := w.resolveRiskEvidenceLink(ctx, risk, evidence.UUID); err != nil {
+		if err := w.resolveRiskEvidenceLink(ctx, *risk.ID, evidence.UUID); err != nil {
 			w.logger.Errorw("Failed to resolve risk evidence link",
 				"error", err, "risk_id", risk.ID, "evidence_id", evidence.UUID)
 			errs = append(errs, err)
@@ -1138,8 +1138,23 @@ func (w *RiskEvidenceWorker) shouldRemoveEvidenceLink(risk *risks.Risk, evidence
 
 // resolveRiskEvidenceLink removes a single evidence link from a risk and, if no links remain,
 // transitions the risk to "remediated" (or emits an audit event for "risk-accepted" risks).
-func (w *RiskEvidenceWorker) resolveRiskEvidenceLink(ctx context.Context, risk *risks.Risk, evidenceStreamID uuid.UUID) error {
+// It locks the risk first, so resolutions of the risk's other streams count its links only
+// after this one commits, and it acts on the status read under the lock.
+func (w *RiskEvidenceWorker) resolveRiskEvidenceLink(ctx context.Context, riskID uuid.UUID, evidenceStreamID uuid.UUID) error {
 	return w.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var risk risks.Risk
+		err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "status").First(&risk, "id = ?", riskID).Error
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil
+		}
+		if err != nil {
+			return fmt.Errorf("failed to lock risk: %w", err)
+		}
+		// Closed or remediated since the caller read it: nothing to resolve.
+		if risks.IsTerminalRiskStatus(risk.Status) {
+			return nil
+		}
+
 		// Delete the evidence link.
 		result := tx.Delete(&risks.RiskEvidenceLink{}, "risk_id = ? AND evidence_id = ?", *risk.ID, evidenceStreamID)
 		if result.Error != nil {
@@ -1187,7 +1202,7 @@ func (w *RiskEvidenceWorker) resolveRiskEvidenceLink(ctx context.Context, risk *
 		// Transition to remediated.
 		oldStatus := risk.Status
 		occurredAt := time.Now().UTC()
-		if err := tx.Model(risk).Update("status", string(risks.RiskStatusRemediated)).Error; err != nil {
+		if err := tx.Model(&risk).Update("status", string(risks.RiskStatusRemediated)).Error; err != nil {
 			return fmt.Errorf("failed to transition risk to remediated: %w", err)
 		}
 		if err := w.emitRiskEventAt(ctx, tx, *risk.ID, string(risks.RiskEventTypeStatusChange), map[string]interface{}{

internal/service/worker/risk_resolution_race_integration_test.go

@@ -0,0 +1,124 @@
+//go:build integration
+
+package worker
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/compliance-framework/api/internal/service/relational"
+	"github.com/compliance-framework/api/internal/service/relational/risks"
+	"github.com/compliance-framework/api/internal/tests"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/suite"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+)
+
+type RiskResolutionRaceIntegrationSuite struct {
+	tests.IntegrationTestSuite
+}
+
+func TestRiskResolutionRaceIntegrationSuite(t *testing.T) {
+	suite.Run(t, new(RiskResolutionRaceIntegrationSuite))
+}
+
+func (s *RiskResolutionRaceIntegrationSuite) SetupTest() {
+	s.Require().NoError(s.Migrator.Refresh())
+}
+
+// Two streams of a risk pass at once: the second job runs between the first's count of the
+// risk's remaining links and its commit, and the risk must still end up remediated.
+func (s *RiskResolutionRaceIntegrationSuite) TestConcurrentPassingStreamsRemediateTheRisk() {
+	ctx := context.Background()
+	worker := NewRiskEvidenceWorker(s.DB, zap.NewNop().Sugar())
+	first := createStreamEvidence(s.T(), s.DB, uuid.New(), relational.EvidenceStatusSatisfied, time.Now())
+	second := createStreamEvidence(s.T(), s.DB, uuid.New(), relational.EvidenceStatusSatisfied, time.Now())
+	risk := createRiskWithEvidenceLink(s.T(), s.DB, string(risks.RiskStatusOpen), first.UUID, nil)
+	s.Require().NoError(s.DB.Create(&risks.RiskEvidenceLink{RiskID: *risk.ID, EvidenceID: second.UUID}).Error)
+
+	const cb = "test:second_stream_resolves"
+	var fired atomic.Bool
+	var secondErr error
+	secondDone := make(chan struct{})
+	s.Require().NoError(s.DB.Callback().Query().After("gorm:query").Register(cb, func(db *gorm.DB) {
+		_, isCount := db.Statement.Dest.(*int64)
+		if !isCount || db.Statement.Table != (risks.RiskEvidenceLink{}).TableName() || !fired.CompareAndSwap(false, true) {
+			return
+		}
+		go func() {
+			defer close(secondDone)
+			secondErr = worker.Work(ctx, makeWorkerJob(RiskProcessEvidenceArgs{EvidenceID: *second.ID}))
+		}()
+		// Let the second job finish, or block on a lock the first job holds.
+		s.Eventually(func() bool {
+			select {
+			case <-secondDone:
+				return true
+			default:
+			}
+			var waiting int64
+			err := s.DB.Raw(`SELECT count(*) FROM pg_stat_activity
+				WHERE datname = current_database() AND wait_event_type = 'Lock'`).Scan(&waiting).Error
+			return err == nil && waiting > 0
+		}, 10*time.Second, 10*time.Millisecond)
+	}))
+	defer func() { s.Require().NoError(s.DB.Callback().Query().Remove(cb)) }()
+
+	s.Require().NoError(worker.Work(ctx, makeWorkerJob(RiskProcessEvidenceArgs{EvidenceID: *first.ID})))
+	s.Require().True(fired.Load())
+	<-secondDone
+	s.Require().NoError(secondErr)
+
+	s.assertRisk(*risk.ID, risks.RiskStatusRemediated, 0)
+}
+
+// A user closes the risk after resolution read it: the link stays and the risk stays closed.
+func (s *RiskResolutionRaceIntegrationSuite) TestResolutionLeavesRiskClosedMeanwhile() {
+	riskID := s.resolveWhileStatusChanges(risks.RiskStatusClosed)
+	s.assertRisk(riskID, risks.RiskStatusClosed, 1)
+}
+
+// A user accepts the risk after resolution read it: the link goes and the risk stays
+// accepted, with an evidence_recovered event in place of remediation.
+func (s *RiskResolutionRaceIntegrationSuite) TestResolutionKeepsRiskAcceptedMeanwhile() {
+	riskID := s.resolveWhileStatusChanges(risks.RiskStatusRiskAccepted)
+	s.assertRisk(riskID, risks.RiskStatusRiskAccepted, 0)
+	var recovered int64
+	s.Require().NoError(s.DB.Model(&risks.RiskEvent{}).Where("risk_id = ? AND event_type = ?",
+		riskID, string(risks.RiskEventTypeEvidenceRecovered)).Count(&recovered).Error)
+	s.Equal(int64(1), recovered)
+}
+
+// resolveWhileStatusChanges resolves a passing evidence of an open risk, setting the risk's
+// status right after resolution's unlocked read of the linked risks.
+func (s *RiskResolutionRaceIntegrationSuite) resolveWhileStatusChanges(status risks.RiskStatus) uuid.UUID {
+	worker := NewRiskEvidenceWorker(s.DB, zap.NewNop().Sugar())
+	evidence := createStreamEvidence(s.T(), s.DB, uuid.New(), relational.EvidenceStatusSatisfied, time.Now())
+	risk := createRiskWithEvidenceLink(s.T(), s.DB, string(risks.RiskStatusOpen), evidence.UUID, nil)
+
+	const cb = "test:status_changes_after_read"
+	var changed atomic.Bool
+	s.Require().NoError(s.DB.Callback().Query().After("gorm:query").Register(cb, func(db *gorm.DB) {
+		if _, ok := db.Statement.Dest.(*[]risks.Risk); ok && changed.CompareAndSwap(false, true) {
+			s.Require().NoError(s.DB.Model(risk).Update("status", string(status)).Error)
+		}
+	}))
+	defer func() { s.Require().NoError(s.DB.Callback().Query().Remove(cb)) }()
+
+	s.Require().NoError(worker.handleEvidenceResolution(context.Background(), evidence))
+	s.Require().True(changed.Load())
+	return *risk.ID
+}
+
+// assertRisk checks the stored risk's status and how many evidence links it has left.
+func (s *RiskResolutionRaceIntegrationSuite) assertRisk(id uuid.UUID, status risks.RiskStatus, links int64) {
+	var stored risks.Risk
+	s.Require().NoError(s.DB.First(&stored, "id = ?", id).Error)
+	s.Equal(string(status), stored.Status)
+	var count int64
+	s.Require().NoError(s.DB.Model(&risks.RiskEvidenceLink{}).Where("risk_id = ?", id).Count(&count).Error)
+	s.Equal(links, count)
+}