Skip to content

Commit 7e1aa14

Browse files
yokowuyokowu
committed
优化虚拟机回收活动校验
1 parent e11c007 commit 7e1aa14

2 files changed

Lines changed: 91 additions & 24 deletions

File tree

‎backend/biz/vmidle/usecase/policy_test.go‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,46 @@ func TestRefreshCachesNotFoundVM(t *testing.T) {
109109
}
110110
}
111111

112+
func TestShouldSkipRecycleForRecentTaskLastActiveAtSkipsClickHouse(t *testing.T) {
113+
ctx := context.Background()
114+
now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC)
115+
taskID := uuid.New()
116+
vm := &db.VirtualMachine{ID: "vm-recent-pg-activity", UserID: uuid.New()}
117+
vm.Edges.Tasks = []*db.Task{{
118+
ID: taskID,
119+
Status: consts.TaskStatusProcessing,
120+
LastActiveAt: now.Add(-time.Minute),
121+
}}
122+
redisClient := newTestRedis(t)
123+
logger := slog.Default()
124+
repo := &refreshHostRepoStub{vm: vm}
125+
activity := &taskLogActivityStub{err: errors.New("unexpected clickhouse query")}
126+
r := &vmIdleRefresher{
127+
cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}},
128+
redis: redisClient,
129+
logger: logger,
130+
hostRepo: repo,
131+
taskLogActivity: activity,
132+
sleepQueue: delayqueue.NewVMSleepQueue(redisClient, logger),
133+
notifyQueue: delayqueue.NewVMNotifyQueue(redisClient, logger),
134+
recycleQueue: delayqueue.NewVMRecycleQueue(redisClient, logger),
135+
}
136+
137+
skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now)
138+
if err != nil {
139+
t.Fatal(err)
140+
}
141+
if !skip {
142+
t.Fatal("expected recent task last_active_at to skip recycle")
143+
}
144+
if activity.calls != 0 {
145+
t.Fatalf("clickhouse calls = %d, want 0", activity.calls)
146+
}
147+
if repo.getVirtualMachineCalls != 1 {
148+
t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls)
149+
}
150+
}
151+
112152
func TestShouldSkipRecycleForRecentTaskLogRefreshesVM(t *testing.T) {
113153
ctx := context.Background()
114154
now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC)
@@ -129,7 +169,7 @@ func TestShouldSkipRecycleForRecentTaskLogRefreshesVM(t *testing.T) {
129169
recycleQueue: delayqueue.NewVMRecycleQueue(redisClient, logger),
130170
}
131171

132-
skip, err := r.shouldSkipRecycleForRecentTaskLog(ctx, vm, now)
172+
skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now)
133173
if err != nil {
134174
t.Fatal(err)
135175
}
@@ -155,7 +195,7 @@ func TestShouldSkipRecycleForRecentTaskLogAllowsStaleVM(t *testing.T) {
155195
taskLogActivity: &taskLogActivityStub{latest: map[uuid.UUID]time.Time{taskID: now.Add(-2 * time.Hour)}},
156196
}
157197

158-
skip, err := r.shouldSkipRecycleForRecentTaskLog(ctx, vm, now)
198+
skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now)
159199
if err != nil {
160200
t.Fatal(err)
161201
}
@@ -202,9 +242,11 @@ type refreshHostRepoStub struct {
202242
type taskLogActivityStub struct {
203243
latest map[uuid.UUID]time.Time
204244
err error
245+
calls int
205246
}
206247

207248
func (s *taskLogActivityStub) LatestTaskLogTime(_ context.Context, taskID uuid.UUID) (time.Time, bool, error) {
249+
s.calls++
208250
if s.err != nil {
209251
return time.Time{}, false, s.err
210252
}

‎backend/biz/vmidle/usecase/vmidle.go‎

Lines changed: 47 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -464,7 +464,7 @@ func (r *vmIdleRefresher) recycleConsumer() {
464464
return nil
465465
}
466466

467-
skip, err := r.shouldSkipRecycleForRecentTaskLog(ctx, vm, time.Now())
467+
skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, time.Now())
468468
if err != nil {
469469
return err
470470
}
@@ -498,8 +498,8 @@ func (r *vmIdleRefresher) recycleConsumer() {
498498
}
499499
}
500500

