diff --git a/backend/modules/evaluation/domain/service/expt_manage_execution_impl.go b/backend/modules/evaluation/domain/service/expt_manage_execution_impl.go index da8dffc20..0714f8b0e 100644 --- a/backend/modules/evaluation/domain/service/expt_manage_execution_impl.go +++ b/backend/modules/evaluation/domain/service/expt_manage_execution_impl.go @@ -303,12 +303,48 @@ func withRetryYieldExt(ext map[string]string, enabled bool) map[string]string { return ext } -func (e *ExptMangerImpl) Run(ctx context.Context, exptID, runID, spaceID int64, itemRetryNum int, session *entity.Session, runMode entity.ExptRunMode, ext map[string]string) error { - if err := NewQuotaService(e.quotaRepo, e.configer).AllowExptRun(ctx, exptID, spaceID, session); err != nil { - return err +func (e *ExptMangerImpl) prepareRun(ctx context.Context, exptID, runID, spaceID int64, session *entity.Session) (expt *entity.Experiment, err error) { + defer func() { + if err != nil { + e.cleanupUnscheduledRun(ctx, exptID, runID, err) + } + }() + + // Read dependencies before reserving quota so read failures cannot leak a quota slot. + expt, err = e.GetDetail(ctx, exptID, spaceID, session) + if err != nil { + return nil, err } + err = NewQuotaService(e.quotaRepo, e.configer).AllowExptRun(ctx, exptID, spaceID, session) + return expt, err +} - expt, err := e.GetDetail(ctx, exptID, spaceID, session) +func (e *ExptMangerImpl) cleanupUnscheduledRun(ctx context.Context, exptID, runID int64, cause error) { + if runID <= 0 { + logs.CtxError(ctx, "[ExptEval][RunLock] invalid cleanup run [expt_id=%v run_id=%v]", exptID, runID) + return + } + + // Only the owned run before MQ publication may be unlocked here. + // Attempt this bare DEL synchronously once; a delayed retry could delete a new run's lock. + ctx = context.WithoutCancel(ctx) + stateCtx, cancelState := context.WithTimeout(ctx, exptRunLogPersistTimeout) + stateErr := e.runLogRepo.Update(stateCtx, exptID, runID, map[string]any{ + "status": int64(entity.ExptStatus_Failed), "status_message": []byte(cause.Error()), + }) + cancelState() + + unlockCtx, cancelUnlock := context.WithTimeout(ctx, 2*time.Second) + released, unlockErr := e.mutex.UnlockForce(unlockCtx, e.makeExptMutexLockKey(exptID)) + cancelUnlock() + if stateErr != nil || unlockErr != nil { + logs.CtxError(ctx, "[ExptEval][RunLock] cleanup failed [expt_id=%v run_id=%v state_err=%v unlock_err=%v]", exptID, runID, stateErr, unlockErr) + } + logs.CtxInfo(ctx, "[ExptEval][RunLock] cleanup result [expt_id=%v run_id=%v released=%v]", exptID, runID, released) +} + +func (e *ExptMangerImpl) Run(ctx context.Context, exptID, runID, spaceID int64, itemRetryNum int, session *entity.Session, runMode entity.ExptRunMode, ext map[string]string) error { + expt, err := e.prepareRun(ctx, exptID, runID, spaceID, session) if err != nil { return err } @@ -414,11 +450,7 @@ func buildExptNotifyParam(expt *entity.Experiment, toStatus entity.ExptStatus) ( } func (e *ExptMangerImpl) RetryItems(ctx context.Context, exptID, runID, spaceID int64, itemRetryNum int, itemIDs []int64, session *entity.Session, ext map[string]string) error { - if err := NewQuotaService(e.quotaRepo, e.configer).AllowExptRun(ctx, exptID, spaceID, session); err != nil { - return err - } - - expt, err := e.GetDetail(ctx, exptID, spaceID, session) + expt, err := e.prepareRun(ctx, exptID, runID, spaceID, session) if err != nil { return err } @@ -1679,15 +1711,22 @@ func (e *ExptMangerImpl) unlockCompletingRun(ctx context.Context, exptID, exptRu return err } -func (e *ExptMangerImpl) LogRun(ctx context.Context, exptID, exptRunID int64, mode entity.ExptRunMode, spaceID int64, itemIDs []int64, session *entity.Session) error { +func (e *ExptMangerImpl) LogRun(ctx context.Context, exptID, exptRunID int64, mode entity.ExptRunMode, spaceID int64, itemIDs []int64, session *entity.Session) (err error) { duration := time.Duration(e.configer.GetExptExecConf(ctx, spaceID).GetZombieIntervalSecond()) * time.Second locked, err := e.mutex.LockBackoff(ctx, e.makeExptMutexLockKey(exptID), duration, time.Second) if err != nil { return err } if !locked { + logs.CtxInfo(ctx, "[ExptEval][RunLock] lock occupied [expt_id=%v run_id=%v]", exptID, exptRunID) return errorx.NewByCode(errno.ExperimentRunningExistedCode) } + ownedRunID := exptRunID + defer func(runID int64) { + if err != nil { + e.cleanupUnscheduledRun(ctx, exptID, runID, err) + } + }(ownedRunID) defer e.mtr.EmitExptExecRun(spaceID, int64(mode)) @@ -1733,6 +1772,14 @@ func (e *ExptMangerImpl) LogRetryItemsRun(ctx context.Context, exptID int64, mod if err != nil { return 0, false, err } + if locked { + ownedRunID := runID + defer func(runID int64) { + if err != nil { + e.cleanupUnscheduledRun(ctx, exptID, runID, err) + } + }(ownedRunID) + } var rl *entity.ExptRunLog retried = !locked @@ -1740,9 +1787,10 @@ func (e *ExptMangerImpl) LogRetryItemsRun(ctx context.Context, exptID int64, mod if retried { runID, err = strconv.ParseInt(existedRunID, 10, 64) if err != nil { - logs.CtxError(ctx, "parsing expt run lock value to runid failed, raw: %v", existedRunID) + logs.CtxInfo(ctx, "[ExptEval][RunLock] lock occupied [expt_id=%v holder=%s]", exptID, existedRunID) return 0, false, errorx.NewByCode(errno.ExperimentRunningExistedCode) } + logs.CtxDebug(ctx, "[ExptEval][RunLock] joining existing run [expt_id=%v run_id=%v]", exptID, runID) completing, err := e.ExistCompletingRunLock(ctx, exptID, runID, spaceID) if err != nil { diff --git a/backend/modules/evaluation/domain/service/expt_manage_execution_impl_test.go b/backend/modules/evaluation/domain/service/expt_manage_execution_impl_test.go index 29186029a..d04d4e96d 100755 --- a/backend/modules/evaluation/domain/service/expt_manage_execution_impl_test.go +++ b/backend/modules/evaluation/domain/service/expt_manage_execution_impl_test.go @@ -153,6 +153,8 @@ func TestExptMangerImpl_Run(t *testing.T) { EXPECT(). CreateOrUpdate(ctx, int64(789), gomock.Any(), session). Return(errors.New("quota exceeded")) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), int64(123), int64(456), gomock.Any()).Return(nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(123)).Return(true, nil) }, wantErr: true, }, @@ -622,6 +624,8 @@ func TestExptMangerImpl_LogRun(t *testing.T) { EXPECT(). Create(ctx, gomock.Any()). Return(errors.New("create failed")) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), int64(123), int64(456), gomock.Any()).Return(nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(123)).Return(true, nil) }, wantErr: true, }, @@ -802,6 +806,8 @@ func TestExptMangerImpl_LogRetryItemsRun(t *testing.T) { Return(true, "1006", nil) mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo). EXPECT().Save(ctx, gomock.Any()).Return(errors.New("save failed")) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), exptID, int64(1006), gomock.Any()).Return(nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(exptID)).Return(true, nil) }, wantErr: true, }, @@ -898,6 +904,8 @@ func TestExptMangerImpl_RetryItems(t *testing.T) { setup: func() { mgr.quotaRepo.(*repoMocks.MockQuotaRepo). EXPECT().CreateOrUpdate(ctx, spaceID, gomock.Any(), session).Return(errors.New("quota exceeded")) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), exptID, runID, gomock.Any()).Return(nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(exptID)).Return(true, nil) mgr.configer.(*componentMocks.MockIConfiger). EXPECT().GetExptExecConf(ctx, spaceID).AnyTimes(). Return(&entity.ExptExecConf{SpaceExptConcurLimit: 10}) diff --git a/backend/modules/evaluation/domain/service/expt_run_admission_cleanup_test.go b/backend/modules/evaluation/domain/service/expt_run_admission_cleanup_test.go new file mode 100644 index 000000000..a87dbb07a --- /dev/null +++ b/backend/modules/evaluation/domain/service/expt_run_admission_cleanup_test.go @@ -0,0 +1,232 @@ +// Copyright (c) 2025 coze-dev Authors +// SPDX-License-Identifier: Apache-2.0 + +package service + +import ( + "context" + "errors" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + idgenMocks "github.com/coze-dev/coze-loop/backend/infra/idgen/mocks" + lockMocks "github.com/coze-dev/coze-loop/backend/infra/lock/mocks" + lwtMocks "github.com/coze-dev/coze-loop/backend/infra/platestwrite/mocks" + metricsMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/component/metrics/mocks" + componentMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/component/mocks" + "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/entity" + eventsMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/events/mocks" + repoMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/repo/mocks" + svcMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/service/mocks" + "github.com/coze-dev/coze-loop/backend/modules/evaluation/pkg/errno" + "github.com/coze-dev/coze-loop/backend/pkg/errorx" +) + +func TestRunAdmissionCleanup(t *testing.T) { + for _, tc := range []struct { + name string + mode entity.ExptRunMode + publishFail bool + }{ + {"submit", entity.EvaluationModeSubmit, false}, + {"retry_all", entity.EvaluationModeRetryAll, false}, + {"retry_failure", entity.EvaluationModeFailRetry, false}, + {"retry_items", entity.EvaluationModeRetryItems, false}, + {"publish_failure_preserves_run", entity.EvaluationModeFailRetry, true}, + {"publish_failure_preserves_retry_items_run", entity.EvaluationModeRetryItems, true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + ctx := context.Background() + session := &entity.Session{UserID: "u"} + const exptID, runID, spaceID = int64(1), int64(2), int64(3) + limit, logCalls := 2, 2 + publishErr := errors.New("publish timed out") + if tc.publishFail { + limit, logCalls = 3, 1 + mgr.publisher.(*eventsMocks.MockExptEventPublisher).EXPECT().PublishExptScheduleEvent(ctx, gomock.Any(), gomock.Any()).Return(publishErr) + } + configer := componentMocks.NewMockIConfiger(ctrl) + configer.EXPECT().GetExptExecConf(ctx, spaceID).Return(&entity.ExptExecConf{ + SpaceExptConcurLimit: limit, ZombieIntervalSecond: 72 * 60 * 60, + }).AnyTimes() + configer.EXPECT().GetRetryYieldEnabled(ctx, spaceID).Return(false).AnyTimes() + mgr.configer = configer + quota := &entity.QuotaSpaceExpt{ExptID2RunTime: map[int64]int64{10: time.Now().Unix(), 11: time.Now().Unix()}} + mgr.quotaRepo.(*repoMocks.MockQuotaRepo).EXPECT().CreateOrUpdate(ctx, spaceID, gomock.Any(), session). + DoAndReturn(func(_ context.Context, _ int64, update func(*entity.QuotaSpaceExpt) (*entity.QuotaSpaceExpt, bool, error), _ *entity.Session) error { + _, changed, err := update(quota) + require.Equal(t, tc.publishFail, changed) + return err + }) + mgr.lwt.(*lwtMocks.MockILatestWriteTracker).EXPECT().CheckWriteFlagByID(ctx, gomock.Any(), exptID).Return(false).AnyTimes() + mgr.exptRepo.(*repoMocks.MockIExperimentRepo).EXPECT().MGetByID(ctx, []int64{exptID}, spaceID). + Return([]*entity.Experiment{{ID: exptID, SpaceID: spaceID, Status: entity.ExptStatus_Failed}}, nil).AnyTimes() + mgr.evaluationSetService.(*svcMocks.MockIEvaluationSetService).EXPECT(). + GetEvaluationSet(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Nil()).Return(&entity.EvaluationSet{}, nil).AnyTimes() + mgr.exptResultService.(*svcMocks.MockExptResultService).EXPECT(). + MGetStats(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mgr.exptAggrResultService.(*svcMocks.MockExptAggrResultService).EXPECT(). + BatchGetExptAggrResultByExperimentIDs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + + var runLog *entity.ExptRunLog + locked := false + lockKey := mgr.makeExptMutexLockKey(exptID) + mgr.mtr.(*metricsMocks.MockExptMetric).EXPECT().EmitExptExecRun(spaceID, int64(tc.mode)).Times(logCalls) + var latestRunID int64 + mgr.exptRepo.(*repoMocks.MockIExperimentRepo).EXPECT().Update(ctx, gomock.Any()). + DoAndReturn(func(_ context.Context, expt *entity.Experiment) error { latestRunID = expt.LatestRunID; return nil }).Times(logCalls) + if tc.mode == entity.EvaluationModeRetryItems { + mgr.idgenerator.(*idgenMocks.MockIIDGenerator).EXPECT().GenID(ctx).Return(runID, nil).Times(logCalls) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().BackoffLockWithValue(ctx, lockKey, strconv.FormatInt(runID, 10), gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, string, string, time.Duration, time.Duration) (bool, string, error) { + if locked { + return false, strconv.FormatInt(runID, 10), nil + } + locked = true + return true, "", nil + }).Times(logCalls) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Save(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, got *entity.ExptRunLog) error { runLog = got; return nil }).Times(logCalls) + } else { + mgr.mutex.(*lockMocks.MockILocker).EXPECT().LockBackoff(ctx, lockKey, gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, string, time.Duration, time.Duration) (bool, error) { + if locked { + return false, nil + } + locked = true + return true, nil + }).Times(logCalls) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Create(ctx, gomock.Any()). + DoAndReturn(func(_ context.Context, got *entity.ExptRunLog) error { runLog = got; return nil }).Times(logCalls) + } + if !tc.publishFail { + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), exptID, runID, gomock.Any()). + DoAndReturn(func(_ context.Context, _, _ int64, fields map[string]any) error { + runLog.Status = fields["status"].(int64) + runLog.StatusMessage = fields["status_message"].([]byte) + return nil + }) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), lockKey). + DoAndReturn(func(context.Context, string) (bool, error) { + require.Equal(t, int64(entity.ExptStatus_Failed), runLog.Status) + locked = false + return true, nil + }) + } + + var err error + if tc.mode == entity.EvaluationModeRetryItems { + _, retried, logErr := mgr.LogRetryItemsRun(ctx, exptID, tc.mode, spaceID, []int64{4}, session) + require.NoError(t, logErr) + require.False(t, retried) + err = mgr.RetryItems(ctx, exptID, runID, spaceID, 0, []int64{4}, session, nil) + } else { + require.NoError(t, mgr.LogRun(ctx, exptID, runID, tc.mode, spaceID, nil, session)) + err = mgr.Run(ctx, exptID, runID, spaceID, 0, session, tc.mode, nil) + } + if tc.publishFail { + require.ErrorIs(t, err, publishErr) + require.True(t, locked) + require.Equal(t, int64(entity.ExptStatus_Pending), runLog.Status) + return + } + status, ok := errorx.FromStatusError(err) + require.True(t, ok) + require.Equal(t, int32(errno.ExperimentRunningCountLimitCode), status.Code()) + require.Len(t, quota.ExptID2RunTime, 2) + require.NotContains(t, quota.ExptID2RunTime, exptID) + require.False(t, locked) + require.Equal(t, int64(entity.ExptStatus_Failed), runLog.Status) + require.Contains(t, string(runLog.StatusMessage), "max limit: 2") + require.Equal(t, runID, latestRunID, "admission cleanup must leave the latest run pointing at the failed attempt") + if tc.mode == entity.EvaluationModeRetryItems { + _, retried, logErr := mgr.LogRetryItemsRun(ctx, exptID, tc.mode, spaceID, []int64{4}, session) + require.NoError(t, logErr) + require.False(t, retried) + } else { + require.NoError(t, mgr.LogRun(ctx, exptID, runID+1, tc.mode, spaceID, nil, session)) + } + require.True(t, locked) + }) + } +} + +func TestRunPreparationReadFailure(t *testing.T) { + for _, tc := range []struct { + name string + updateFails, updateTimesOut, updateCommittedError bool + lockPresent, released bool + unlockErr error + }{ + {name: "success", lockPresent: true, released: true}, + {name: "update_failure_still_unlocks", updateFails: true, lockPresent: true, released: true}, + {name: "update_committed_then_error_still_unlocks", updateCommittedError: true, lockPresent: true, released: true}, + {name: "update_timeout_leaves_full_unlock_budget", updateTimesOut: true, lockPresent: true, released: true}, + {name: "lock_already_absent"}, + {name: "unlock_error_without_delete_preserves_original", lockPresent: true, unlockErr: errors.New("redis unavailable")}, + {name: "update_and_unlock_fail_preserves_original", updateFails: true, lockPresent: true, unlockErr: errors.New("redis unavailable")}, + } { + t.Run(tc.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + const exptID, runID, spaceID = int64(1), int64(2), int64(3) + mgr.lwt.(*lwtMocks.MockILatestWriteTracker).EXPECT().CheckWriteFlagByID(ctx, gomock.Any(), exptID).Return(false) + mgr.exptRepo.(*repoMocks.MockIExperimentRepo).EXPECT().MGetByID(ctx, []int64{exptID}, spaceID). + DoAndReturn(func(context.Context, []int64, int64) ([]*entity.Experiment, error) { + cancel() + return nil, context.Canceled + }) + runLog := &entity.ExptRunLog{ExptID: exptID, ExptRunID: runID, Status: int64(entity.ExptStatus_Pending)} + locked := tc.lockPresent + var updateCtx context.Context + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), exptID, runID, gomock.Any()). + DoAndReturn(func(cleanupCtx context.Context, _, _ int64, fields map[string]any) error { + updateCtx = cleanupCtx + require.NoError(t, cleanupCtx.Err()) + deadline, ok := cleanupCtx.Deadline() + require.True(t, ok) + require.InDelta(t, 5, time.Until(deadline).Seconds(), 1) + if tc.updateTimesOut { + <-cleanupCtx.Done() + return cleanupCtx.Err() + } + if tc.updateFails { + return errors.New("run log unavailable") + } + runLog.Status = fields["status"].(int64) + runLog.StatusMessage = fields["status_message"].([]byte) + if tc.updateCommittedError { + return errors.New("run log response lost") + } + return nil + }) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(exptID)). + DoAndReturn(func(unlockCtx context.Context, _ string) (bool, error) { + require.NotSame(t, updateCtx, unlockCtx) + require.NoError(t, unlockCtx.Err()) + deadline, ok := unlockCtx.Deadline() + require.True(t, ok) + require.InDelta(t, 2, time.Until(deadline).Seconds(), 0.5) + if tc.unlockErr == nil { + locked = false + } + return tc.released, tc.unlockErr + }) + err := mgr.Run(ctx, exptID, runID, spaceID, 0, &entity.Session{UserID: "u"}, entity.EvaluationModeFailRetry, nil) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, tc.lockPresent && tc.unlockErr != nil, locked) + if !tc.updateFails && !tc.updateTimesOut { + require.Equal(t, int64(entity.ExptStatus_Failed), runLog.Status) + require.Equal(t, context.Canceled.Error(), string(runLog.StatusMessage)) + } + }) + } +} diff --git a/backend/modules/evaluation/domain/service/expt_run_logging_cleanup_test.go b/backend/modules/evaluation/domain/service/expt_run_logging_cleanup_test.go new file mode 100644 index 000000000..c10b6892b --- /dev/null +++ b/backend/modules/evaluation/domain/service/expt_run_logging_cleanup_test.go @@ -0,0 +1,265 @@ +// Copyright (c) 2025 coze-dev Authors +// SPDX-License-Identifier: Apache-2.0 + +package service + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + idgenMocks "github.com/coze-dev/coze-loop/backend/infra/idgen/mocks" + lockMocks "github.com/coze-dev/coze-loop/backend/infra/lock/mocks" + metricsMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/component/metrics/mocks" + "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/entity" + repoMocks "github.com/coze-dev/coze-loop/backend/modules/evaluation/domain/repo/mocks" +) + +func TestLogRunPersistenceFailureReleasesOwnedLock(t *testing.T) { + for _, tc := range []struct { + name, failure string + retryItems, committed bool + }{ + {"create_run_log", "persist", false, false}, + {"create_run_log_response_lost", "persist", false, true}, + {"update_latest_run", "latest", false, false}, + {"update_latest_run_response_lost", "latest", false, true}, + {"save_retry_items_run", "persist", true, false}, + {"save_lost_response_new_run", "persist", true, true}, + {"update_retry_items_latest_run", "latest", true, false}, + {"update_retry_items_latest_run_response_lost", "latest", true, true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + ctx := context.Background() + const exptID, runID, nextRunID, spaceID = int64(1), int64(2), int64(3), int64(4) + const oldRunID = int64(5) + latestRunID := oldRunID + session := &entity.Session{UserID: "u"} + writeErr := errors.New("persistence unavailable") + locked, attempts := false, 0 + var saved *entity.ExptRunLog + persist := func(_ context.Context, run *entity.ExptRunLog) error { + if attempts == 1 && tc.failure == "persist" { + if tc.committed { + saved = run + } + return writeErr + } + saved = run + return nil + } + acquire := func() bool { + if locked { + return false + } + attempts++ + locked = true + return true + } + if tc.retryItems { + mgr.idgenerator.(*idgenMocks.MockIIDGenerator).EXPECT().GenID(ctx). + DoAndReturn(func(context.Context) (int64, error) { return runID + int64(attempts), nil }).Times(2) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().BackoffLockWithValue(ctx, mgr.makeExptMutexLockKey(exptID), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, string, string, time.Duration, time.Duration) (bool, string, error) { + return acquire(), "2", nil + }).Times(2) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Save(ctx, gomock.Any()).DoAndReturn(persist).Times(2) + mgr.mtr.(*metricsMocks.MockExptMetric).EXPECT().EmitExptExecRun(spaceID, int64(entity.EvaluationModeRetryItems)) + } else { + mgr.mutex.(*lockMocks.MockILocker).EXPECT().LockBackoff(ctx, mgr.makeExptMutexLockKey(exptID), gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, string, time.Duration, time.Duration) (bool, error) { return acquire(), nil }).Times(2) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Create(ctx, gomock.Any()).DoAndReturn(persist).Times(2) + mgr.mtr.(*metricsMocks.MockExptMetric).EXPECT().EmitExptExecRun(spaceID, int64(entity.EvaluationModeFailRetry)).Times(2) + } + updates := 1 + if tc.failure == "latest" { + updates++ + } + mgr.exptRepo.(*repoMocks.MockIExperimentRepo).EXPECT().Update(ctx, gomock.Any()). + DoAndReturn(func(_ context.Context, expt *entity.Experiment) error { + if attempts == 1 { + if tc.committed { + latestRunID = expt.LatestRunID + } + return writeErr + } + latestRunID = expt.LatestRunID + return nil + }).Times(updates) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), exptID, runID, gomock.Any()). + DoAndReturn(func(cleanupCtx context.Context, _, _ int64, fields map[string]any) error { + require.NoError(t, cleanupCtx.Err()) + require.Equal(t, map[string]any{"status": int64(entity.ExptStatus_Failed), "status_message": []byte(writeErr.Error())}, fields) + if saved == nil { + return nil + } + saved.Status = fields["status"].(int64) + saved.StatusMessage = fields["status_message"].([]byte) + return nil + }) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(exptID)). + DoAndReturn(func(context.Context, string) (bool, error) { locked = false; return true, nil }) + if tc.retryItems { + gotRunID, _, err := mgr.LogRetryItemsRun(ctx, exptID, entity.EvaluationModeRetryItems, spaceID, []int64{4}, session) + require.ErrorIs(t, err, writeErr) + require.Zero(t, gotRunID, "cleanup must retain the owned run ID even though the failed API returns zero") + } else { + require.ErrorIs(t, mgr.LogRun(ctx, exptID, runID, entity.EvaluationModeFailRetry, spaceID, nil, session), writeErr) + } + require.False(t, locked) + wantLatest := oldRunID + if tc.failure == "latest" && tc.committed { + wantLatest = runID + } + require.Equal(t, wantLatest, latestRunID, "cleanup must not overwrite the persisted experiment pointer") + if saved != nil { + require.Equal(t, int64(entity.ExptStatus_Failed), saved.Status) + require.Equal(t, writeErr.Error(), string(saved.StatusMessage)) + } + if tc.retryItems { + gotRunID, retried, err := mgr.LogRetryItemsRun(ctx, exptID, entity.EvaluationModeRetryItems, spaceID, []int64{4}, session) + require.NoError(t, err) + require.False(t, retried) + require.Equal(t, nextRunID, gotRunID) + } else { + require.NoError(t, mgr.LogRun(ctx, exptID, nextRunID, entity.EvaluationModeFailRetry, spaceID, nil, session)) + } + require.True(t, locked) + require.Equal(t, nextRunID, latestRunID) + }) + } +} + +func TestLogRetryItemsPersistenceFailurePreservesExistingRun(t *testing.T) { + for _, tc := range []struct { + name string + committed bool + }{ + {"save_failure_existing_run", false}, + {"save_lost_response_existing_run", true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + ctx := context.Background() + const exptID, existingRunID, proposedRunID, spaceID = int64(1), int64(2), int64(3), int64(4) + mgr.idgenerator.(*idgenMocks.MockIIDGenerator).EXPECT().GenID(ctx).Return(proposedRunID, nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().BackoffLockWithValue(ctx, mgr.makeExptMutexLockKey(exptID), "3", gomock.Any(), gomock.Any()).Return(false, "2", nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().Exists(ctx, gomock.Any()).Return(false, nil) + stored := &entity.ExptRunLog{ExptID: exptID, ExptRunID: existingRunID, Status: int64(entity.ExptStatus_Processing)} + loaded := *stored + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Get(ctx, exptID, existingRunID).Return(&loaded, nil) + writeErr := errors.New("append failed") + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Save(ctx, &loaded). + DoAndReturn(func(_ context.Context, run *entity.ExptRunLog) error { + if tc.committed { + *stored = *run + } + return writeErr + }) + runID, retried, err := mgr.LogRetryItemsRun(ctx, exptID, entity.EvaluationModeRetryItems, spaceID, []int64{5}, &entity.Session{UserID: "u"}) + require.ErrorIs(t, err, writeErr) + require.Zero(t, runID) + require.False(t, retried) + require.Equal(t, int64(entity.ExptStatus_Processing), stored.Status) + if tc.committed { + require.Len(t, stored.ItemIds, 1) + require.Equal(t, []int64{5}, stored.ItemIds[0].ItemIDs) + } else { + require.Empty(t, stored.ItemIds) + } + }) + } +} + +func TestLogRunLockErrorDoesNotCleanup(t *testing.T) { + for _, retryItems := range []bool{false, true} { + t.Run(map[bool]string{false: "LogRun", true: "LogRetryItemsRun"}[retryItems], func(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + ctx := context.Background() + lockErr := errors.New("lock outcome unknown") + if retryItems { + mgr.idgenerator.(*idgenMocks.MockIIDGenerator).EXPECT().GenID(ctx).Return(int64(2), nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().BackoffLockWithValue(ctx, gomock.Any(), "2", gomock.Any(), gomock.Any()).Return(true, "", lockErr) + _, _, err := mgr.LogRetryItemsRun(ctx, 1, entity.EvaluationModeRetryItems, 3, nil, &entity.Session{UserID: "u"}) + require.ErrorIs(t, err, lockErr) + } else { + mgr.mutex.(*lockMocks.MockILocker).EXPECT().LockBackoff(ctx, gomock.Any(), gomock.Any(), gomock.Any()).Return(true, lockErr) + require.ErrorIs(t, mgr.LogRun(ctx, 1, 2, entity.EvaluationModeFailRetry, 3, nil, &entity.Session{UserID: "u"}), lockErr) + } + }) + } +} + +func TestCleanupUnscheduledRunRejectsInvalidRunID(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + for _, runID := range []int64{0, -1} { + mgr.cleanupUnscheduledRun(context.Background(), 1, runID, errors.New("run failed")) + } +} + +func TestLogRunLatestFailureDoesNotRollbackSupersededQuota(t *testing.T) { + for _, retryItems := range []bool{false, true} { + for _, committed := range []bool{false, true} { + t.Run(fmt.Sprintf("retry_items_%v_committed_%v", retryItems, committed), func(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := newTestExptManager(ctrl) + ctx := context.Background() + const exptID, oldRunID, newRunID, spaceID = int64(1), int64(2), int64(3), int64(4) + guard := &fakeGuard{} + mgr.centralGuard = guard + expt := &entity.Experiment{ + ID: exptID, SpaceID: spaceID, LatestRunID: oldRunID, + SchedulerScope: testScope, ExptDispatchMode: entity.ExptDispatchModeEnforce, + } + if retryItems { + mgr.idgenerator.(*idgenMocks.MockIIDGenerator).EXPECT().GenID(ctx).Return(newRunID, nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().BackoffLockWithValue(ctx, gomock.Any(), "3", gomock.Any(), gomock.Any()).Return(true, "", nil) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Save(ctx, gomock.Any()).Return(nil) + } else { + mgr.mutex.(*lockMocks.MockILocker).EXPECT().LockBackoff(ctx, gomock.Any(), gomock.Any(), gomock.Any()).Return(true, nil) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Create(ctx, gomock.Any()).Return(nil) + mgr.mtr.(*metricsMocks.MockExptMetric).EXPECT().EmitExptExecRun(spaceID, int64(entity.EvaluationModeFailRetry)) + } + mgr.exptRepo.(*repoMocks.MockIExperimentRepo).EXPECT().GetByID(ctx, exptID, spaceID).Return(expt, nil) + mgr.itemResultRepo.(*repoMocks.MockIExptItemResultRepo).EXPECT(). + ScanItemRunLogs(ctx, exptID, oldRunID, gomock.Any(), gomock.Any(), gomock.Any(), spaceID). + Return([]*entity.ExptItemResultRunLog{{ItemID: 11}}, int64(1), nil) + writeErr := errors.New("latest run write failed") + mgr.exptRepo.(*repoMocks.MockIExperimentRepo).EXPECT().Update(ctx, gomock.Any()). + DoAndReturn(func(_ context.Context, update *entity.Experiment) error { + // Releasing the old quota precedes this write; these operations are not a transaction. + require.Len(t, guard.releases(), 1) + if committed { + expt.LatestRunID = update.LatestRunID + } + return writeErr + }) + mgr.runLogRepo.(*repoMocks.MockIExptRunLogRepo).EXPECT().Update(gomock.Any(), exptID, newRunID, gomock.Any()).Return(nil) + mgr.mutex.(*lockMocks.MockILocker).EXPECT().UnlockForce(gomock.Any(), mgr.makeExptMutexLockKey(exptID)).Return(true, nil) + if retryItems { + _, _, err := mgr.LogRetryItemsRun(ctx, exptID, entity.EvaluationModeRetryItems, spaceID, nil, &entity.Session{UserID: "u"}) + require.ErrorIs(t, err, writeErr) + } else { + require.ErrorIs(t, mgr.LogRun(ctx, exptID, newRunID, entity.EvaluationModeFailRetry, spaceID, nil, &entity.Session{UserID: "u"}), writeErr) + } + require.Len(t, guard.releases(), 1) + require.Equal(t, oldRunID, guard.releases()[0].RunID) + wantLatest := oldRunID + if committed { + wantLatest = newRunID + } + require.Equal(t, wantLatest, expt.LatestRunID) + }) + } + } +}