501-
func (r *vmIdleRefresher) shouldSkipRecycleForRecentTaskLog(ctx context.Context, vm *db.VirtualMachine, now time.Time) (bool, error) {
502-
if r.taskLogActivity == nil || vm == nil {
501+
func (r *vmIdleRefresher) shouldSkipRecycleForRecentActivity(ctx context.Context, vm *db.VirtualMachine, now time.Time) (bool, error) {
502+
if vm == nil {
503503
return false, nil
504504
}
505505
policy, err := r.resolvePolicyForVM(ctx, vm)
@@ -510,40 +510,55 @@ func (r *vmIdleRefresher) shouldSkipRecycleForRecentTaskLog(ctx context.Context,
510510
return false, nil
511511
}
512512

513-
taskIDs, err := r.activeTaskIDs(ctx, vm)
513+
activeTasks, err := r.activeRecycleTasks(ctx, vm)
514514
if err != nil {
515515
return false, err
516516
}
517-
if len(taskIDs) == 0 {
517+
if len(activeTasks) == 0 {
518518
return false, nil
519519
}
520520

521521
cutoff := now.Add(-time.Duration(policy.EffectiveRecycleSeconds) * time.Second)
522-
for _, taskID := range taskIDs {
523-
latest, ok, err := r.taskLogActivity.LatestTaskLogTime(ctx, taskID)
522+
for _, tk := range activeTasks {
523+
if tk.LastActiveAt.IsZero() || !tk.LastActiveAt.After(cutoff) {
524+
continue
525+
}
526+
return r.skipRecycleAndRefresh(ctx, vm, tk.ID, tk.LastActiveAt, cutoff, "pg_task_last_active_at")
527+
}
528+
529+
if r.taskLogActivity == nil {
530+
return false, nil
531+
}
532+
for _, tk := range activeTasks {
533+
latest, ok, err := r.taskLogActivity.LatestTaskLogTime(ctx, tk.ID)
524534
if err != nil {
525-
return false, fmt.Errorf("get latest task log time for task %s: %w", taskID, err)
535+
return false, fmt.Errorf("get latest task log time for task %s: %w", tk.ID, err)
526536
}
527537
if !ok || !latest.After(cutoff) {
528538
continue
529539
}
530540

531-
r.logger.InfoContext(ctx, "skip vm recycle because task has recent log",
532-
"vm_id", vm.ID,
533-
"task_id", taskID.String(),
534-
"latest_log_at", latest.UTC().Format(time.RFC3339Nano),
535-
"cutoff", cutoff.UTC().Format(time.RFC3339Nano))
536-
if err := r.Refresh(ctx, vm.ID); err != nil {
537-
return false, fmt.Errorf("refresh vm idle timer after recent task log: %w", err)
538-
}
539-
return true, nil
541+
return r.skipRecycleAndRefresh(ctx, vm, tk.ID, latest, cutoff, "clickhouse_task_log")
540542
}
541543
return false, nil
542544
}
543545

544-
func (r *vmIdleRefresher) activeTaskIDs(ctx context.Context, vm *db.VirtualMachine) ([]uuid.UUID, error) {
546+
func (r *vmIdleRefresher) skipRecycleAndRefresh(ctx context.Context, vm *db.VirtualMachine, taskID uuid.UUID, activeAt, cutoff time.Time, source string) (bool, error) {
547+
r.logger.InfoContext(ctx, "skip vm recycle because task has recent activity",
548+
"vm_id", vm.ID,
549+
"task_id", taskID.String(),
550+
"source", source,
551+
"active_at", activeAt.UTC().Format(time.RFC3339Nano),
552+
"cutoff", cutoff.UTC().Format(time.RFC3339Nano))
553+
if err := r.Refresh(ctx, vm.ID); err != nil {
554+
return false, fmt.Errorf("refresh vm idle timer after recent task activity: %w", err)
555+
}
556+
return true, nil
557+
}
558+
559+
func (r *vmIdleRefresher) activeRecycleTasks(ctx context.Context, vm *db.VirtualMachine) ([]*db.Task, error) {
545560
seen := make(map[uuid.UUID]struct{}, len(vm.Edges.Tasks))
546-
taskIDs := make([]uuid.UUID, 0, len(vm.Edges.Tasks))
561+
activeTasks := make([]*db.Task, 0, len(vm.Edges.Tasks))
547562
for _, tk := range vm.Edges.Tasks {
548563
if tk == nil {
549564
continue
@@ -555,10 +570,10 @@ func (r *vmIdleRefresher) activeTaskIDs(ctx context.Context, vm *db.VirtualMachi
555570
continue
556571
}
557572
seen[tk.ID] = struct{}{}
558-
taskIDs = append(taskIDs, tk.ID)
573+
activeTasks = append(activeTasks, tk)
559574
}
560575
if len(vm.Edges.Tasks) > 0 {
561-
return taskIDs, nil
576+
return activeTasks, nil
562577
}
563578

564579
taskIDStr, err := r.hostRepo.GetTaskIDByVMID(ctx, vm.ID)
@@ -572,7 +587,17 @@ func (r *vmIdleRefresher) activeTaskIDs(ctx context.Context, vm *db.VirtualMachi
572587
if err != nil {
573588
return nil, fmt.Errorf("invalid task id %q: %w", taskIDStr, err)
574589
}
575-
return []uuid.UUID{taskID}, nil
590+
if r.taskRepo == nil {
591+
return []*db.Task{{ID: taskID}}, nil
592+
}
593+
tk, err := r.taskRepo.GetByID(ctx, taskID)
594+
if err != nil {
595+
return nil, fmt.Errorf("get task %s: %w", taskID, err)
596+
}
597+
if tk.Status == consts.TaskStatusFinished || tk.Status == consts.TaskStatusError {
598+
return nil, nil
599+
}
600+
return []*db.Task{tk}, nil
576601
}
577602

578603
func (r *vmIdleRefresher) markRecycledTasksFinished(ctx context.Context, vm *db.VirtualMachine) error {

0 commit comments

Comments
 (0)