From aa3783eee11113fa5ae9f64e0b5dfba49bdef2eb Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Mon, 13 Jul 2026 21:18:20 +0800 Subject: [PATCH 01/10] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20control=20=E4=BF=9D?= =?UTF-8?q?=E6=B4=BB=E5=BB=B6=E9=95=BF=20VM=20=E5=9B=9E=E6=94=B6=E6=97=B6?= =?UTF-8?q?=E9=97=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/biz/host/handler/v1/internal.go | 2 +- .../handler/v1/internal_vm_activity_test.go | 4 +- backend/biz/task/handler/v1/task_control.go | 8 +- .../biz/task/handler/v1/task_control_test.go | 4 +- backend/biz/task/usecase/task.go | 5 +- .../biz/task/usecase/task_activity_test.go | 126 ++++++++++- .../usecase/task_create_vm_preinsert_test.go | 3 +- backend/biz/vmidle/usecase/policy_test.go | 213 ++++++++++++++++-- backend/biz/vmidle/usecase/vmidle.go | 56 +++-- backend/domain/host.go | 2 +- 10 files changed, 377 insertions(+), 46 deletions(-) diff --git a/backend/biz/host/handler/v1/internal.go b/backend/biz/host/handler/v1/internal.go index bb6ca4c27..f76449c67 100644 --- a/backend/biz/host/handler/v1/internal.go +++ b/backend/biz/host/handler/v1/internal.go @@ -122,7 +122,7 @@ func (h *InternalHostHandler) VMActivity(c *web.Context, req VMActivityReq) erro if strings.TrimSpace(req.VMID) == "" { return errors.New("vm_id is required") } - if err := h.idleRefresher.Refresh(c.Request().Context(), req.VMID); err != nil { + if err := h.idleRefresher.RecordActivity(c.Request().Context(), req.VMID); err != nil { h.logger.WarnContext(c.Request().Context(), "failed to refresh vm idle timers on activity", "vm_id", req.VMID, "error", err) return err } diff --git a/backend/biz/host/handler/v1/internal_vm_activity_test.go b/backend/biz/host/handler/v1/internal_vm_activity_test.go index eca2875a4..469e988b0 100644 --- a/backend/biz/host/handler/v1/internal_vm_activity_test.go +++ b/backend/biz/host/handler/v1/internal_vm_activity_test.go @@ -69,7 +69,9 @@ type internalVMIdleRefresherStub struct { ch chan string } -func (s *internalVMIdleRefresherStub) Refresh(_ context.Context, vmID string) error { +func (s *internalVMIdleRefresherStub) KeepAwake(context.Context, string) error { return nil } + +func (s *internalVMIdleRefresherStub) RecordActivity(_ context.Context, vmID string) error { select { case s.ch <- vmID: default: diff --git a/backend/biz/task/handler/v1/task_control.go b/backend/biz/task/handler/v1/task_control.go index fbfd3c2db..19cce13ab 100644 --- a/backend/biz/task/handler/v1/task_control.go +++ b/backend/biz/task/handler/v1/task_control.go @@ -148,7 +148,7 @@ func (h *TaskHandler) Control(c *web.Context, req domain.TaskControlReq) error { // 连接建立：刷新空闲计时器 if vm := task.VirtualMachine; vm != nil { - if err := h.idleRefresher.Refresh(c.Request().Context(), vm.ID); err != nil { + if err := h.idleRefresher.KeepAwake(c.Request().Context(), vm.ID); err != nil { logger.WarnContext(c.Request().Context(), "failed to refresh idle timers on connect", "error", err) } @@ -174,7 +174,7 @@ func (h *TaskHandler) Control(c *web.Context, req domain.TaskControlReq) error { if vm := task.VirtualMachine; vm != nil && !h.controlConns.Has(taskID) { ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() - if err := h.idleRefresher.Refresh(ctx, vm.ID); err != nil { + if err := h.idleRefresher.KeepAwake(ctx, vm.ID); err != nil { logger.WarnContext(ctx, "failed to refresh idle timers on disconnect", "error", err) } } @@ -227,7 +227,7 @@ func (h *TaskHandler) controlPing(ctx context.Context, wsConn *ws.WebsocketManag // controlKeepAlive 定期刷新空闲计时器，防止 VM 被误判空闲 func (h *TaskHandler) controlKeepAlive(ctx context.Context, taskID uuid.UUID, vmID string) error { - if err := h.idleRefresher.Refresh(ctx, vmID); err != nil { + if err := h.idleRefresher.KeepAwake(ctx, vmID); err != nil { h.logger.WarnContext(ctx, "keepalive refresh failed", "vmID", vmID, "error", err) } if err := h.taskActivity.Refresh(ctx, taskID); err != nil { @@ -243,7 +243,7 @@ func (h *TaskHandler) controlKeepAlive(ctx context.Context, taskID uuid.UUID, vm case <-ctx.Done(): return ctx.Err() case <-idleTicker.C: - if err := h.idleRefresher.Refresh(ctx, vmID); err != nil { + if err := h.idleRefresher.KeepAwake(ctx, vmID); err != nil { h.logger.WarnContext(ctx, "keepalive refresh failed", "vmID", vmID, "error", err) } case <-activityTicker.C: diff --git a/backend/biz/task/handler/v1/task_control_test.go b/backend/biz/task/handler/v1/task_control_test.go index 7e807df45..67f7720aa 100644 --- a/backend/biz/task/handler/v1/task_control_test.go +++ b/backend/biz/task/handler/v1/task_control_test.go @@ -66,7 +66,7 @@ type testVMIdleRefresher struct { ch chan string } -func (r *testVMIdleRefresher) Refresh(_ context.Context, vmID string) error { +func (r *testVMIdleRefresher) KeepAwake(_ context.Context, vmID string) error { select { case r.ch <- vmID: default: @@ -74,6 +74,8 @@ func (r *testVMIdleRefresher) Refresh(_ context.Context, vmID string) error { return nil } +func (r *testVMIdleRefresher) RecordActivity(context.Context, string) error { return nil } + type testTaskActivityRefresher struct { ch chan uuid.UUID } diff --git a/backend/biz/task/usecase/task.go b/backend/biz/task/usecase/task.go index a3ffe22bf..c4741abde 100644 --- a/backend/biz/task/usecase/task.go +++ b/backend/biz/task/usecase/task.go @@ -432,6 +432,9 @@ func (a *TaskUsecase) Continue(ctx context.Context, user *domain.User, id uuid.U }); err != nil { return err } + if err := a.idleRefresher.RecordActivity(ctx, tk.VirtualMachine.ID); err != nil { + a.logger.WarnContext(ctx, "failed to refresh vm idle timers on user input", "task_id", id, "vm_id", tk.VirtualMachine.ID, "error", err) + } // 缓存最近一次 user-input，供通知推送使用 a.redis.Set(ctx, fmt.Sprintf("mcai:task:%s:last_input", id.String()), string(req.Content), 24*time.Hour) @@ -721,7 +724,7 @@ func (a *TaskUsecase) refreshCreatedTaskState(ctx context.Context, taskID uuid.U if vmID == "" { return } - if err := a.idleRefresher.Refresh(ctx, vmID); err != nil { + if err := a.idleRefresher.RecordActivity(ctx, vmID); err != nil { a.logger.WarnContext(ctx, "failed to refresh vm idle timers on create", "task_id", taskID, "vm_id", vmID, "error", err) } } diff --git a/backend/biz/task/usecase/task_activity_test.go b/backend/biz/task/usecase/task_activity_test.go index 845e180db..8c5b278f7 100644 --- a/backend/biz/task/usecase/task_activity_test.go +++ b/backend/biz/task/usecase/task_activity_test.go @@ -6,8 +6,16 @@ import ( "io" "log/slog" "testing" + "time" + "github.com/alicebob/miniredis/v2" "github.com/google/uuid" + "github.com/redis/go-redis/v9" + + "github.com/chaitin/MonkeyCode/backend/config" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/domain" + "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" ) func TestRefreshCreatedTaskStateAlwaysRefreshesIdleTimer(t *testing.T) { @@ -29,14 +37,88 @@ func TestRefreshCreatedTaskStateAlwaysRefreshesIdleTimer(t *testing.T) { if taskRefresher.taskID != taskID { t.Fatalf("task id = %s, want %s", taskRefresher.taskID, taskID) } - if !idleRefresher.called { + if !idleRefresher.recordActivityCalled { t.Fatal("expected vm idle refresher to be called") } + if idleRefresher.keepAwakeCalled { + t.Fatal("did not expect task creation to use keep awake") + } if idleRefresher.vmID != vmID { t.Fatalf("vm id = %s, want %s", idleRefresher.vmID, vmID) } } +func TestContinueRecordsVMActivityAfterTaskflowAcceptsInput(t *testing.T) { + manager := &continueTaskManagerStub{} + idleRefresher := &vmIdleRefresherStub{} + u, user, taskID := newContinueTaskUsecase(t, manager, idleRefresher) + + if err := u.Continue(context.Background(), user, taskID, domain.ContinueTaskReq{Content: []byte("继续执行")}); err != nil { + t.Fatal(err) + } + if !manager.called { + t.Fatal("expected taskflow continue to be called") + } + if !idleRefresher.recordActivityCalled || idleRefresher.vmID != "vm-continue" { + t.Fatalf("record activity called = %v, vm id = %q", idleRefresher.recordActivityCalled, idleRefresher.vmID) + } + if idleRefresher.keepAwakeCalled { + t.Fatal("did not expect user input to use keep awake") + } +} + +func TestContinueDoesNotRecordVMActivityWhenTaskflowRejectsInput(t *testing.T) { + wantErr := errors.New("taskflow unavailable") + manager := &continueTaskManagerStub{err: wantErr} + idleRefresher := &vmIdleRefresherStub{} + u, user, taskID := newContinueTaskUsecase(t, manager, idleRefresher) + + err := u.Continue(context.Background(), user, taskID, domain.ContinueTaskReq{Content: []byte("继续执行")}) + if !errors.Is(err, wantErr) { + t.Fatalf("Continue() error = %v, want %v", err, wantErr) + } + if idleRefresher.recordActivityCalled || idleRefresher.keepAwakeCalled { + t.Fatal("did not expect rejected input to refresh vm idle state") + } +} + +func newContinueTaskUsecase(t *testing.T, manager *continueTaskManagerStub, idleRefresher *vmIdleRefresherStub) (*TaskUsecase, *domain.User, uuid.UUID) { + t.Helper() + user := &domain.User{ID: uuid.New()} + taskID := uuid.New() + vm := &db.VirtualMachine{ + ID: "vm-continue", + EnvironmentID: "env-continue", + CreatedAt: time.Now(), + Edges: db.VirtualMachineEdges{ + Host: &db.Host{ID: "host-continue"}, + }, + } + repo := &continueTaskRepoStub{task: &db.Task{ + ID: taskID, + UserID: user.ID, + CreatedAt: time.Now(), + Edges: db.TaskEdges{Vms: []*db.VirtualMachine{vm}}, + }} + redisClient := newTaskActivityRedis(t) + return &TaskUsecase{ + cfg: &config.Config{}, + repo: repo, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + taskflow: &continueTaskflowStub{manager: manager}, + redis: redisClient, + idleRefresher: idleRefresher, + }, user, taskID +} + +func newTaskActivityRedis(t *testing.T) *redis.Client { + t.Helper() + srv := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + return client +} + type taskActivityRefresherStub struct { taskID uuid.UUID forceCalled bool @@ -54,13 +136,49 @@ func (s *taskActivityRefresherStub) ForceRefresh(_ context.Context, taskID uuid. } type vmIdleRefresherStub struct { - vmID string + vmID string + keepAwakeCalled bool + recordActivityCalled bool + err error +} + +func (s *vmIdleRefresherStub) KeepAwake(_ context.Context, vmID string) error { + s.vmID = vmID + s.keepAwakeCalled = true + return s.err +} + +func (s *vmIdleRefresherStub) RecordActivity(_ context.Context, vmID string) error { + s.vmID = vmID + s.recordActivityCalled = true + return s.err +} + +type continueTaskRepoStub struct { + domain.TaskRepo + task *db.Task +} + +func (s *continueTaskRepoStub) Info(context.Context, *domain.User, uuid.UUID, bool) (*db.Task, error) { + return s.task, nil +} + +type continueTaskflowStub struct { + taskflow.Clienter + manager taskflow.TaskManager +} + +func (s *continueTaskflowStub) TaskManager() taskflow.TaskManager { + return s.manager +} + +type continueTaskManagerStub struct { + taskflow.TaskManager called bool err error } -func (s *vmIdleRefresherStub) Refresh(_ context.Context, vmID string) error { - s.vmID = vmID +func (s *continueTaskManagerStub) Continue(context.Context, taskflow.TaskReq) error { s.called = true return s.err } diff --git a/backend/biz/task/usecase/task_create_vm_preinsert_test.go b/backend/biz/task/usecase/task_create_vm_preinsert_test.go index ae469b72d..0d8302048 100644 --- a/backend/biz/task/usecase/task_create_vm_preinsert_test.go +++ b/backend/biz/task/usecase/task_create_vm_preinsert_test.go @@ -192,6 +192,7 @@ func (noopTaskActivityRefresher) ForceRefresh(context.Context, uuid.UUID) error type noopVMIdleRefresher struct{} -func (noopVMIdleRefresher) Refresh(context.Context, string) error { return nil } +func (noopVMIdleRefresher) KeepAwake(context.Context, string) error { return nil } +func (noopVMIdleRefresher) RecordActivity(context.Context, string) error { return nil } var _ vmidle.VMIdleRefresher = noopVMIdleRefresher{} diff --git a/backend/biz/vmidle/usecase/policy_test.go b/backend/biz/vmidle/usecase/policy_test.go index 16979c3a6..c9c38c7d3 100644 --- a/backend/biz/vmidle/usecase/policy_test.go +++ b/backend/biz/vmidle/usecase/policy_test.go @@ -23,26 +23,30 @@ import ( "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" ) -func TestVMIdleSchedulePlanFromPolicy(t *testing.T) { +func TestVMIdleSchedulePlanUsesSingleNow(t *testing.T) { + now := time.Date(2026, 7, 13, 10, 0, 0, 123, time.UTC) policy := &domain.TeamTaskVMIdlePolicy{ TeamID: uuid.New(), - SleepEnabled: false, - EffectiveSleepSeconds: 0, + SleepEnabled: true, + EffectiveSleepSeconds: 600, RecycleEnabled: true, EffectiveRecycleSeconds: 3600, } schedules := []notifySchedule{{name: "default", lead: 600 * time.Second, leadSeconds: 600}} - got := buildVMIdleSchedulePlan(policy, schedules) - if got.SleepAt != nil { - t.Fatalf("sleep should be disabled: %#v", got.SleepAt) + got := buildVMIdleSchedulePlan(policy, schedules, now) + if got.SleepAt == nil || !got.SleepAt.Equal(now.Add(600*time.Second)) { + t.Fatalf("sleep at = %v, want %v", got.SleepAt, now.Add(600*time.Second)) } - if got.RecycleAt == nil { - t.Fatal("recycle should be scheduled") + if got.RecycleAt == nil || !got.RecycleAt.Equal(now.Add(3600*time.Second)) { + t.Fatalf("recycle at = %v, want %v", got.RecycleAt, now.Add(3600*time.Second)) } if len(got.NotifyJobs) != 1 || got.NotifyJobs[0].MemberSuffix != "default" { t.Fatalf("notify jobs = %#v", got.NotifyJobs) } + if !got.NotifyJobs[0].RunAt.Equal(now.Add(3000 * time.Second)) { + t.Fatalf("notify at = %v, want %v", got.NotifyJobs[0].RunAt, now.Add(3000*time.Second)) + } } func TestResolvePolicyForVMFallsBackToGlobalWhenTeamMissing(t *testing.T) { @@ -63,7 +67,7 @@ func TestResolvePolicyForVMFallsBackToGlobalWhenTeamMissing(t *testing.T) { } } -func TestRefreshDebouncesBeforeVMQuery(t *testing.T) { +func TestRecordActivityDebouncesBeforeVMQuery(t *testing.T) { ctx := context.Background() redisClient := newTestRedis(t) repo := &refreshHostRepoStub{ @@ -76,18 +80,50 @@ func TestRefreshDebouncesBeforeVMQuery(t *testing.T) { hostRepo: repo, } - if err := r.Refresh(ctx, "vm-activity"); err != nil { - t.Fatalf("first refresh: %v", err) + if err := r.RecordActivity(ctx, "vm-activity"); err != nil { + t.Fatalf("first activity: %v", err) } - if err := r.Refresh(ctx, "vm-activity"); err != nil { - t.Fatalf("second refresh: %v", err) + if err := r.RecordActivity(ctx, "vm-activity"); err != nil { + t.Fatalf("second activity: %v", err) } if repo.getVirtualMachineCalls != 1 { t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls) } } -func TestRefreshCachesNotFoundVM(t *testing.T) { +func TestKeepAwakeAndRecordActivityDebounceIndependently(t *testing.T) { + ctx := context.Background() + redisClient := newTestRedis(t) + repo := &refreshHostRepoStub{ + vm: &db.VirtualMachine{ID: "vm-activity", UserID: uuid.New()}, + } + r := &vmIdleRefresher{ + cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 604800}}, + redis: redisClient, + logger: slog.Default(), + hostRepo: repo, + } + + if err := r.KeepAwake(ctx, "vm-activity"); err != nil { + t.Fatalf("keep awake: %v", err) + } + if err := r.RecordActivity(ctx, "vm-activity"); err != nil { + t.Fatalf("record activity: %v", err) + } + if repo.getVirtualMachineCalls != 2 { + t.Fatalf("GetVirtualMachine calls = %d, want 2", repo.getVirtualMachineCalls) + } + for _, key := range []string{ + "vm:idle:debounce:vm-activity:keep-awake", + "vm:idle:debounce:vm-activity:activity", + } { + if exists, err := redisClient.Exists(ctx, key).Result(); err != nil || exists != 1 { + t.Fatalf("debounce key %q exists = %d, err = %v", key, exists, err) + } + } +} + +func TestRecordActivityCachesNotFoundVM(t *testing.T) { ctx := context.Background() redisClient := newTestRedis(t) repo := &refreshHostRepoStub{err: newVirtualMachineNotFoundErr(t)} @@ -98,17 +134,138 @@ func TestRefreshCachesNotFoundVM(t *testing.T) { hostRepo: repo, } - if err := r.Refresh(ctx, "missing-vm"); err != nil { - t.Fatalf("first refresh: %v", err) + if err := r.RecordActivity(ctx, "missing-vm"); err != nil { + t.Fatalf("first activity: %v", err) } - if err := r.Refresh(ctx, "missing-vm"); err != nil { - t.Fatalf("second refresh: %v", err) + if err := r.RecordActivity(ctx, "missing-vm"); err != nil { + t.Fatalf("second activity: %v", err) } if repo.getVirtualMachineCalls != 1 { t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls) } } +func TestKeepAwakeOnlyUpdatesSleepSchedule(t *testing.T) { + ctx := context.Background() + redisClient := newTestRedis(t) + vmID := "vm-keep-awake" + vm := &db.VirtualMachine{ID: vmID, UserID: uuid.New(), HostID: "host-1", EnvironmentID: "env-1"} + vm.Edges.Tasks = []*db.Task{{ID: uuid.New()}} + r := newQueueTestVMIdleRefresher(redisClient, &refreshHostRepoStub{vm: vm}) + oldRecycleAt := time.Date(2026, 7, 20, 10, 0, 0, 0, time.UTC) + payload := &domain.VmIdleInfo{VmID: vmID, RecycleAt: oldRecycleAt} + oldNotifyAt := oldRecycleAt.Add(-10 * time.Minute) + if _, err := r.notifyQueue.Enqueue(ctx, notifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { + t.Fatal(err) + } + if _, err := r.recycleQueue.Enqueue(ctx, recycleQueueKey, payload, oldRecycleAt, vmID); err != nil { + t.Fatal(err) + } + + before := time.Now() + if err := r.KeepAwake(ctx, vmID); err != nil { + t.Fatal(err) + } + after := time.Now() + + assertVMIdleJobBetween(t, r.sleepQueue.RedisDelayQueue, ctx, sleepQueueKey, vmID, before.Add(10*time.Minute), after.Add(10*time.Minute)) + assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, notifyQueueKey, vmID+":default", oldNotifyAt) + assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, recycleQueueKey, vmID, oldRecycleAt) +} + +func TestRecordActivityUpdatesAllSchedules(t *testing.T) { + ctx := context.Background() + redisClient := newTestRedis(t) + vmID := "vm-record-activity" + vm := &db.VirtualMachine{ID: vmID, UserID: uuid.New(), HostID: "host-1", EnvironmentID: "env-1"} + vm.Edges.Tasks = []*db.Task{{ID: uuid.New()}} + r := newQueueTestVMIdleRefresher(redisClient, &refreshHostRepoStub{vm: vm}) + + before := time.Now() + if err := r.RecordActivity(ctx, vmID); err != nil { + t.Fatal(err) + } + after := time.Now() + + assertVMIdleJobBetween(t, r.sleepQueue.RedisDelayQueue, ctx, sleepQueueKey, vmID, before.Add(10*time.Minute), after.Add(10*time.Minute)) + assertVMIdleJobBetween(t, r.notifyQueue.RedisDelayQueue, ctx, notifyQueueKey, vmID+":default", before.Add(50*time.Minute), after.Add(50*time.Minute)) + assertVMIdleJobBetween(t, r.recycleQueue.RedisDelayQueue, ctx, recycleQueueKey, vmID, before.Add(time.Hour), after.Add(time.Hour)) +} + +func TestKeepAwakeWithSleepDisabledOnlyRemovesSleepSchedule(t *testing.T) { + ctx := context.Background() + redisClient := newTestRedis(t) + vmID := "vm-sleep-disabled" + vm := &db.VirtualMachine{ID: vmID, UserID: uuid.New(), HostID: "host-1", EnvironmentID: "env-1"} + vm.Edges.Tasks = []*db.Task{{ID: uuid.New()}} + r := newQueueTestVMIdleRefresher(redisClient, &refreshHostRepoStub{vm: vm}) + r.teamPolicyRepo = &refreshTeamPolicyRepoStub{team: &db.Team{ + ID: uuid.New(), + TaskConcurrencyLimit: 3, + TaskVMSleepEnabled: false, + TaskVMRecycleEnabled: true, + TaskVMRecycleSeconds: 3600, + }} + oldRecycleAt := time.Date(2026, 7, 20, 10, 0, 0, 0, time.UTC) + payload := &domain.VmIdleInfo{VmID: vmID, RecycleAt: oldRecycleAt} + oldSleepAt := oldRecycleAt.Add(-50 * time.Minute) + oldNotifyAt := oldRecycleAt.Add(-10 * time.Minute) + if _, err := r.sleepQueue.Enqueue(ctx, sleepQueueKey, payload, oldSleepAt, vmID); err != nil { + t.Fatal(err) + } + if _, err := r.notifyQueue.Enqueue(ctx, notifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { + t.Fatal(err) + } + if _, err := r.recycleQueue.Enqueue(ctx, recycleQueueKey, payload, oldRecycleAt, vmID); err != nil { + t.Fatal(err) + } + + if err := r.KeepAwake(ctx, vmID); err != nil { + t.Fatal(err) + } + if _, _, ok, err := r.sleepQueue.GetJobInfo(ctx, sleepQueueKey, vmID); err != nil || ok { + t.Fatalf("sleep job ok = %v, err = %v, want removed", ok, err) + } + assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, notifyQueueKey, vmID+":default", oldNotifyAt) + assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, recycleQueueKey, vmID, oldRecycleAt) +} + +func newQueueTestVMIdleRefresher(redisClient *redis.Client, repo domain.HostRepo) *vmIdleRefresher { + logger := slog.Default() + return &vmIdleRefresher{ + cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}}, + redis: redisClient, + logger: logger, + hostRepo: repo, + sleepQueue: delayqueue.NewVMSleepQueue(redisClient, logger), + notifyQueue: delayqueue.NewVMNotifyQueue(redisClient, logger), + recycleQueue: delayqueue.NewVMRecycleQueue(redisClient, logger), + schedules: []notifySchedule{{name: "default", lead: 10 * time.Minute, leadSeconds: 600}}, + } +} + +func assertVMIdleJobAt(t *testing.T, queue *delayqueue.RedisDelayQueue[*domain.VmIdleInfo], ctx context.Context, queueKey, id string, want time.Time) { + t.Helper() + _, got, ok, err := queue.GetJobInfo(ctx, queueKey, id) + if err != nil || !ok { + t.Fatalf("job %q ok = %v, err = %v", id, ok, err) + } + if !got.Equal(want) { + t.Fatalf("job %q run at = %v, want %v", id, got, want) + } +} + +func assertVMIdleJobBetween(t *testing.T, queue *delayqueue.RedisDelayQueue[*domain.VmIdleInfo], ctx context.Context, queueKey, id string, earliest, latest time.Time) { + t.Helper() + _, got, ok, err := queue.GetJobInfo(ctx, queueKey, id) + if err != nil || !ok { + t.Fatalf("job %q ok = %v, err = %v", id, ok, err) + } + if got.Before(earliest.Add(-time.Second)) || got.After(latest.Add(time.Second)) { + t.Fatalf("job %q run at = %v, want between %v and %v", id, got, earliest, latest) + } +} + func TestShouldSkipRecycleForRecentTaskLastActiveAtSkipsClickHouse(t *testing.T) { ctx := context.Background() now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC) @@ -239,6 +396,26 @@ type refreshHostRepoStub struct { getVirtualMachineCalls int } +type refreshTeamPolicyRepoStub struct { + team *db.Team +} + +func (s *refreshTeamPolicyRepoStub) GetTeam(context.Context, uuid.UUID) (*db.Team, error) { + return nil, errors.New("not implemented") +} + +func (s *refreshTeamPolicyRepoStub) GetTeamByUserID(context.Context, uuid.UUID) (*db.Team, error) { + return s.team, nil +} + +func (s *refreshTeamPolicyRepoStub) UpdateTaskVMIdlePolicy(context.Context, uuid.UUID, *domain.UpdateTeamTaskVMIdlePolicyReq) (*db.Team, error) { + return nil, errors.New("not implemented") +} + +func (s *refreshTeamPolicyRepoStub) GetMember(context.Context, uuid.UUID, uuid.UUID) (*db.TeamMember, error) { + return nil, errors.New("not implemented") +} + type taskLogActivityStub struct { latest map[uuid.UUID]time.Time err error diff --git a/backend/biz/vmidle/usecase/vmidle.go b/backend/biz/vmidle/usecase/vmidle.go index 5a4b0bdd8..7b9778505 100644 --- a/backend/biz/vmidle/usecase/vmidle.go +++ b/backend/biz/vmidle/usecase/vmidle.go @@ -25,7 +25,8 @@ import ( ) type VMIdleRefresher interface { - Refresh(ctx context.Context, vmID string) error + KeepAwake(ctx context.Context, vmID string) error + RecordActivity(ctx context.Context, vmID string) error } const ( @@ -129,6 +130,13 @@ type vmIdleSchedulePlan struct { NotifyJobs []vmIdleNotifyJob } +type vmIdleRefreshMode string + +const ( + vmIdleKeepAwakeMode vmIdleRefreshMode = "keep-awake" + vmIdleActivityMode vmIdleRefreshMode = "activity" +) + func NewVMIdleRefresher(i *do.Injector) (VMIdleRefresher, error) { cfg := do.MustInvoke[*config.Config](i) taskLogActivity, err := do.Invoke[*clickhouse.Client](i) @@ -212,8 +220,7 @@ func (r *vmIdleRefresher) resolvePolicyForVM(ctx context.Context, vm *db.Virtual return domain.ResolveTeamTaskVMIdlePolicy(team, r.cfg.VMIdle) } -func buildVMIdleSchedulePlan(policy *domain.TeamTaskVMIdlePolicy, schedules []notifySchedule) vmIdleSchedulePlan { - now := time.Now() +func buildVMIdleSchedulePlan(policy *domain.TeamTaskVMIdlePolicy, schedules []notifySchedule, now time.Time) vmIdleSchedulePlan { var plan vmIdleSchedulePlan if policy.SleepEnabled { sleepAt := now.Add(time.Duration(policy.EffectiveSleepSeconds) * time.Second) @@ -237,7 +244,15 @@ func buildVMIdleSchedulePlan(policy *domain.TeamTaskVMIdlePolicy, schedules []no return plan } -func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { +func (r *vmIdleRefresher) KeepAwake(ctx context.Context, vmID string) error { + return r.refresh(ctx, vmID, vmIdleKeepAwakeMode) +} + +func (r *vmIdleRefresher) RecordActivity(ctx context.Context, vmID string) error { + return r.refresh(ctx, vmID, vmIdleActivityMode) +} + +func (r *vmIdleRefresher) refresh(ctx context.Context, vmID string, mode vmIdleRefreshMode) error { notFoundKey := fmt.Sprintf("vm:idle:not-found:%s", vmID) if exists, err := r.redis.Exists(ctx, notFoundKey).Result(); err == nil && exists > 0 { return nil @@ -245,7 +260,7 @@ func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { r.logger.WarnContext(ctx, "redis not found cache check failed", "vmID", vmID, "error", err) } - debounceKey := fmt.Sprintf("vm:idle:debounce:%s", vmID) + debounceKey := fmt.Sprintf("vm:idle:debounce:%s:%s", vmID, mode) ok, err := r.redis.SetNX(ctx, debounceKey, "1", vmIdleDebounceTTL).Result() if err != nil { r.logger.ErrorContext(ctx, "redis SetNX failed", "vmID", vmID, "error", err) @@ -278,7 +293,7 @@ func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { if err != nil { return err } - plan := buildVMIdleSchedulePlan(policy, r.schedules) + plan := buildVMIdleSchedulePlan(policy, r.schedules, time.Now()) recycleAt := time.Time{} if plan.RecycleAt != nil { recycleAt = *plan.RecycleAt @@ -290,16 +305,29 @@ func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { EnvID: vm.EnvironmentID, RecycleAt: recycleAt, } + if mode == vmIdleKeepAwakeMode { + return r.applySleepPlan(ctx, vmID, payload, plan) + } + return r.applyActivityPlan(ctx, vmID, payload, policy, plan) +} - var errs []error +func (r *vmIdleRefresher) applySleepPlan(ctx context.Context, vmID string, payload *domain.VmIdleInfo, plan vmIdleSchedulePlan) error { if plan.SleepAt != nil { if _, err := r.sleepQueue.Enqueue(ctx, sleepQueueKey, payload, *plan.SleepAt, vmID); err != nil { r.logger.ErrorContext(ctx, "failed to enqueue sleep", "error", err, "vmID", vmID) - errs = append(errs, fmt.Errorf("enqueue sleep: %w", err)) + return fmt.Errorf("enqueue sleep: %w", err) } } else if err := r.sleepQueue.Remove(ctx, sleepQueueKey, vmID); err != nil { r.logger.ErrorContext(ctx, "failed to remove sleep", "error", err, "vmID", vmID) - errs = append(errs, fmt.Errorf("remove sleep: %w", err)) + return fmt.Errorf("remove sleep: %w", err) + } + return nil +} + +func (r *vmIdleRefresher) applyActivityPlan(ctx context.Context, vmID string, payload *domain.VmIdleInfo, policy *domain.TeamTaskVMIdlePolicy, plan vmIdleSchedulePlan) error { + var errs []error + if err := r.applySleepPlan(ctx, vmID, payload, plan); err != nil { + errs = append(errs, err) } for _, s := range r.schedules { member := fmt.Sprintf("%s:%s", vmID, s.name) @@ -523,7 +551,7 @@ func (r *vmIdleRefresher) shouldSkipRecycleForRecentActivity(ctx context.Context if tk.LastActiveAt.IsZero() || !tk.LastActiveAt.After(cutoff) { continue } - return r.skipRecycleAndRefresh(ctx, vm, tk.ID, tk.LastActiveAt, cutoff, "pg_task_last_active_at") + return r.skipRecycleAndRecordActivity(ctx, vm, tk.ID, tk.LastActiveAt, cutoff, "pg_task_last_active_at") } if r.taskLogActivity == nil { @@ -538,20 +566,20 @@ func (r *vmIdleRefresher) shouldSkipRecycleForRecentActivity(ctx context.Context continue } - return r.skipRecycleAndRefresh(ctx, vm, tk.ID, latest, cutoff, "clickhouse_task_log") + return r.skipRecycleAndRecordActivity(ctx, vm, tk.ID, latest, cutoff, "clickhouse_task_log") } return false, nil } -func (r *vmIdleRefresher) skipRecycleAndRefresh(ctx context.Context, vm *db.VirtualMachine, taskID uuid.UUID, activeAt, cutoff time.Time, source string) (bool, error) { +func (r *vmIdleRefresher) skipRecycleAndRecordActivity(ctx context.Context, vm *db.VirtualMachine, taskID uuid.UUID, activeAt, cutoff time.Time, source string) (bool, error) { r.logger.InfoContext(ctx, "skip vm recycle because task has recent activity", "vm_id", vm.ID, "task_id", taskID.String(), "source", source, "active_at", activeAt.UTC().Format(time.RFC3339Nano), "cutoff", cutoff.UTC().Format(time.RFC3339Nano)) - if err := r.Refresh(ctx, vm.ID); err != nil { - return false, fmt.Errorf("refresh vm idle timer after recent task activity: %w", err) + if err := r.RecordActivity(ctx, vm.ID); err != nil { + return false, fmt.Errorf("record vm activity after recent task activity: %w", err) } return true, nil } diff --git a/backend/domain/host.go b/backend/domain/host.go index c25c270ad..663f66fe5 100644 --- a/backend/domain/host.go +++ b/backend/domain/host.go @@ -89,7 +89,7 @@ type VmIdleInfo struct { EnvID string `json:"env_id"` TaskID string `json:"task_id,omitempty"` // 关联的任务 ID，用于通知 Name string `json:"name,omitempty"` // 任务名称，用于通知内容 - // RecycleAt 是本次 Refresh 算出的预计回收时间。每次用户活动都会延长这个值， + // RecycleAt 是本次 RecordActivity 算出的预计回收时间。每次用户活动都会延长这个值， // consumer 把它编进 RefID，让每个回收窗口都能产生不同的 dedup key（否则 // dispatcher 会按 (subID, eventType, RefID) 把同一 task 的后续推送全部静默）。 RecycleAt time.Time `json:"recycle_at"` From 2e7dc429b263f650353b57284214314ca5bffb8c Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Mon, 13 Jul 2026 22:41:19 +0800 Subject: [PATCH 02/10] =?UTF-8?q?=E7=BB=9F=E4=B8=80=20VM=20=E5=9B=9E?= =?UTF-8?q?=E6=94=B6=E6=B5=81=E7=A8=8B=E5=B9=B6=E6=8F=90=E4=BE=9B=E8=BF=90?= =?UTF-8?q?=E7=BB=B4=E6=9F=A5=E8=AF=A2=E8=83=BD=E5=8A=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/biz/vmidle/usecase/policy_test.go | 143 +------- backend/biz/vmidle/usecase/vmidle.go | 207 +---------- backend/pkg/delayqueue/delayqueue.go | 22 +- backend/pkg/delayqueue/delayqueue_test.go | 47 +++ backend/pkg/lifecycle/vmrecyclehook.go | 136 ++----- backend/pkg/lifecycle/vmrecyclehook_test.go | 94 +++++ backend/pkg/lifecycle/vmtaskhook.go | 2 - backend/pkg/lifecycle/vmtaskhook_test.go | 6 +- backend/pkg/loki/client.go | 62 ++-- backend/pkg/loki/client_test.go | 67 +++- backend/pkg/register.go | 4 + backend/pkg/tasklog/clickhouse_provider.go | 25 ++ .../pkg/tasklog/clickhouse_provider_test.go | 47 +++ backend/pkg/tasklog/gateway.go | 8 + backend/pkg/tasklog/gateway_test.go | 20 ++ backend/pkg/tasklog/loki_provider.go | 7 + backend/pkg/tasklog/provider.go | 1 + backend/pkg/vmrecycle/analyzer.go | 333 ++++++++++++++++++ backend/pkg/vmrecycle/analyzer_test.go | 231 ++++++++++++ backend/pkg/vmrecycle/vmrecycle.go | 232 ++++++++++++ backend/pkg/vmrecycle/vmrecycle_test.go | 331 +++++++++++++++++ 21 files changed, 1552 insertions(+), 473 deletions(-) create mode 100644 backend/pkg/delayqueue/delayqueue_test.go create mode 100644 backend/pkg/lifecycle/vmrecyclehook_test.go create mode 100644 backend/pkg/vmrecycle/analyzer.go create mode 100644 backend/pkg/vmrecycle/analyzer_test.go create mode 100644 backend/pkg/vmrecycle/vmrecycle.go create mode 100644 backend/pkg/vmrecycle/vmrecycle_test.go diff --git a/backend/biz/vmidle/usecase/policy_test.go b/backend/biz/vmidle/usecase/policy_test.go index c9c38c7d3..61afcf102 100644 --- a/backend/biz/vmidle/usecase/policy_test.go +++ b/backend/biz/vmidle/usecase/policy_test.go @@ -14,13 +14,13 @@ import ( "github.com/redis/go-redis/v9" "github.com/chaitin/MonkeyCode/backend/config" - "github.com/chaitin/MonkeyCode/backend/consts" "github.com/chaitin/MonkeyCode/backend/db" "github.com/chaitin/MonkeyCode/backend/db/enttest" "github.com/chaitin/MonkeyCode/backend/db/virtualmachine" "github.com/chaitin/MonkeyCode/backend/domain" "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" + "github.com/chaitin/MonkeyCode/backend/pkg/vmrecycle" ) func TestVMIdleSchedulePlanUsesSingleNow(t *testing.T) { @@ -155,10 +155,10 @@ func TestKeepAwakeOnlyUpdatesSleepSchedule(t *testing.T) { oldRecycleAt := time.Date(2026, 7, 20, 10, 0, 0, 0, time.UTC) payload := &domain.VmIdleInfo{VmID: vmID, RecycleAt: oldRecycleAt} oldNotifyAt := oldRecycleAt.Add(-10 * time.Minute) - if _, err := r.notifyQueue.Enqueue(ctx, notifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { + if _, err := r.notifyQueue.Enqueue(ctx, vmrecycle.NotifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { t.Fatal(err) } - if _, err := r.recycleQueue.Enqueue(ctx, recycleQueueKey, payload, oldRecycleAt, vmID); err != nil { + if _, err := r.recycleQueue.Enqueue(ctx, vmrecycle.RecycleQueueKey, payload, oldRecycleAt, vmID); err != nil { t.Fatal(err) } @@ -168,9 +168,9 @@ func TestKeepAwakeOnlyUpdatesSleepSchedule(t *testing.T) { } after := time.Now() - assertVMIdleJobBetween(t, r.sleepQueue.RedisDelayQueue, ctx, sleepQueueKey, vmID, before.Add(10*time.Minute), after.Add(10*time.Minute)) - assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, notifyQueueKey, vmID+":default", oldNotifyAt) - assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, recycleQueueKey, vmID, oldRecycleAt) + assertVMIdleJobBetween(t, r.sleepQueue.RedisDelayQueue, ctx, vmrecycle.SleepQueueKey, vmID, before.Add(10*time.Minute), after.Add(10*time.Minute)) + assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, vmrecycle.NotifyQueueKey, vmID+":default", oldNotifyAt) + assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, vmrecycle.RecycleQueueKey, vmID, oldRecycleAt) } func TestRecordActivityUpdatesAllSchedules(t *testing.T) { @@ -187,9 +187,9 @@ func TestRecordActivityUpdatesAllSchedules(t *testing.T) { } after := time.Now() - assertVMIdleJobBetween(t, r.sleepQueue.RedisDelayQueue, ctx, sleepQueueKey, vmID, before.Add(10*time.Minute), after.Add(10*time.Minute)) - assertVMIdleJobBetween(t, r.notifyQueue.RedisDelayQueue, ctx, notifyQueueKey, vmID+":default", before.Add(50*time.Minute), after.Add(50*time.Minute)) - assertVMIdleJobBetween(t, r.recycleQueue.RedisDelayQueue, ctx, recycleQueueKey, vmID, before.Add(time.Hour), after.Add(time.Hour)) + assertVMIdleJobBetween(t, r.sleepQueue.RedisDelayQueue, ctx, vmrecycle.SleepQueueKey, vmID, before.Add(10*time.Minute), after.Add(10*time.Minute)) + assertVMIdleJobBetween(t, r.notifyQueue.RedisDelayQueue, ctx, vmrecycle.NotifyQueueKey, vmID+":default", before.Add(50*time.Minute), after.Add(50*time.Minute)) + assertVMIdleJobBetween(t, r.recycleQueue.RedisDelayQueue, ctx, vmrecycle.RecycleQueueKey, vmID, before.Add(time.Hour), after.Add(time.Hour)) } func TestKeepAwakeWithSleepDisabledOnlyRemovesSleepSchedule(t *testing.T) { @@ -210,24 +210,24 @@ func TestKeepAwakeWithSleepDisabledOnlyRemovesSleepSchedule(t *testing.T) { payload := &domain.VmIdleInfo{VmID: vmID, RecycleAt: oldRecycleAt} oldSleepAt := oldRecycleAt.Add(-50 * time.Minute) oldNotifyAt := oldRecycleAt.Add(-10 * time.Minute) - if _, err := r.sleepQueue.Enqueue(ctx, sleepQueueKey, payload, oldSleepAt, vmID); err != nil { + if _, err := r.sleepQueue.Enqueue(ctx, vmrecycle.SleepQueueKey, payload, oldSleepAt, vmID); err != nil { t.Fatal(err) } - if _, err := r.notifyQueue.Enqueue(ctx, notifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { + if _, err := r.notifyQueue.Enqueue(ctx, vmrecycle.NotifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { t.Fatal(err) } - if _, err := r.recycleQueue.Enqueue(ctx, recycleQueueKey, payload, oldRecycleAt, vmID); err != nil { + if _, err := r.recycleQueue.Enqueue(ctx, vmrecycle.RecycleQueueKey, payload, oldRecycleAt, vmID); err != nil { t.Fatal(err) } if err := r.KeepAwake(ctx, vmID); err != nil { t.Fatal(err) } - if _, _, ok, err := r.sleepQueue.GetJobInfo(ctx, sleepQueueKey, vmID); err != nil || ok { + if _, _, ok, err := r.sleepQueue.GetJobInfo(ctx, vmrecycle.SleepQueueKey, vmID); err != nil || ok { t.Fatalf("sleep job ok = %v, err = %v, want removed", ok, err) } - assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, notifyQueueKey, vmID+":default", oldNotifyAt) - assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, recycleQueueKey, vmID, oldRecycleAt) + assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, vmrecycle.NotifyQueueKey, vmID+":default", oldNotifyAt) + assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, vmrecycle.RecycleQueueKey, vmID, oldRecycleAt) } func newQueueTestVMIdleRefresher(redisClient *redis.Client, repo domain.HostRepo) *vmIdleRefresher { @@ -266,104 +266,6 @@ func assertVMIdleJobBetween(t *testing.T, queue *delayqueue.RedisDelayQueue[*dom } } -func TestShouldSkipRecycleForRecentTaskLastActiveAtSkipsClickHouse(t *testing.T) { - ctx := context.Background() - now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC) - taskID := uuid.New() - vm := &db.VirtualMachine{ID: "vm-recent-pg-activity", UserID: uuid.New()} - vm.Edges.Tasks = []*db.Task{{ - ID: taskID, - Status: consts.TaskStatusProcessing, - LastActiveAt: now.Add(-time.Minute), - }} - redisClient := newTestRedis(t) - logger := slog.Default() - repo := &refreshHostRepoStub{vm: vm} - activity := &taskLogActivityStub{err: errors.New("unexpected clickhouse query")} - r := &vmIdleRefresher{ - cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}}, - redis: redisClient, - logger: logger, - hostRepo: repo, - taskLogActivity: activity, - sleepQueue: delayqueue.NewVMSleepQueue(redisClient, logger), - notifyQueue: delayqueue.NewVMNotifyQueue(redisClient, logger), - recycleQueue: delayqueue.NewVMRecycleQueue(redisClient, logger), - } - - skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now) - if err != nil { - t.Fatal(err) - } - if !skip { - t.Fatal("expected recent task last_active_at to skip recycle") - } - if activity.calls != 0 { - t.Fatalf("clickhouse calls = %d, want 0", activity.calls) - } - if repo.getVirtualMachineCalls != 1 { - t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls) - } -} - -func TestShouldSkipRecycleForRecentTaskLogRefreshesVM(t *testing.T) { - ctx := context.Background() - now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC) - taskID := uuid.New() - vm := &db.VirtualMachine{ID: "vm-recent-log", UserID: uuid.New()} - vm.Edges.Tasks = []*db.Task{{ID: taskID, Status: consts.TaskStatusProcessing}} - redisClient := newTestRedis(t) - logger := slog.Default() - repo := &refreshHostRepoStub{vm: vm} - r := &vmIdleRefresher{ - cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}}, - redis: redisClient, - logger: logger, - hostRepo: repo, - taskLogActivity: &taskLogActivityStub{latest: map[uuid.UUID]time.Time{taskID: now.Add(-time.Minute)}}, - sleepQueue: delayqueue.NewVMSleepQueue(redisClient, logger), - notifyQueue: delayqueue.NewVMNotifyQueue(redisClient, logger), - recycleQueue: delayqueue.NewVMRecycleQueue(redisClient, logger), - } - - skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now) - if err != nil { - t.Fatal(err) - } - if !skip { - t.Fatal("expected recent task log to skip recycle") - } - if repo.getVirtualMachineCalls != 1 { - t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls) - } -} - -func TestShouldSkipRecycleForRecentTaskLogAllowsStaleVM(t *testing.T) { - ctx := context.Background() - now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC) - taskID := uuid.New() - vm := &db.VirtualMachine{ID: "vm-stale-log", UserID: uuid.New()} - vm.Edges.Tasks = []*db.Task{{ID: taskID, Status: consts.TaskStatusProcessing}} - repo := &refreshHostRepoStub{vm: vm} - r := &vmIdleRefresher{ - cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}}, - logger: slog.Default(), - hostRepo: repo, - taskLogActivity: &taskLogActivityStub{latest: map[uuid.UUID]time.Time{taskID: now.Add(-2 * time.Hour)}}, - } - - skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now) - if err != nil { - t.Fatal(err) - } - if skip { - t.Fatal("expected stale task log to allow recycle") - } - if repo.getVirtualMachineCalls != 0 { - t.Fatalf("GetVirtualMachine calls = %d, want 0", repo.getVirtualMachineCalls) - } -} - func newTestRedis(t *testing.T) *redis.Client { t.Helper() srv := miniredis.RunT(t) @@ -416,21 +318,6 @@ func (s *refreshTeamPolicyRepoStub) GetMember(context.Context, uuid.UUID, uuid.U return nil, errors.New("not implemented") } -type taskLogActivityStub struct { - latest map[uuid.UUID]time.Time - err error - calls int -} - -func (s *taskLogActivityStub) LatestTaskLogTime(_ context.Context, taskID uuid.UUID) (time.Time, bool, error) { - s.calls++ - if s.err != nil { - return time.Time{}, false, s.err - } - latest, ok := s.latest[taskID] - return latest, ok, nil -} - func (s *refreshHostRepoStub) List(context.Context, uuid.UUID) ([]*db.Host, error) { return nil, errors.New("not implemented") } diff --git a/backend/biz/vmidle/usecase/vmidle.go b/backend/biz/vmidle/usecase/vmidle.go index 7b9778505..5da275928 100644 --- a/backend/biz/vmidle/usecase/vmidle.go +++ b/backend/biz/vmidle/usecase/vmidle.go @@ -17,11 +17,10 @@ import ( "github.com/chaitin/MonkeyCode/backend/consts" "github.com/chaitin/MonkeyCode/backend/db" "github.com/chaitin/MonkeyCode/backend/domain" - "github.com/chaitin/MonkeyCode/backend/pkg/clickhouse" "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" - "github.com/chaitin/MonkeyCode/backend/pkg/entx" "github.com/chaitin/MonkeyCode/backend/pkg/notify/dispatcher" "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" + "github.com/chaitin/MonkeyCode/backend/pkg/vmrecycle" ) type VMIdleRefresher interface { @@ -30,10 +29,6 @@ type VMIdleRefresher interface { } const ( - sleepQueueKey = "vm:idle:sleep" - notifyQueueKey = "vm:idle:notify" - recycleQueueKey = "vm:idle:recycle" - vmIdleDebounceTTL = 30 * time.Second vmNotFoundTTL = 30 * time.Second ) @@ -106,18 +101,14 @@ type vmIdleRefresher struct { hostRepo domain.HostRepo taskRepo domain.TaskRepo teamPolicyRepo domain.TeamPolicyRepo - taskLogActivity taskLogActivityReader notifyDispatcher *dispatcher.Dispatcher sleepQueue *delayqueue.VMSleepQueue notifyQueue *delayqueue.VMNotifyQueue recycleQueue *delayqueue.VMRecycleQueue + recycler vmrecycle.Recycler schedules []notifySchedule } -type taskLogActivityReader interface { - LatestTaskLogTime(ctx context.Context, taskID uuid.UUID) (time.Time, bool, error) -} - type vmIdleNotifyJob struct { MemberSuffix string RunAt time.Time @@ -139,14 +130,6 @@ const ( func NewVMIdleRefresher(i *do.Injector) (VMIdleRefresher, error) { cfg := do.MustInvoke[*config.Config](i) - taskLogActivity, err := do.Invoke[*clickhouse.Client](i) - if err != nil { - return nil, err - } - var activityReader taskLogActivityReader - if taskLogActivity != nil { - activityReader = taskLogActivity - } r := &vmIdleRefresher{ cfg: cfg, redis: do.MustInvoke[*redis.Client](i), @@ -155,11 +138,11 @@ func NewVMIdleRefresher(i *do.Injector) (VMIdleRefresher, error) { hostRepo: do.MustInvoke[domain.HostRepo](i), taskRepo: do.MustInvoke[domain.TaskRepo](i), teamPolicyRepo: do.MustInvoke[domain.TeamPolicyRepo](i), - taskLogActivity: activityReader, notifyDispatcher: do.MustInvoke[*dispatcher.Dispatcher](i), sleepQueue: do.MustInvoke[*delayqueue.VMSleepQueue](i), notifyQueue: do.MustInvoke[*delayqueue.VMNotifyQueue](i), recycleQueue: do.MustInvoke[*delayqueue.VMRecycleQueue](i), + recycler: do.MustInvoke[vmrecycle.Recycler](i), schedules: buildNotifySchedules(cfg.VMIdle), } @@ -313,11 +296,11 @@ func (r *vmIdleRefresher) refresh(ctx context.Context, vmID string, mode vmIdleR func (r *vmIdleRefresher) applySleepPlan(ctx context.Context, vmID string, payload *domain.VmIdleInfo, plan vmIdleSchedulePlan) error { if plan.SleepAt != nil { - if _, err := r.sleepQueue.Enqueue(ctx, sleepQueueKey, payload, *plan.SleepAt, vmID); err != nil { + if _, err := r.sleepQueue.Enqueue(ctx, vmrecycle.SleepQueueKey, payload, *plan.SleepAt, vmID); err != nil { r.logger.ErrorContext(ctx, "failed to enqueue sleep", "error", err, "vmID", vmID) return fmt.Errorf("enqueue sleep: %w", err) } - } else if err := r.sleepQueue.Remove(ctx, sleepQueueKey, vmID); err != nil { + } else if err := r.sleepQueue.Remove(ctx, vmrecycle.SleepQueueKey, vmID); err != nil { r.logger.ErrorContext(ctx, "failed to remove sleep", "error", err, "vmID", vmID) return fmt.Errorf("remove sleep: %w", err) } @@ -332,7 +315,7 @@ func (r *vmIdleRefresher) applyActivityPlan(ctx context.Context, vmID string, pa for _, s := range r.schedules { member := fmt.Sprintf("%s:%s", vmID, s.name) if !policy.RecycleEnabled { - if err := r.notifyQueue.Remove(ctx, notifyQueueKey, member); err != nil { + if err := r.notifyQueue.Remove(ctx, vmrecycle.NotifyQueueKey, member); err != nil { r.logger.ErrorContext(ctx, "failed to remove notify", "error", err, "vm_id", vmID, "tier", s.name) errs = append(errs, fmt.Errorf("remove notify %s: %w", s.name, err)) } @@ -341,7 +324,7 @@ func (r *vmIdleRefresher) applyActivityPlan(ctx context.Context, vmID string, pa for _, job := range plan.NotifyJobs { // member key 带 tier name 后缀，让同一 VM 的不同档作业互不覆盖。 member := fmt.Sprintf("%s:%s", vmID, job.MemberSuffix) - if _, err := r.notifyQueue.Enqueue(ctx, notifyQueueKey, payload, job.RunAt, member); err != nil { + if _, err := r.notifyQueue.Enqueue(ctx, vmrecycle.NotifyQueueKey, payload, job.RunAt, member); err != nil { r.logger.ErrorContext(ctx, "failed to enqueue notify", "error", err, "vm_id", vmID, "tier", job.MemberSuffix) errs = append(errs, fmt.Errorf("enqueue notify %s: %w", job.MemberSuffix, err)) continue @@ -353,11 +336,11 @@ func (r *vmIdleRefresher) applyActivityPlan(ctx context.Context, vmID string, pa "fire_at", job.RunAt.Format(time.RFC3339)) } if plan.RecycleAt != nil { - if _, err := r.recycleQueue.Enqueue(ctx, recycleQueueKey, payload, *plan.RecycleAt, vmID); err != nil { + if _, err := r.recycleQueue.Enqueue(ctx, vmrecycle.RecycleQueueKey, payload, *plan.RecycleAt, vmID); err != nil { r.logger.ErrorContext(ctx, "failed to enqueue recycle", "error", err, "vmID", vmID) errs = append(errs, fmt.Errorf("enqueue recycle: %w", err)) } - } else if err := r.recycleQueue.Remove(ctx, recycleQueueKey, vmID); err != nil { + } else if err := r.recycleQueue.Remove(ctx, vmrecycle.RecycleQueueKey, vmID); err != nil { r.logger.ErrorContext(ctx, "failed to remove recycle", "error", err, "vmID", vmID) errs = append(errs, fmt.Errorf("remove recycle: %w", err)) } @@ -367,7 +350,7 @@ func (r *vmIdleRefresher) applyActivityPlan(ctx context.Context, vmID string, pa func (r *vmIdleRefresher) sleepConsumer() { logger := r.logger.With("fn", "sleepConsumer") for { - err := r.sleepQueue.StartConsumer(context.Background(), sleepQueueKey, + err := r.sleepQueue.StartConsumer(context.Background(), vmrecycle.SleepQueueKey, func(ctx context.Context, job *delayqueue.Job[*domain.VmIdleInfo]) error { logger.InfoContext(ctx, "vm idle sleep triggered", "vmID", job.Payload.VmID) vm, err := r.hostRepo.GetVirtualMachine(ctx, job.Payload.VmID) @@ -399,7 +382,7 @@ func (r *vmIdleRefresher) sleepConsumer() { func (r *vmIdleRefresher) notifyConsumer() { logger := r.logger.With("fn", "notifyConsumer") for { - err := r.notifyQueue.StartConsumer(context.Background(), notifyQueueKey, + err := r.notifyQueue.StartConsumer(context.Background(), vmrecycle.NotifyQueueKey, func(ctx context.Context, job *delayqueue.Job[*domain.VmIdleInfo]) error { // job.ID 形如 ":"，从中解出 tier 名再 lookup schedule。 s, ok := r.lookupSchedule(job.ID) @@ -476,179 +459,17 @@ func (r *vmIdleRefresher) lookupSchedule(jobID string) (notifySchedule, bool) { func (r *vmIdleRefresher) recycleConsumer() { logger := r.logger.With("fn", "recycleConsumer") for { - err := r.recycleQueue.StartConsumer(context.Background(), recycleQueueKey, + err := r.recycleQueue.StartConsumer(context.Background(), vmrecycle.RecycleQueueKey, func(ctx context.Context, job *delayqueue.Job[*domain.VmIdleInfo]) error { logger.InfoContext(ctx, "vm recycle triggered", "vmID", job.Payload.VmID) - - ctx = entx.SkipSoftDelete(ctx) - vm, err := r.hostRepo.GetVirtualMachine(ctx, job.Payload.VmID) - if err != nil { - if db.IsNotFound(err) { - return nil - } - return fmt.Errorf("get vm %s: %w", job.Payload.VmID, err) - } - if vm.IsRecycled { - return nil - } - - skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, time.Now()) - if err != nil { - return err - } - if skip { - return nil - } - - if err := r.hostRepo.UpdateVirtualMachine(ctx, vm.ID, func(vmuo *db.VirtualMachineUpdateOne) error { - vmuo.SetIsRecycled(true) - return nil - }); err != nil { - return err - } - - if err := r.markRecycledTasksFinished(ctx, vm); err != nil { - return err - } - - if err := r.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{ - UserID: vm.UserID.String(), - HostID: vm.HostID, - ID: vm.EnvironmentID, - }); err != nil { - return fmt.Errorf("delete vm %s: %w", vm.ID, err) - } - - return nil + _, err := r.recycler.Recycle(ctx, job.Payload.VmID) + return err }) logger.Warn("recycle consumer error, retrying...", "error", err) time.Sleep(10 * time.Second) } } -func (r *vmIdleRefresher) shouldSkipRecycleForRecentActivity(ctx context.Context, vm *db.VirtualMachine, now time.Time) (bool, error) { - if vm == nil { - return false, nil - } - policy, err := r.resolvePolicyForVM(ctx, vm) - if err != nil { - return false, err - } - if !policy.RecycleEnabled || policy.EffectiveRecycleSeconds <= 0 { - return false, nil - } - - activeTasks, err := r.activeRecycleTasks(ctx, vm) - if err != nil { - return false, err - } - if len(activeTasks) == 0 { - return false, nil - } - - cutoff := now.Add(-time.Duration(policy.EffectiveRecycleSeconds) * time.Second) - for _, tk := range activeTasks { - if tk.LastActiveAt.IsZero() || !tk.LastActiveAt.After(cutoff) { - continue - } - return r.skipRecycleAndRecordActivity(ctx, vm, tk.ID, tk.LastActiveAt, cutoff, "pg_task_last_active_at") - } - - if r.taskLogActivity == nil { - return false, nil - } - for _, tk := range activeTasks { - latest, ok, err := r.taskLogActivity.LatestTaskLogTime(ctx, tk.ID) - if err != nil { - return false, fmt.Errorf("get latest task log time for task %s: %w", tk.ID, err) - } - if !ok || !latest.After(cutoff) { - continue - } - - return r.skipRecycleAndRecordActivity(ctx, vm, tk.ID, latest, cutoff, "clickhouse_task_log") - } - return false, nil -} - -func (r *vmIdleRefresher) skipRecycleAndRecordActivity(ctx context.Context, vm *db.VirtualMachine, taskID uuid.UUID, activeAt, cutoff time.Time, source string) (bool, error) { - r.logger.InfoContext(ctx, "skip vm recycle because task has recent activity", - "vm_id", vm.ID, - "task_id", taskID.String(), - "source", source, - "active_at", activeAt.UTC().Format(time.RFC3339Nano), - "cutoff", cutoff.UTC().Format(time.RFC3339Nano)) - if err := r.RecordActivity(ctx, vm.ID); err != nil { - return false, fmt.Errorf("record vm activity after recent task activity: %w", err) - } - return true, nil -} - -func (r *vmIdleRefresher) activeRecycleTasks(ctx context.Context, vm *db.VirtualMachine) ([]*db.Task, error) { - seen := make(map[uuid.UUID]struct{}, len(vm.Edges.Tasks)) - activeTasks := make([]*db.Task, 0, len(vm.Edges.Tasks)) - for _, tk := range vm.Edges.Tasks { - if tk == nil { - continue - } - if tk.Status == consts.TaskStatusFinished || tk.Status == consts.TaskStatusError { - continue - } - if _, ok := seen[tk.ID]; ok { - continue - } - seen[tk.ID] = struct{}{} - activeTasks = append(activeTasks, tk) - } - if len(vm.Edges.Tasks) > 0 { - return activeTasks, nil - } - - taskIDStr, err := r.hostRepo.GetTaskIDByVMID(ctx, vm.ID) - if err != nil { - return nil, fmt.Errorf("get task id by vm %s: %w", vm.ID, err) - } - if taskIDStr == "" { - return nil, nil - } - taskID, err := uuid.Parse(taskIDStr) - if err != nil { - return nil, fmt.Errorf("invalid task id %q: %w", taskIDStr, err) - } - if r.taskRepo == nil { - return []*db.Task{{ID: taskID}}, nil - } - tk, err := r.taskRepo.GetByID(ctx, taskID) - if err != nil { - return nil, fmt.Errorf("get task %s: %w", taskID, err) - } - if tk.Status == consts.TaskStatusFinished || tk.Status == consts.TaskStatusError { - return nil, nil - } - return []*db.Task{tk}, nil -} - -func (r *vmIdleRefresher) markRecycledTasksFinished(ctx context.Context, vm *db.VirtualMachine) error { - var errs []error - for _, tk := range vm.Edges.Tasks { - if tk == nil { - continue - } - if tk.Status == consts.TaskStatusFinished || tk.Status == consts.TaskStatusError { - continue - } - err := r.taskRepo.Update(ctx, nil, tk.ID, func(up *db.TaskUpdateOne) error { - up.SetStatus(consts.TaskStatusFinished) - up.SetCompletedAt(time.Now()) - return nil - }) - if err != nil { - errs = append(errs, fmt.Errorf("update task %s: %w", tk.ID, err)) - } - } - return errors.Join(errs...) -} - func (r *vmIdleRefresher) buildRecycleNotifyEvent(ctx context.Context, vm *db.VirtualMachine, expiresAt time.Time) (*domain.NotifyEvent, error) { // 直接按 virtualmachine_id 查 task_virtualmachines 拿 task_id， // 不再依赖 vm.Edges.Tasks 的 eager load 链。 diff --git a/backend/pkg/delayqueue/delayqueue.go b/backend/pkg/delayqueue/delayqueue.go index 92baf2ae7..dabfc02ba 100644 --- a/backend/pkg/delayqueue/delayqueue.go +++ b/backend/pkg/delayqueue/delayqueue.go @@ -117,15 +117,10 @@ func (q *RedisDelayQueue[T]) Enqueue(ctx context.Context, queue string, payload // GetJobInfo 查询任务信息 func (q *RedisDelayQueue[T]) GetJobInfo(ctx context.Context, queue, id string) (*Job[T], time.Time, bool, error) { - zkey := q.zsetKey(queue) - score, err := q.rdb.ZScore(ctx, zkey, id).Result() - if err == redis.Nil { - return nil, time.Time{}, false, nil - } - if err != nil { - return nil, time.Time{}, false, err + runAt, ok, err := q.GetRunAt(ctx, queue, id) + if err != nil || !ok { + return nil, time.Time{}, ok, err } - runAt := time.UnixMilli(int64(score)) job, err := q.loadJob(ctx, queue, id) if err != nil { @@ -137,6 +132,17 @@ func (q *RedisDelayQueue[T]) GetJobInfo(ctx context.Context, queue, id string) ( return job, runAt, true, nil } +func (q *RedisDelayQueue[T]) GetRunAt(ctx context.Context, queue, id string) (time.Time, bool, error) { + score, err := q.rdb.ZScore(ctx, q.zsetKey(queue), id).Result() + if errors.Is(err, redis.Nil) { + return time.Time{}, false, nil + } + if err != nil { + return time.Time{}, false, err + } + return time.UnixMilli(int64(score)), true, nil +} + // Remove 移除任务 func (q *RedisDelayQueue[T]) Remove(ctx context.Context, queue, id string) error { if err := q.rdb.ZRem(ctx, q.zsetKey(queue), id).Err(); err != nil { diff --git a/backend/pkg/delayqueue/delayqueue_test.go b/backend/pkg/delayqueue/delayqueue_test.go new file mode 100644 index 000000000..f57f87bd1 --- /dev/null +++ b/backend/pkg/delayqueue/delayqueue_test.go @@ -0,0 +1,47 @@ +package delayqueue + +import ( + "context" + "io" + "log/slog" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" +) + +func TestGetRunAtReadsScoreWithoutPayload(t *testing.T) { + ctx := context.Background() + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + queue := NewRedisDelayQueue[string](rdb, slog.New(slog.NewTextHandler(io.Discard, nil)), WithPrefix[string]("test")) + runAt := time.Date(2026, 7, 13, 12, 0, 0, 123*int(time.Millisecond), time.UTC) + if _, err := queue.Enqueue(ctx, "recycle", "payload", runAt, "vm-1"); err != nil { + t.Fatal(err) + } + if err := rdb.Del(ctx, queue.jobKey("recycle", "vm-1")).Err(); err != nil { + t.Fatal(err) + } + + got, ok, err := queue.GetRunAt(ctx, "recycle", "vm-1") + if err != nil || !ok { + t.Fatalf("GetRunAt() ok = %v, err = %v", ok, err) + } + if !got.Equal(runAt) { + t.Fatalf("run at = %v, want %v", got, runAt) + } +} + +func TestGetRunAtReturnsMissing(t *testing.T) { + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + queue := NewRedisDelayQueue[string](rdb, slog.Default()) + + _, ok, err := queue.GetRunAt(context.Background(), "recycle", "missing") + if err != nil || ok { + t.Fatalf("GetRunAt() ok = %v, err = %v", ok, err) + } +} diff --git a/backend/pkg/lifecycle/vmrecyclehook.go b/backend/pkg/lifecycle/vmrecyclehook.go index a54c81372..01dcc537a 100644 --- a/backend/pkg/lifecycle/vmrecyclehook.go +++ b/backend/pkg/lifecycle/vmrecyclehook.go @@ -2,52 +2,29 @@ package lifecycle import ( "context" + "errors" "fmt" "log/slog" "time" - "github.com/redis/go-redis/v9" "github.com/samber/do" - "github.com/chaitin/MonkeyCode/backend/db" "github.com/chaitin/MonkeyCode/backend/domain" "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" - "github.com/chaitin/MonkeyCode/backend/pkg/entx" - "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" + "github.com/chaitin/MonkeyCode/backend/pkg/vmrecycle" ) -const ( - vmSleepQueueKey = "vm:idle:sleep" - vmNotifyQueueKey = "vm:idle:notify" - vmWechat2hQueueKey = "vm:idle:notify:wechat:2h" - vmWechat15mQueueKey = "vm:idle:notify:wechat:15m" - vmRecycleQueueKey = "vm:idle:recycle" - vmExpireQueueKey = "vm:expire" -) - -// VMRecycleHook VM 回收 Hook，负责删除 VM、清理队列和 Redis 键、标记 DB type VMRecycleHook struct { - taskflow taskflow.Clienter - redis *redis.Client - hostRepo domain.HostRepo - vmSleepQueue *delayqueue.VMSleepQueue - vmNotifyQueue *delayqueue.VMNotifyQueue - vmRecycleQueue *delayqueue.VMRecycleQueue - vmExpireQueue *delayqueue.VMExpireQueue - logger *slog.Logger + recycler vmrecycle.Recycler + recycleQueue *delayqueue.VMRecycleQueue + logger *slog.Logger } -// NewVMRecycleHook 创建 VM 回收 Hook func NewVMRecycleHook(i *do.Injector) *VMRecycleHook { return &VMRecycleHook{ - taskflow: do.MustInvoke[taskflow.Clienter](i), - redis: do.MustInvoke[*redis.Client](i), - hostRepo: do.MustInvoke[domain.HostRepo](i), - vmSleepQueue: do.MustInvoke[*delayqueue.VMSleepQueue](i), - vmNotifyQueue: do.MustInvoke[*delayqueue.VMNotifyQueue](i), - vmRecycleQueue: do.MustInvoke[*delayqueue.VMRecycleQueue](i), - vmExpireQueue: do.MustInvoke[*delayqueue.VMExpireQueue](i), - logger: do.MustInvoke[*slog.Logger](i).With("hook", "vm-recycle-hook"), + recycler: do.MustInvoke[vmrecycle.Recycler](i), + recycleQueue: do.MustInvoke[*delayqueue.VMRecycleQueue](i), + logger: do.MustInvoke[*slog.Logger](i).With("hook", "vm-recycle-hook"), } } @@ -55,96 +32,25 @@ func (h *VMRecycleHook) Name() string { return "vm-recycle-hook" } func (h *VMRecycleHook) Priority() int { return 100 } func (h *VMRecycleHook) Async() bool { return false } -func (h *VMRecycleHook) OnStateChange(ctx context.Context, vmID string, from, to VMState, metadata VMMetadata) error { +func (h *VMRecycleHook) OnStateChange(ctx context.Context, vmID string, _, to VMState, metadata VMMetadata) error { if to != VMStateRecycled { return nil } - logger := h.logger.With("vm_id", vmID, "task_id", metadata.TaskID) - logger.InfoContext(ctx, "recycling VM") - - // 1. 查询 VM 完整信息 - ctx = entx.SkipSoftDelete(ctx) - vm, err := h.hostRepo.GetVirtualMachine(ctx, vmID) - if err != nil { - logger.ErrorContext(ctx, "failed to get VM info", "error", err) - return nil // VM 不存在则跳过 - } - - if vm.IsRecycled { - logger.InfoContext(ctx, "VM already recycled, skipping") + if _, err := h.recycler.Recycle(ctx, vmID); err == nil { return nil - } - - // 2. 删除 VM - if err := h.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{ - UserID: metadata.UserID.String(), - HostID: vm.HostID, - ID: vm.EnvironmentID, - }); err != nil { - logger.ErrorContext(ctx, "failed to delete VM, falling back to recycle queue", "error", err) - h.enqueueRetry(ctx, vm, metadata) - return nil - } - - // 3-6. 清理操作（失败仅记录日志） - h.cleanup(ctx, logger, vm, metadata) - - return nil -} - -// enqueueRetry 将 VM 投入 vmRecycleQueue 进行重试 -func (h *VMRecycleHook) enqueueRetry(ctx context.Context, vm *db.VirtualMachine, metadata VMMetadata) { - taskID := "" - if metadata.TaskID != nil { - taskID = metadata.TaskID.String() - } - payload := &domain.VmIdleInfo{ - UID: metadata.UserID, - VmID: vm.ID, - HostID: vm.HostID, - EnvID: vm.EnvironmentID, - TaskID: taskID, - } - if _, err := h.vmRecycleQueue.Enqueue(ctx, vmRecycleQueueKey, payload, time.Now(), vm.ID); err != nil { - h.logger.ErrorContext(ctx, "failed to enqueue VM for retry", "vm_id", vm.ID, "error", err) - } -} - -// cleanup 清理 delay queue、Redis 键、标记 DB -func (h *VMRecycleHook) cleanup(ctx context.Context, logger *slog.Logger, vm *db.VirtualMachine, metadata VMMetadata) { - // 3. 清理 delay queue 条目 - _ = h.vmSleepQueue.Remove(ctx, vmSleepQueueKey, vm.ID) - _ = h.vmNotifyQueue.Remove(ctx, vmNotifyQueueKey, vm.ID) - _ = h.vmNotifyQueue.Remove(ctx, vmWechat2hQueueKey, vm.ID) - _ = h.vmNotifyQueue.Remove(ctx, vmWechat15mQueueKey, vm.ID) - _ = h.vmRecycleQueue.Remove(ctx, vmRecycleQueueKey, vm.ID) - _ = h.vmExpireQueue.Remove(ctx, vmExpireQueueKey, vm.ID) - - // 4. 清理 task 相关 Redis 键 - if metadata.TaskID != nil { - taskIDStr := metadata.TaskID.String() - if err := h.redis.Del(ctx, - fmt.Sprintf("task:create_req:%s", taskIDStr), - fmt.Sprintf("mcai:task:%s:last_input", taskIDStr), - ).Err(); err != nil { - logger.WarnContext(ctx, "failed to clean task redis keys", "error", err) + } else { + h.logger.WarnContext(ctx, "vm recycle failed, enqueueing retry", "vm_id", vmID, "error", err) + payload := &domain.VmIdleInfo{UID: metadata.UserID, VmID: vmID} + if metadata.TaskID != nil { + payload.TaskID = metadata.TaskID.String() + } + if _, queueErr := h.recycleQueue.Enqueue(ctx, vmrecycle.RecycleQueueKey, payload, time.Now(), vmID); queueErr != nil { + return errors.Join( + fmt.Errorf("recycle vm %s: %w", vmID, err), + fmt.Errorf("enqueue vm %s recycle retry: %w", vmID, queueErr), + ) } - } - - // 5. DB 标记 is_recycled = true - if err := h.hostRepo.UpdateVirtualMachine(ctx, vm.ID, func(vmuo *db.VirtualMachineUpdateOne) error { - vmuo.SetIsRecycled(true) return nil - }); err != nil { - logger.WarnContext(ctx, "failed to mark VM as recycled", "error", err) } - - // 6. 清理 lifecycle Redis 键（最后执行） - lifecycleKey := fmt.Sprintf("lifecycle:%s", vm.ID) - if err := h.redis.Del(ctx, lifecycleKey).Err(); err != nil { - logger.WarnContext(ctx, "failed to clean lifecycle key", "error", err) - } - - logger.InfoContext(ctx, "VM recycled successfully") } diff --git a/backend/pkg/lifecycle/vmrecyclehook_test.go b/backend/pkg/lifecycle/vmrecyclehook_test.go new file mode 100644 index 000000000..fe407e948 --- /dev/null +++ b/backend/pkg/lifecycle/vmrecyclehook_test.go @@ -0,0 +1,94 @@ +package lifecycle + +import ( + "context" + "errors" + "io" + "log/slog" + "testing" + + "github.com/alicebob/miniredis/v2" + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + + "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" + "github.com/chaitin/MonkeyCode/backend/pkg/vmrecycle" +) + +func TestVMRecycleHookCallsRecyclerSynchronously(t *testing.T) { + recycler := &lifecycleRecyclerStub{result: vmrecycle.Result{VMID: "vm-1", Status: vmrecycle.StatusRecycled}} + hook, _ := newVMRecycleHookTest(t, recycler) + + if err := hook.OnStateChange(context.Background(), "vm-1", VMStateRunning, VMStateRecycled, VMMetadata{}); err != nil { + t.Fatal(err) + } + if recycler.calls != 1 || recycler.vmID != "vm-1" { + t.Fatalf("recycler calls = %d, vm id = %q", recycler.calls, recycler.vmID) + } +} + +func TestVMRecycleHookEnqueuesRetryWhenRecyclerFails(t *testing.T) { + wantErr := errors.New("delete failed") + recycler := &lifecycleRecyclerStub{err: wantErr} + hook, queue := newVMRecycleHookTest(t, recycler) + taskID := uuid.New() + userID := uuid.New() + + if err := hook.OnStateChange(context.Background(), "vm-2", VMStateRunning, VMStateRecycled, VMMetadata{TaskID: &taskID, UserID: userID}); err != nil { + t.Fatal(err) + } + job, _, ok, err := queue.GetJobInfo(context.Background(), vmrecycle.RecycleQueueKey, "vm-2") + if err != nil || !ok { + t.Fatalf("retry job ok = %v, err = %v", ok, err) + } + if job.Payload.VmID != "vm-2" || job.Payload.UID != userID || job.Payload.TaskID != taskID.String() { + t.Fatalf("retry payload = %+v", job.Payload) + } +} + +func TestVMRecycleHookReturnsBothRecyclerAndQueueFailures(t *testing.T) { + recycleErr := errors.New("delete failed") + recycler := &lifecycleRecyclerStub{err: recycleErr} + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + queue := delayqueue.NewVMRecycleQueue(rdb, slog.Default()) + if err := rdb.Close(); err != nil { + t.Fatal(err) + } + hook := &VMRecycleHook{ + recycler: recycler, + recycleQueue: queue, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + + err := hook.OnStateChange(context.Background(), "vm-3", VMStateRunning, VMStateRecycled, VMMetadata{}) + if !errors.Is(err, recycleErr) { + t.Fatalf("error = %v, want recycle error", err) + } +} + +func newVMRecycleHookTest(t *testing.T, recycler vmrecycle.Recycler) (*VMRecycleHook, *delayqueue.VMRecycleQueue) { + t.Helper() + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + queue := delayqueue.NewVMRecycleQueue(rdb, slog.Default()) + return &VMRecycleHook{ + recycler: recycler, + recycleQueue: queue, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }, queue +} + +type lifecycleRecyclerStub struct { + result vmrecycle.Result + err error + calls int + vmID string +} + +func (s *lifecycleRecyclerStub) Recycle(_ context.Context, vmID string) (vmrecycle.Result, error) { + s.calls++ + s.vmID = vmID + return s.result, s.err +} diff --git a/backend/pkg/lifecycle/vmtaskhook.go b/backend/pkg/lifecycle/vmtaskhook.go index cd2283558..e92631b9d 100644 --- a/backend/pkg/lifecycle/vmtaskhook.go +++ b/backend/pkg/lifecycle/vmtaskhook.go @@ -38,8 +38,6 @@ func (h *VMTaskHook) OnStateChange(ctx context.Context, _ string, _ VMState, to target = consts.TaskStatusProcessing case VMStateFailed: target = consts.TaskStatusError - case VMStateRecycled: - target = consts.TaskStatusFinished default: return nil } diff --git a/backend/pkg/lifecycle/vmtaskhook_test.go b/backend/pkg/lifecycle/vmtaskhook_test.go index 0d57b6904..a217f02e6 100644 --- a/backend/pkg/lifecycle/vmtaskhook_test.go +++ b/backend/pkg/lifecycle/vmtaskhook_test.go @@ -58,7 +58,7 @@ func TestVMTaskHook_OnStateChange_FailedTransitionsTaskToError(t *testing.T) { } } -func TestVMTaskHook_OnStateChange_RecycledTransitionsTaskToFinished(t *testing.T) { +func TestVMTaskHook_OnStateChange_RecycledDoesNotTransitionTask(t *testing.T) { mr, err := miniredis.Run() if err != nil { t.Fatalf("miniredis.Run() error = %v", err) @@ -101,7 +101,7 @@ func TestVMTaskHook_OnStateChange_RecycledTransitionsTaskToFinished(t *testing.T if err != nil { t.Fatalf("taskLifecycle.GetState() error = %v", err) } - if state != consts.TaskStatusFinished { - t.Fatalf("task state = %s, want %s", state, consts.TaskStatusFinished) + if state != consts.TaskStatusProcessing { + t.Fatalf("task state = %s, want %s", state, consts.TaskStatusProcessing) } } diff --git a/backend/pkg/loki/client.go b/backend/pkg/loki/client.go index 0dd5a6bfa..4067e1559 100644 --- a/backend/pkg/loki/client.go +++ b/backend/pkg/loki/client.go @@ -597,7 +597,19 @@ func (c *Client) tailWebSocketSession( // 优先从 Loki structured metadata (labels) 中读取 event 字段，回退到解析 JSON body。 // end 为搜索的结束时间上界，零值表示 time.Now()。 func (c *Client) FindLastEvent(ctx context.Context, taskID string, event string, start, end time.Time) (time.Time, error) { + latest, _, err := c.FindLastEventIn(ctx, taskID, []string{event}, start, end) + return latest, err +} + +func (c *Client) FindLastEventIn(ctx context.Context, taskID string, events []string, start, end time.Time) (time.Time, bool, error) { const pageSize = 200 + if len(events) == 0 { + return time.Time{}, false, nil + } + eventSet := make(map[string]struct{}, len(events)) + for _, event := range events { + eventSet[event] = struct{}{} + } if end.IsZero() { end = time.Now() @@ -609,36 +621,40 @@ func (c *Client) FindLastEvent(ctx context.Context, taskID string, event string, for { entries, err := c.QueryByTaskID(ctx, taskID, start, end, pageSize, "backward") if err != nil { - return time.Time{}, fmt.Errorf("FindLastEvent query failed: %w", err) + return time.Time{}, false, fmt.Errorf("FindLastEventIn query failed: %w", err) } - - for _, entry := range entries { - // 优先从 labels（structured metadata）读取 - if ev, ok := entry.Labels["event"]; ok { - if ev == event { - return entry.Timestamp, nil - } - continue - } - // 回退：解析 JSON body - var chunk struct { - Event string `json:"event"` - } - if err := json.Unmarshal([]byte(entry.Line), &chunk); err != nil { - continue - } - if chunk.Event == event { - return entry.Timestamp, nil - } + if latest, ok := latestMatchingEvent(entries, eventSet); ok { + return latest, true, nil } if len(entries) < pageSize { - return time.Time{}, nil + return time.Time{}, false, nil } - // 继续往前扫描 - end = entries[len(entries)-1].Timestamp + end = entries[len(entries)-1].Timestamp.Add(-time.Nanosecond) + } +} + +func latestMatchingEvent(entries []LogEntry, events map[string]struct{}) (time.Time, bool) { + for _, entry := range entries { + if event, ok := entry.Labels["event"]; ok { + _, matched := events[event] + if matched { + return entry.Timestamp, true + } + continue + } + var chunk struct { + Event string `json:"event"` + } + if err := json.Unmarshal([]byte(entry.Line), &chunk); err != nil { + continue + } + if _, ok := events[chunk.Event]; ok { + return entry.Timestamp, true + } } + return time.Time{}, false } // FindLatestRoundStart 定位 attach 模式下最新论次的起点。 diff --git a/backend/pkg/loki/client_test.go b/backend/pkg/loki/client_test.go index cb9121544..f5e8b85ed 100644 --- a/backend/pkg/loki/client_test.go +++ b/backend/pkg/loki/client_test.go @@ -1,7 +1,11 @@ package loki import ( + "context" "encoding/json" + "net/http" + "net/http/httptest" + "strconv" "testing" "time" ) @@ -54,7 +58,7 @@ func TestFindLatestRoundStart(t *testing.T) { {Timestamp: base.Add(2 * time.Second), Line: marshalChunk(t, "task-started")}, {Timestamp: base.Add(3 * time.Second), Line: marshalChunk(t, "task-running")}, }, - want: taskCreatedAt, + want: taskCreatedAt, }, } @@ -98,3 +102,64 @@ func TestFilterEntriesByTimeWindow(t *testing.T) { } } } + +func TestLatestMatchingEventUsesLabelsAndJSONFallback(t *testing.T) { + base := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + events := map[string]struct{}{"user-input": {}, "task-event": {}} + entries := []LogEntry{ + {Timestamp: base.Add(3 * time.Minute), Labels: map[string]string{"event": "task-running"}, Line: marshalChunk(t, "task-event")}, + {Timestamp: base.Add(2 * time.Minute), Labels: map[string]string{"event": "task-event"}}, + {Timestamp: base.Add(time.Minute), Line: marshalChunk(t, "user-input")}, + } + + got, ok := latestMatchingEvent(entries, events) + if !ok || !got.Equal(base.Add(2*time.Minute)) { + t.Fatalf("latest = %v, ok = %v", got, ok) + } +} + +func TestLatestMatchingEventReturnsMissing(t *testing.T) { + entries := []LogEntry{{Timestamp: time.Now(), Line: marshalChunk(t, "task-running")}} + if _, ok := latestMatchingEvent(entries, map[string]struct{}{"task-event": {}}); ok { + t.Fatal("expected no matching event") + } +} + +func TestFindLastEventInPagesBackward(t *testing.T) { + base := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + requests := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + values := make([][]string, 0, 200) + event := "task-running" + if requests == 1 { + for i := range 200 { + ts := base.Add(-time.Duration(i) * time.Second) + values = append(values, []string{strconv.FormatInt(ts.UnixNano(), 10), marshalChunk(t, "task-running")}) + } + } else { + event = "task-event" + values = append(values, []string{strconv.FormatInt(base.Add(-201*time.Second).UnixNano(), 10), marshalChunk(t, "task-event")}) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "status": "success", + "data": map[string]any{ + "resultType": "streams", + "result": []any{map[string]any{ + "stream": map[string]string{"task_id": "task-1", "event": event}, + "values": values, + }}, + }, + }) + })) + defer server.Close() + + client := NewClient(server.URL) + got, ok, err := client.FindLastEventIn(context.Background(), "task-1", []string{"user-input", "task-event"}, base.Add(-time.Hour), base.Add(time.Second)) + if err != nil || !ok { + t.Fatalf("FindLastEventIn() ok = %v, err = %v", ok, err) + } + if !got.Equal(base.Add(-201*time.Second)) || requests != 2 { + t.Fatalf("latest = %v, requests = %d", got, requests) + } +} diff --git a/backend/pkg/register.go b/backend/pkg/register.go index 2b295e0d3..94bb2d854 100644 --- a/backend/pkg/register.go +++ b/backend/pkg/register.go @@ -34,6 +34,7 @@ import ( "github.com/chaitin/MonkeyCode/backend/pkg/tasker" "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" "github.com/chaitin/MonkeyCode/backend/pkg/tasklog" + "github.com/chaitin/MonkeyCode/backend/pkg/vmrecycle" "github.com/chaitin/MonkeyCode/backend/pkg/ws" ) @@ -191,6 +192,9 @@ func RegisterInfra(i *do.Injector, w ...*web.Web) error { return delayqueue.NewVMExpireQueue(r, l), nil }) + do.Provide(i, vmrecycle.NewRecycler) + do.Provide(i, vmrecycle.NewAnalyzer) + // Channel Registry（通知渠道） do.Provide(i, func(i *do.Injector) (*msgpush.WechatClient, error) { cfg := do.MustInvoke[*config.Config](i) diff --git a/backend/pkg/tasklog/clickhouse_provider.go b/backend/pkg/tasklog/clickhouse_provider.go index 3b76704ac..1a5b34216 100644 --- a/backend/pkg/tasklog/clickhouse_provider.go +++ b/backend/pkg/tasklog/clickhouse_provider.go @@ -5,6 +5,7 @@ import ( "database/sql" "fmt" "strconv" + "strings" "time" "github.com/google/uuid" @@ -24,6 +25,30 @@ func (p *ClickHouseProvider) Name() string { return "clickhouse" } +func (p *ClickHouseProvider) LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string) (time.Time, bool, error) { + if p.client == nil { + return time.Time{}, false, ErrProviderUnavailable + } + if len(events) == 0 { + return time.Time{}, false, nil + } + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(events)), ",") + query := fmt.Sprintf(`SELECT maxOrNull(ts) FROM %s WHERE task_id = ? AND ts >= ? AND ts <= ? AND event IN (%s)`, p.client.Table(), placeholders) + args := make([]any, 0, len(events)+3) + args = append(args, taskID, start, end) + for _, event := range events { + args = append(args, event) + } + var latest sql.NullTime + if err := p.client.QueryRowContext(ctx, query, args...).Scan(&latest); err != nil { + return time.Time{}, false, err + } + if !latest.Valid { + return time.Time{}, false, nil + } + return latest.Time.UTC(), true, nil +} + func (p *ClickHouseProvider) QueryLatestTurn(ctx context.Context, taskID uuid.UUID, taskCreatedAt, end time.Time) (*QueryLatestTurnResp, error) { if p.client == nil { return nil, ErrProviderUnavailable diff --git a/backend/pkg/tasklog/clickhouse_provider_test.go b/backend/pkg/tasklog/clickhouse_provider_test.go index 8107730f1..c507e296c 100644 --- a/backend/pkg/tasklog/clickhouse_provider_test.go +++ b/backend/pkg/tasklog/clickhouse_provider_test.go @@ -13,6 +13,53 @@ import ( "github.com/chaitin/MonkeyCode/backend/pkg/tasklog" ) +func TestClickHouseProviderLatestEventTimeFiltersEvents(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + taskID := uuid.New() + start := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC) + end := time.Date(2026, 7, 13, 0, 0, 0, 0, time.UTC) + want := end.Add(-time.Minute) + mock.ExpectQuery("SELECT maxOrNull\\(ts\\)[\\s\\S]*event IN \\(\\?,\\?\\)"). + WithArgs(taskID, start, end, "user-input", "task-event"). + WillReturnRows(sqlmock.NewRows([]string{"max"}).AddRow(want)) + + provider := tasklog.NewClickHouseProvider(clickhouse.NewWithDBAndTable(db, "task_logs_test")) + got, ok, err := provider.LatestEventTime(context.Background(), taskID, start, end, []string{"user-input", "task-event"}) + if err != nil || !ok { + t.Fatalf("LatestEventTime() ok = %v, err = %v", ok, err) + } + if !got.Equal(want) { + t.Fatalf("latest = %v, want %v", got, want) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestClickHouseProviderLatestEventTimeHandlesNoMatch(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + taskID := uuid.New() + start := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC) + end := time.Date(2026, 7, 13, 0, 0, 0, 0, time.UTC) + mock.ExpectQuery("SELECT maxOrNull\\(ts\\)"). + WithArgs(taskID, start, end, "task-event"). + WillReturnRows(sqlmock.NewRows([]string{"max"}).AddRow(nil)) + + provider := tasklog.NewClickHouseProvider(clickhouse.NewWithDBAndTable(db, "task_logs_test")) + _, ok, err := provider.LatestEventTime(context.Background(), taskID, start, end, []string{"task-event"}) + if err != nil || ok { + t.Fatalf("LatestEventTime() ok = %v, err = %v", ok, err) + } +} + func TestClickHouseProviderQueryLatestTurnUsesTurnSeqCursor(t *testing.T) { db, mock, err := sqlmock.New() if err != nil { diff --git a/backend/pkg/tasklog/gateway.go b/backend/pkg/tasklog/gateway.go index d452dbfe2..5a34013fe 100644 --- a/backend/pkg/tasklog/gateway.go +++ b/backend/pkg/tasklog/gateway.go @@ -16,6 +16,14 @@ type Gateway struct { ClickHouse Provider } +func (g *Gateway) LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string, store consts.LogStore) (time.Time, bool, error) { + p, err := g.providerByStore(store) + if err != nil { + return time.Time{}, false, err + } + return p.LatestEventTime(ctx, taskID, start, end, events) +} + func (g *Gateway) QueryLatestTurn(ctx context.Context, taskID uuid.UUID, taskCreatedAt, end time.Time, store consts.LogStore) (*QueryLatestTurnResp, error) { p, err := g.providerByStore(store) if err != nil { diff --git a/backend/pkg/tasklog/gateway_test.go b/backend/pkg/tasklog/gateway_test.go index ef2ca7518..fa7def7fa 100644 --- a/backend/pkg/tasklog/gateway_test.go +++ b/backend/pkg/tasklog/gateway_test.go @@ -16,6 +16,12 @@ type gatewayProviderStub struct { name string queryLatestTurnCalled bool queryTurnsCalled bool + latestEventTimeCalled bool +} + +func (s *gatewayProviderStub) LatestEventTime(context.Context, uuid.UUID, time.Time, time.Time, []string) (time.Time, bool, error) { + s.latestEventTimeCalled = true + return time.Now(), true, nil } func (s *gatewayProviderStub) Name() string { @@ -70,6 +76,20 @@ func TestGatewayClickHouseStoreUsesClickHouse(t *testing.T) { } } +func TestGatewayLatestEventTimeUsesConfiguredStore(t *testing.T) { + loki := &gatewayProviderStub{name: "loki"} + clickHouse := &gatewayProviderStub{name: "clickhouse"} + gateway := &Gateway{Loki: loki, ClickHouse: clickHouse} + + _, _, err := gateway.LatestEventTime(context.Background(), uuid.New(), time.Now().Add(-time.Hour), time.Now(), []string{"task-event"}, consts.LogStoreClickHouse) + if err != nil { + t.Fatal(err) + } + if !clickHouse.latestEventTimeCalled || loki.latestEventTimeCalled { + t.Fatalf("loki called = %v, clickhouse called = %v", loki.latestEventTimeCalled, clickHouse.latestEventTimeCalled) + } +} + func TestGatewayUnknownStoreReturnsError(t *testing.T) { loki := &gatewayProviderStub{name: "loki"} clickHouse := &gatewayProviderStub{name: "clickhouse"} diff --git a/backend/pkg/tasklog/loki_provider.go b/backend/pkg/tasklog/loki_provider.go index 1691331ca..94a8336dc 100644 --- a/backend/pkg/tasklog/loki_provider.go +++ b/backend/pkg/tasklog/loki_provider.go @@ -26,6 +26,13 @@ func (p *LokiProvider) Name() string { return "loki" } +func (p *LokiProvider) LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string) (time.Time, bool, error) { + if p.client == nil { + return time.Time{}, false, ErrProviderUnavailable + } + return p.client.FindLastEventIn(ctx, taskID.String(), events, start, end) +} + func (p *LokiProvider) QueryWindow(ctx context.Context, taskID uuid.UUID, start, end time.Time) ([]Entry, error) { if p.client == nil { return nil, ErrProviderUnavailable diff --git a/backend/pkg/tasklog/provider.go b/backend/pkg/tasklog/provider.go index 8f7ca4a80..f35c3a95c 100644 --- a/backend/pkg/tasklog/provider.go +++ b/backend/pkg/tasklog/provider.go @@ -9,6 +9,7 @@ import ( type Provider interface { Name() string + LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string) (time.Time, bool, error) QueryLatestTurn(ctx context.Context, taskID uuid.UUID, taskCreatedAt, end time.Time) (*QueryLatestTurnResp, error) QueryTurns(ctx context.Context, taskID uuid.UUID, taskCreatedAt time.Time, opts QueryTurnsOpts) (*QueryTurnsResp, error) QueryUserInputs(ctx context.Context, taskID uuid.UUID, taskCreatedAt time.Time, cursor string, limit int) (*QueryUserInputsResp, error) diff --git a/backend/pkg/vmrecycle/analyzer.go b/backend/pkg/vmrecycle/analyzer.go new file mode 100644 index 000000000..ffe385a97 --- /dev/null +++ b/backend/pkg/vmrecycle/analyzer.go @@ -0,0 +1,333 @@ +package vmrecycle + +import ( + "context" + "fmt" + "sort" + "sync" + "time" + + "entgo.io/ent/dialect/sql" + "github.com/google/uuid" + "github.com/samber/do" + + "github.com/chaitin/MonkeyCode/backend/config" + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/db/task" + "github.com/chaitin/MonkeyCode/backend/db/virtualmachine" + "github.com/chaitin/MonkeyCode/backend/domain" + "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" + "github.com/chaitin/MonkeyCode/backend/pkg/entx" + "github.com/chaitin/MonkeyCode/backend/pkg/tasklog" +) + +const ( + analyzePageSize = 100 + analyzeWorkerCount = 8 +) + +var activityEvents = []string{ + "user-input", + "task-started", + "task-running", + "task-ended", + "task-error", + "task-event", +} + +type Decision string + +const ( + DecisionCandidate Decision = "candidate" + DecisionUnavailable Decision = "unavailable" + DecisionHistoricalUncertain Decision = "historical_uncertain" + DecisionNotDue Decision = "not_due" + DecisionDisabled Decision = "disabled" + DecisionAlreadyRecycled Decision = "already_recycled" + DecisionNotFound Decision = "not_found" +) + +type TaskInfo struct { + ID uuid.UUID + Status consts.TaskStatus + LogStore consts.LogStore + CreatedAt time.Time +} + +type Analysis struct { + Target string + VMID string + Tasks []TaskInfo + LastActivityAt time.Time + RecycleSeconds int + DueAt time.Time + RedisRecycleAt *time.Time + Overdue time.Duration + Decision Decision + Reason string + AlreadyRecycled bool +} + +type ScanReport struct { + Items []Analysis + Counts map[Decision]int +} + +type Analyzer interface { + Scan(ctx context.Context) (ScanReport, error) + AnalyzeTargets(ctx context.Context, taskIDs []uuid.UUID, vmIDs []string) ([]Analysis, error) +} + +type activityReader interface { + LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string, store consts.LogStore) (time.Time, bool, error) +} + +type deadlineReader interface { + GetRunAt(ctx context.Context, queue, id string) (time.Time, bool, error) +} + +type analyzer struct { + cfg *config.Config + db *db.Client + teamPolicyRepo domain.TeamPolicyRepo + activity activityReader + deadlines deadlineReader + now func() time.Time +} + +func NewAnalyzer(i *do.Injector) (Analyzer, error) { + return &analyzer{ + cfg: do.MustInvoke[*config.Config](i), + db: do.MustInvoke[*db.Client](i), + teamPolicyRepo: do.MustInvoke[domain.TeamPolicyRepo](i), + activity: do.MustInvoke[*tasklog.Gateway](i), + deadlines: do.MustInvoke[*delayqueue.VMRecycleQueue](i), + now: time.Now, + }, nil +} + +func (a *analyzer) Scan(ctx context.Context) (ScanReport, error) { + report := ScanReport{Counts: make(map[Decision]int)} + cursor := "" + for { + query := a.db.VirtualMachine.Query(). + Where(virtualmachine.IsRecycled(false), virtualmachine.HasTasks()). + WithTasks(). + Order(db.Asc(virtualmachine.FieldID)). + Limit(analyzePageSize) + if cursor != "" { + query = query.Where(func(s *sql.Selector) { + s.Where(sql.GT(s.C(virtualmachine.FieldID), cursor)) + }) + } + vms, err := query.All(ctx) + if err != nil { + return report, err + } + if len(vms) == 0 { + break + } + items := a.analyzeBatch(ctx, vms) + for _, item := range items { + report.Counts[item.Decision]++ + if item.Decision == DecisionCandidate || item.Decision == DecisionUnavailable || item.Decision == DecisionHistoricalUncertain { + report.Items = append(report.Items, item) + } + } + cursor = vms[len(vms)-1].ID + if len(vms) < analyzePageSize { + break + } + } + sort.Slice(report.Items, func(i, j int) bool { return report.Items[i].VMID < report.Items[j].VMID }) + return report, nil +} + +func (a *analyzer) AnalyzeTargets(ctx context.Context, taskIDs []uuid.UUID, vmIDs []string) ([]Analysis, error) { + ctx = entx.SkipSoftDelete(ctx) + seen := make(map[string]struct{}, len(taskIDs)+len(vmIDs)) + items := make([]Analysis, 0, len(taskIDs)+len(vmIDs)) + for _, taskID := range taskIDs { + tk, err := a.db.Task.Query().Where(task.ID(taskID)).WithVms().Only(ctx) + if err != nil { + if db.IsNotFound(err) { + items = append(items, Analysis{Target: "task:" + taskID.String(), Decision: DecisionNotFound, Reason: "task not found"}) + continue + } + items = append(items, Analysis{Target: "task:" + taskID.String(), Decision: DecisionUnavailable, Reason: err.Error()}) + continue + } + if len(tk.Edges.Vms) == 0 { + items = append(items, Analysis{Target: "task:" + taskID.String(), Decision: DecisionNotFound, Reason: "task has no VM"}) + continue + } + for _, vm := range tk.Edges.Vms { + seen[vm.ID] = struct{}{} + } + } + for _, vmID := range vmIDs { + seen[vmID] = struct{}{} + } + uniqueIDs := make([]string, 0, len(seen)) + for vmID := range seen { + uniqueIDs = append(uniqueIDs, vmID) + } + sort.Strings(uniqueIDs) + for _, vmID := range uniqueIDs { + vm, err := a.loadVM(ctx, vmID) + if err != nil { + if db.IsNotFound(err) { + items = append(items, Analysis{Target: "vm:" + vmID, VMID: vmID, Decision: DecisionNotFound, Reason: "VM not found"}) + continue + } + items = append(items, Analysis{Target: "vm:" + vmID, VMID: vmID, Decision: DecisionUnavailable, Reason: err.Error()}) + continue + } + item := a.analyzeVM(ctx, vm) + item.Target = "vm:" + vmID + items = append(items, item) + } + return items, nil +} + +func (a *analyzer) analyzeBatch(ctx context.Context, vms []*db.VirtualMachine) []Analysis { + workers := analyzeWorkerCount + if len(vms) < workers { + workers = len(vms) + } + jobs := make(chan *db.VirtualMachine) + results := make(chan Analysis, len(vms)) + var wg sync.WaitGroup + for range workers { + wg.Add(1) + go func() { + defer wg.Done() + for vm := range jobs { + results <- a.analyzeVM(ctx, vm) + } + }() + } + go func() { + for _, vm := range vms { + jobs <- vm + } + close(jobs) + wg.Wait() + close(results) + }() + items := make([]Analysis, 0, len(vms)) + for item := range results { + items = append(items, item) + } + return items +} + +func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysis { + item := Analysis{VMID: vm.ID, AlreadyRecycled: vm.IsRecycled} + for _, tk := range vm.Edges.Tasks { + if tk == nil { + continue + } + store := consts.LogStoreLoki + if tk.LogStore != nil && *tk.LogStore != "" { + store = *tk.LogStore + } + item.Tasks = append(item.Tasks, TaskInfo{ID: tk.ID, Status: tk.Status, LogStore: store, CreatedAt: tk.CreatedAt}) + } + if vm.IsRecycled { + item.Decision = DecisionAlreadyRecycled + item.Reason = "VM already recycled" + return item + } + if len(item.Tasks) == 0 { + item.Decision = DecisionUnavailable + item.Reason = "VM has no task" + return item + } + + policy, err := a.resolvePolicy(ctx, vm.UserID) + if err != nil { + item.Decision = DecisionUnavailable + item.Reason = fmt.Sprintf("resolve recycle policy: %v", err) + return item + } + if !policy.RecycleEnabled { + item.Decision = DecisionDisabled + item.Reason = "recycle disabled" + return item + } + item.RecycleSeconds = policy.EffectiveRecycleSeconds + now := a.now() + for _, tk := range item.Tasks { + lastActivity := tk.CreatedAt + latest, ok, err := a.activity.LatestEventTime(ctx, tk.ID, tk.CreatedAt, now, activityEvents, tk.LogStore) + if err != nil { + item.Decision = DecisionUnavailable + item.Reason = fmt.Sprintf("query task %s activity: %v", tk.ID, err) + return item + } + if ok && latest.After(lastActivity) { + lastActivity = latest + } + if lastActivity.After(item.LastActivityAt) { + item.LastActivityAt = lastActivity + } + } + item.DueAt = item.LastActivityAt.Add(time.Duration(item.RecycleSeconds) * time.Second) + redisAt, ok, err := a.deadlines.GetRunAt(ctx, RecycleQueueKey, vm.ID) + if err != nil { + item.Decision = DecisionUnavailable + item.Reason = fmt.Sprintf("query redis recycle deadline: %v", err) + return item + } + if ok { + item.RedisRecycleAt = &redisAt + } + if now.Before(item.DueAt) { + item.Decision = DecisionNotDue + item.Reason = "activity deadline not reached" + return item + } + item.Overdue = now.Sub(item.DueAt) + if item.RedisRecycleAt == nil { + item.Decision = DecisionHistoricalUncertain + item.Reason = "redis recycle deadline missing" + return item + } + if now.Before(*item.RedisRecycleAt) { + item.Decision = DecisionHistoricalUncertain + item.Reason = "redis recycle deadline is still in the future" + return item + } + item.Decision = DecisionCandidate + item.Reason = "activity and redis deadlines are overdue" + return item +} + +func (a *analyzer) resolvePolicy(ctx context.Context, userID uuid.UUID) (*domain.TeamTaskVMIdlePolicy, error) { + team, err := a.teamPolicyRepo.GetTeamByUserID(ctx, userID) + if err != nil && !db.IsNotFound(err) { + return nil, err + } + if db.IsNotFound(err) { + team = nil + } + return domain.ResolveTeamTaskVMIdlePolicy(team, a.cfg.VMIdle) +} + +func (a *analyzer) loadVM(ctx context.Context, vmID string) (*db.VirtualMachine, error) { + return a.db.VirtualMachine.Query().Where(virtualmachine.ID(vmID)).WithTasks().Only(ctx) +} + +var _ Analyzer = (*analyzer)(nil) +var _ activityReader = (*tasklog.Gateway)(nil) +var _ deadlineReader = (*delayqueue.VMRecycleQueue)(nil) + +func IsExecutable(decision Decision) bool { + return decision == DecisionCandidate +} + +func IsSuccessfulNoop(decision Decision) bool { + return decision == DecisionAlreadyRecycled +} diff --git a/backend/pkg/vmrecycle/analyzer_test.go b/backend/pkg/vmrecycle/analyzer_test.go new file mode 100644 index 000000000..e0e0ed0ba --- /dev/null +++ b/backend/pkg/vmrecycle/analyzer_test.go @@ -0,0 +1,231 @@ +package vmrecycle + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + _ "github.com/mattn/go-sqlite3" + + "github.com/chaitin/MonkeyCode/backend/config" + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/db/enttest" + "github.com/chaitin/MonkeyCode/backend/domain" +) + +func TestAnalyzerCandidateRequiresBothDeadlinesOverdue(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + taskID := uuid.New() + logStore := consts.LogStoreClickHouse + vm := &db.VirtualMachine{ + ID: "vm-candidate", + UserID: uuid.New(), + Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ + ID: taskID, + Status: consts.TaskStatusProcessing, + CreatedAt: now.Add(-2 * time.Hour), + LogStore: &logStore, + }}}, + } + activityAt := now.Add(-90 * time.Minute) + redisAt := now.Add(-time.Minute) + analyzer := newAnalyzerForTest(now, &analyzerActivityStub{latest: map[uuid.UUID]time.Time{taskID: activityAt}}, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: redisAt}}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionCandidate { + t.Fatalf("decision = %s, reason = %s", got.Decision, got.Reason) + } + if !got.LastActivityAt.Equal(activityAt) || !got.DueAt.Equal(activityAt.Add(time.Hour)) { + t.Fatalf("last activity = %v, due at = %v", got.LastActivityAt, got.DueAt) + } + if got.RedisRecycleAt == nil || !got.RedisRecycleAt.Equal(redisAt) { + t.Fatalf("redis recycle at = %v", got.RedisRecycleAt) + } + if got.Overdue != 30*time.Minute { + t.Fatalf("overdue = %v, want 30m", got.Overdue) + } +} + +func TestAnalyzerMarksHistoricalUncertainWhenRedisDeadlineCannotConfirm(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + taskID := uuid.New() + vm := &db.VirtualMachine{ID: "vm-uncertain", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ID: taskID, CreatedAt: now.Add(-2 * time.Hour)}}}} + + tests := []struct { + name string + deadlines *analyzerDeadlineStub + }{ + {name: "missing", deadlines: &analyzerDeadlineStub{}}, + {name: "future", deadlines: &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(time.Hour)}}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + analyzer := newAnalyzerForTest(now, &analyzerActivityStub{}, tt.deadlines) + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionHistoricalUncertain { + t.Fatalf("decision = %s, reason = %s", got.Decision, got.Reason) + } + }) + } +} + +func TestAnalyzerUsesLatestActivityAcrossTasks(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + oldTaskID := uuid.New() + recentTaskID := uuid.New() + vm := &db.VirtualMachine{ID: "vm-multi-task", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{ + {ID: oldTaskID, CreatedAt: now.Add(-3 * time.Hour)}, + {ID: recentTaskID, CreatedAt: now.Add(-2 * time.Hour)}, + }}} + activity := &analyzerActivityStub{latest: map[uuid.UUID]time.Time{ + oldTaskID: now.Add(-2 * time.Hour), + recentTaskID: now.Add(-30 * time.Minute), + }} + analyzer := newAnalyzerForTest(now, activity, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(-time.Hour)}}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionNotDue { + t.Fatalf("decision = %s, want %s", got.Decision, DecisionNotDue) + } + if !got.LastActivityAt.Equal(now.Add(-30 * time.Minute)) { + t.Fatalf("last activity = %v", got.LastActivityAt) + } + if len(activity.events) != len(activityEvents) { + t.Fatalf("activity events = %v", activity.events) + } + if _, ok := activity.events["task-event"]; !ok { + t.Fatalf("activity events = %v, missing task-event", activity.events) + } +} + +func TestAnalyzerReturnsUnavailableOnLogFailure(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + wantErr := errors.New("log query failed") + taskID := uuid.New() + vm := &db.VirtualMachine{ID: "vm-log-error", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ID: taskID, CreatedAt: now.Add(-2 * time.Hour)}}}} + analyzer := newAnalyzerForTest(now, &analyzerActivityStub{err: wantErr}, &analyzerDeadlineStub{}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionUnavailable || got.Reason == "" { + t.Fatalf("analysis = %+v", got) + } +} + +func TestAnalyzerAlreadyRecycledIsSuccessfulNoop(t *testing.T) { + analyzer := newAnalyzerForTest(time.Now(), &analyzerActivityStub{}, &analyzerDeadlineStub{}) + got := analyzer.analyzeVM(context.Background(), &db.VirtualMachine{ID: "vm-recycled", IsRecycled: true}) + if got.Decision != DecisionAlreadyRecycled || !IsSuccessfulNoop(got.Decision) { + t.Fatalf("analysis = %+v", got) + } +} + +func TestAnalyzerDeduplicatesTaskAndVMTargets(t *testing.T) { + ctx := context.Background() + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + client := enttest.Open(t, "sqlite3", fmt.Sprintf("file:vm-recycle-analyzer-%s?mode=memory&cache=shared&_fk=1", uuid.NewString())) + t.Cleanup(func() { _ = client.Close() }) + + userID := uuid.New() + if _, err := client.User.Create(). + SetID(userID). + SetName("tester"). + SetRole(consts.UserRoleIndividual). + SetStatus(consts.UserStatusActive). + Save(ctx); err != nil { + t.Fatal(err) + } + hostID := "host-1" + if _, err := client.Host.Create().SetID(hostID).SetUserID(userID).SetHostname("host").Save(ctx); err != nil { + t.Fatal(err) + } + vmID := "vm-1" + if _, err := client.VirtualMachine.Create().SetID(vmID).SetHostID(hostID).SetUserID(userID).SetName("vm").Save(ctx); err != nil { + t.Fatal(err) + } + taskID := uuid.New() + if _, err := client.Task.Create(). + SetID(taskID). + SetUserID(userID). + SetKind(consts.TaskTypeDevelop). + SetContent("content"). + SetStatus(consts.TaskStatusProcessing). + SetCreatedAt(now.Add(-2 * time.Hour)). + Save(ctx); err != nil { + t.Fatal(err) + } + if _, err := client.TaskVirtualMachine.Create().SetID(uuid.New()).SetTaskID(taskID).SetVirtualmachineID(vmID).Save(ctx); err != nil { + t.Fatal(err) + } + + analyzer := newAnalyzerForTest(now, &analyzerActivityStub{}, &analyzerDeadlineStub{runAt: map[string]time.Time{vmID: now.Add(-time.Minute)}}) + analyzer.db = client + items, err := analyzer.AnalyzeTargets(ctx, []uuid.UUID{taskID, taskID}, []string{vmID, vmID}) + if err != nil { + t.Fatal(err) + } + if len(items) != 1 || items[0].VMID != vmID { + t.Fatalf("items = %+v, want one analysis for %s", items, vmID) + } +} + +func newAnalyzerForTest(now time.Time, activity activityReader, deadlines deadlineReader) *analyzer { + return &analyzer{ + cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}}, + teamPolicyRepo: &analyzerTeamPolicyRepoStub{team: &db.Team{ + ID: uuid.New(), + TaskConcurrencyLimit: 3, + TaskVMSleepEnabled: true, + TaskVMSleepSeconds: 600, + TaskVMRecycleEnabled: true, + TaskVMRecycleSeconds: 0, + }}, + activity: activity, + deadlines: deadlines, + now: func() time.Time { return now }, + } +} + +type analyzerActivityStub struct { + latest map[uuid.UUID]time.Time + err error + events map[string]struct{} +} + +func (s *analyzerActivityStub) LatestEventTime(_ context.Context, taskID uuid.UUID, _, _ time.Time, events []string, _ consts.LogStore) (time.Time, bool, error) { + s.events = make(map[string]struct{}, len(events)) + for _, event := range events { + s.events[event] = struct{}{} + } + if s.err != nil { + return time.Time{}, false, s.err + } + latest, ok := s.latest[taskID] + return latest, ok, nil +} + +type analyzerDeadlineStub struct { + runAt map[string]time.Time + err error +} + +func (s *analyzerDeadlineStub) GetRunAt(_ context.Context, _, id string) (time.Time, bool, error) { + if s.err != nil { + return time.Time{}, false, s.err + } + runAt, ok := s.runAt[id] + return runAt, ok, nil +} + +type analyzerTeamPolicyRepoStub struct { + domain.TeamPolicyRepo + team *db.Team + err error +} + +func (s *analyzerTeamPolicyRepoStub) GetTeamByUserID(context.Context, uuid.UUID) (*db.Team, error) { + return s.team, s.err +} diff --git a/backend/pkg/vmrecycle/vmrecycle.go b/backend/pkg/vmrecycle/vmrecycle.go new file mode 100644 index 000000000..461548a3d --- /dev/null +++ b/backend/pkg/vmrecycle/vmrecycle.go @@ -0,0 +1,232 @@ +package vmrecycle + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + "github.com/samber/do" + + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/domain" + "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" + "github.com/chaitin/MonkeyCode/backend/pkg/entx" + "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" +) + +const ( + SleepQueueKey = "vm:idle:sleep" + NotifyQueueKey = "vm:idle:notify" + RecycleQueueKey = "vm:idle:recycle" + VMExpireQueueKey = "vm:expire" + wechat2hQueueKey = "vm:idle:notify:wechat:2h" + wechat15mQueueKey = "vm:idle:notify:wechat:15m" + recycleLockTTL = 10 * time.Minute +) + +type Status string + +const ( + StatusRecycled Status = "recycled" + StatusAlreadyRecycled Status = "already_recycled" + StatusNotFound Status = "not_found" +) + +var ErrInProgress = errors.New("vm recycle in progress") + +type Result struct { + VMID string + TaskIDs []uuid.UUID + Status Status +} + +type Recycler interface { + Recycle(ctx context.Context, vmID string) (Result, error) +} + +type removableQueue interface { + Remove(ctx context.Context, queue, id string) error +} + +type notifyQueue interface { + removableQueue + RemoveByPrefix(ctx context.Context, queue, prefix string) (int, error) +} + +type recycler struct { + redis *redis.Client + logger *slog.Logger + hostRepo domain.HostRepo + taskRepo domain.TaskRepo + taskflow taskflow.Clienter + sleepQueue removableQueue + notifyQueue notifyQueue + recycleQueue removableQueue + expireQueue removableQueue + now func() time.Time +} + +func NewRecycler(i *do.Injector) (Recycler, error) { + return &recycler{ + redis: do.MustInvoke[*redis.Client](i), + logger: do.MustInvoke[*slog.Logger](i).With("module", "VMRecycler"), + hostRepo: do.MustInvoke[domain.HostRepo](i), + taskRepo: do.MustInvoke[domain.TaskRepo](i), + taskflow: do.MustInvoke[taskflow.Clienter](i), + sleepQueue: do.MustInvoke[*delayqueue.VMSleepQueue](i), + notifyQueue: do.MustInvoke[*delayqueue.VMNotifyQueue](i), + recycleQueue: do.MustInvoke[*delayqueue.VMRecycleQueue](i), + expireQueue: do.MustInvoke[*delayqueue.VMExpireQueue](i), + now: time.Now, + }, nil +} + +func (r *recycler) Recycle(ctx context.Context, vmID string) (Result, error) { + result := Result{VMID: vmID} + token, err := r.acquire(ctx, vmID) + if err != nil { + return result, err + } + defer r.release(ctx, vmID, token) + + ctx = entx.SkipSoftDelete(ctx) + vm, err := r.hostRepo.GetVirtualMachine(ctx, vmID) + if err != nil { + if db.IsNotFound(err) { + result.Status = StatusNotFound + return result, nil + } + return result, fmt.Errorf("get vm %s: %w", vmID, err) + } + result.TaskIDs = taskIDs(vm) + + if vm.IsRecycled { + result.Status = StatusAlreadyRecycled + } else { + if err := r.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{ + UserID: vm.UserID.String(), + HostID: vm.HostID, + ID: vm.EnvironmentID, + }); err != nil { + return result, fmt.Errorf("delete vm %s: %w", vmID, err) + } + if err := r.hostRepo.UpdateVirtualMachine(ctx, vmID, func(up *db.VirtualMachineUpdateOne) error { + up.SetIsRecycled(true) + return nil + }); err != nil { + return result, fmt.Errorf("mark vm %s recycled: %w", vmID, err) + } + result.Status = StatusRecycled + } + + if err := r.cleanup(ctx, vm); err != nil { + return result, err + } + r.logger.InfoContext(ctx, "vm recycled", "vm_id", vmID, "status", result.Status, "task_ids", result.TaskIDs) + return result, nil +} + +func (r *recycler) acquire(ctx context.Context, vmID string) (string, error) { + token := uuid.NewString() + err := r.redis.SetArgs(ctx, recycleLockKey(vmID), token, redis.SetArgs{Mode: "NX", TTL: recycleLockTTL}).Err() + if errors.Is(err, redis.Nil) { + return "", fmt.Errorf("%w: %s", ErrInProgress, vmID) + } + if err != nil { + return "", fmt.Errorf("acquire recycle lock for vm %s: %w", vmID, err) + } + return token, nil +} + +func (r *recycler) release(ctx context.Context, vmID, token string) { + const script = ` +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +end +return 0 +` + if err := r.redis.Eval(ctx, script, []string{recycleLockKey(vmID)}, token).Err(); err != nil && !errors.Is(err, redis.Nil) { + r.logger.WarnContext(ctx, "failed to release vm recycle lock", "vm_id", vmID, "error", err) + } +} + +func (r *recycler) cleanup(ctx context.Context, vm *db.VirtualMachine) error { + var errs []error + completedAt := r.now() + for _, tk := range vm.Edges.Tasks { + if tk == nil || tk.Status == consts.TaskStatusFinished || tk.Status == consts.TaskStatusError { + continue + } + if err := r.taskRepo.Update(ctx, nil, tk.ID, func(up *db.TaskUpdateOne) error { + up.SetStatus(consts.TaskStatusFinished) + up.SetCompletedAt(completedAt) + return nil + }); err != nil { + errs = append(errs, fmt.Errorf("finish task %s: %w", tk.ID, err)) + } + } + + remove := func(name string, queue removableQueue, queueKey, id string) { + if err := queue.Remove(ctx, queueKey, id); err != nil { + errs = append(errs, fmt.Errorf("remove %s: %w", name, err)) + } + } + remove("sleep job", r.sleepQueue, SleepQueueKey, vm.ID) + if _, err := r.notifyQueue.RemoveByPrefix(ctx, NotifyQueueKey, vm.ID+":"); err != nil { + errs = append(errs, fmt.Errorf("remove notify jobs: %w", err)) + } + remove("legacy notify job", r.notifyQueue, NotifyQueueKey, vm.ID) + remove("legacy wechat 2h notify job", r.notifyQueue, wechat2hQueueKey, vm.ID) + remove("legacy wechat 15m notify job", r.notifyQueue, wechat15mQueueKey, vm.ID) + remove("expire job", r.expireQueue, VMExpireQueueKey, vm.ID) + + keys := []string{ + fmt.Sprintf("lifecycle:%s", vm.ID), + fmt.Sprintf("vm:idle:debounce:%s", vm.ID), + fmt.Sprintf("vm:idle:debounce:%s:keep-awake", vm.ID), + fmt.Sprintf("vm:idle:debounce:%s:activity", vm.ID), + fmt.Sprintf("vm:idle:not-found:%s", vm.ID), + } + for _, taskID := range taskIDs(vm) { + keys = append(keys, + fmt.Sprintf("task:create_req:%s", taskID), + fmt.Sprintf("mcai:task:%s:last_input", taskID), + ) + } + if err := r.redis.Del(ctx, keys...).Err(); err != nil { + errs = append(errs, fmt.Errorf("delete recycle redis keys: %w", err)) + } + + if len(errs) > 0 { + return errors.Join(errs...) + } + if err := r.recycleQueue.Remove(ctx, RecycleQueueKey, vm.ID); err != nil { + return fmt.Errorf("remove recycle job: %w", err) + } + return nil +} + +func taskIDs(vm *db.VirtualMachine) []uuid.UUID { + seen := make(map[uuid.UUID]struct{}, len(vm.Edges.Tasks)) + ids := make([]uuid.UUID, 0, len(vm.Edges.Tasks)) + for _, tk := range vm.Edges.Tasks { + if tk == nil { + continue + } + if _, ok := seen[tk.ID]; ok { + continue + } + seen[tk.ID] = struct{}{} + ids = append(ids, tk.ID) + } + return ids +} + +func recycleLockKey(vmID string) string { + return fmt.Sprintf("vm:recycle:lock:%s", vmID) +} diff --git a/backend/pkg/vmrecycle/vmrecycle_test.go b/backend/pkg/vmrecycle/vmrecycle_test.go new file mode 100644 index 000000000..c9cf475b2 --- /dev/null +++ b/backend/pkg/vmrecycle/vmrecycle_test.go @@ -0,0 +1,331 @@ +package vmrecycle + +import ( + "context" + "errors" + "io" + "log/slog" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/domain" + "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" + "github.com/chaitin/MonkeyCode/backend/pkg/taskflow" +) + +func TestRecyclerRecyclesVMAndCleansLocalState(t *testing.T) { + ctx := context.Background() + rdb := newTestRedis(t) + processingTaskID := uuid.New() + finishedTaskID := uuid.New() + vm := &db.VirtualMachine{ + ID: "vm-1", + UserID: uuid.New(), + HostID: "host-1", + EnvironmentID: "env-1", + Edges: db.VirtualMachineEdges{Tasks: []*db.Task{ + {ID: processingTaskID, Status: consts.TaskStatusProcessing}, + {ID: finishedTaskID, Status: consts.TaskStatusFinished}, + }}, + } + hostRepo := &recycleHostRepoStub{vm: vm} + taskRepo := &recycleTaskRepoStub{} + vmClient := &recycleVMStub{} + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + sleepQueue := delayqueue.NewVMSleepQueue(rdb, logger) + notifyQueue := delayqueue.NewVMNotifyQueue(rdb, logger) + recycleQueue := delayqueue.NewVMRecycleQueue(rdb, logger) + expireQueue := delayqueue.NewVMExpireQueue(rdb, logger) + r := &recycler{ + redis: rdb, + logger: logger, + hostRepo: hostRepo, + taskRepo: taskRepo, + taskflow: &recycleTaskflowStub{vm: vmClient}, + sleepQueue: sleepQueue, + notifyQueue: notifyQueue, + recycleQueue: recycleQueue, + expireQueue: expireQueue, + now: func() time.Time { return time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) }, + } + payload := &domain.VmIdleInfo{VmID: vm.ID} + runAt := time.Now().Add(time.Hour) + mustEnqueue(t, sleepQueue, ctx, SleepQueueKey, payload, runAt, vm.ID) + mustEnqueue(t, notifyQueue, ctx, NotifyQueueKey, payload, runAt, vm.ID+":default") + mustEnqueue(t, notifyQueue, ctx, NotifyQueueKey, payload, runAt, vm.ID+":wechat900s") + mustEnqueue(t, notifyQueue, ctx, wechat2hQueueKey, payload, runAt, vm.ID) + mustEnqueue(t, notifyQueue, ctx, wechat15mQueueKey, payload, runAt, vm.ID) + mustEnqueue(t, recycleQueue, ctx, RecycleQueueKey, payload, runAt, vm.ID) + if _, err := expireQueue.Enqueue(ctx, VMExpireQueueKey, &domain.VmExpireInfo{VmID: vm.ID}, runAt, vm.ID); err != nil { + t.Fatal(err) + } + keys := []string{ + "lifecycle:" + vm.ID, + "vm:idle:debounce:" + vm.ID, + "vm:idle:debounce:" + vm.ID + ":keep-awake", + "vm:idle:debounce:" + vm.ID + ":activity", + "vm:idle:not-found:" + vm.ID, + "task:create_req:" + processingTaskID.String(), + "mcai:task:" + processingTaskID.String() + ":last_input", + } + for _, key := range keys { + if err := rdb.Set(ctx, key, "value", time.Hour).Err(); err != nil { + t.Fatal(err) + } + } + + result, err := r.Recycle(ctx, vm.ID) + if err != nil { + t.Fatal(err) + } + if result.Status != StatusRecycled || result.VMID != vm.ID || len(result.TaskIDs) != 2 { + t.Fatalf("result = %+v", result) + } + if vmClient.deleteCalls != 1 || vmClient.lastDelete.ID != vm.EnvironmentID { + t.Fatalf("delete calls = %d, request = %+v", vmClient.deleteCalls, vmClient.lastDelete) + } + if hostRepo.updateCalls != 1 || !vm.IsRecycled { + t.Fatalf("update calls = %d, recycled = %v", hostRepo.updateCalls, vm.IsRecycled) + } + if len(taskRepo.updated) != 1 || taskRepo.updated[0] != processingTaskID { + t.Fatalf("updated tasks = %v, want [%s]", taskRepo.updated, processingTaskID) + } + assertJobMissing(t, sleepQueue.RedisDelayQueue, ctx, SleepQueueKey, vm.ID) + assertJobMissing(t, notifyQueue.RedisDelayQueue, ctx, NotifyQueueKey, vm.ID+":default") + assertJobMissing(t, notifyQueue.RedisDelayQueue, ctx, NotifyQueueKey, vm.ID+":wechat900s") + assertJobMissing(t, notifyQueue.RedisDelayQueue, ctx, wechat2hQueueKey, vm.ID) + assertJobMissing(t, notifyQueue.RedisDelayQueue, ctx, wechat15mQueueKey, vm.ID) + assertJobMissing(t, recycleQueue.RedisDelayQueue, ctx, RecycleQueueKey, vm.ID) + if _, _, ok, err := expireQueue.GetJobInfo(ctx, VMExpireQueueKey, vm.ID); err != nil || ok { + t.Fatalf("expire job ok = %v, err = %v", ok, err) + } + for _, key := range keys { + if exists, err := rdb.Exists(ctx, key).Result(); err != nil || exists != 0 { + t.Fatalf("key %q exists = %d, err = %v", key, exists, err) + } + } +} + +func TestRecyclerAlreadyRecycledSkipsRemoteDeleteAndRepairsCleanup(t *testing.T) { + vm := &db.VirtualMachine{ID: "vm-recycled", IsRecycled: true, Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ID: uuid.New(), Status: consts.TaskStatusProcessing}}}} + r, _, taskRepo, vmClient := newStubRecycler(t, vm) + + result, err := r.Recycle(context.Background(), vm.ID) + if err != nil { + t.Fatal(err) + } + if result.Status != StatusAlreadyRecycled { + t.Fatalf("status = %s, want %s", result.Status, StatusAlreadyRecycled) + } + if vmClient.deleteCalls != 0 { + t.Fatalf("delete calls = %d, want 0", vmClient.deleteCalls) + } + if len(taskRepo.updated) != 1 { + t.Fatalf("updated tasks = %v", taskRepo.updated) + } +} + +func TestRecyclerRemoteFailureDoesNotMarkOrClean(t *testing.T) { + wantErr := errors.New("remote delete failed") + vm := &db.VirtualMachine{ID: "vm-delete-fail", UserID: uuid.New()} + r, hostRepo, _, vmClient := newStubRecycler(t, vm) + vmClient.err = wantErr + queues := r.sleepQueue.(*recycleQueueStub) + + result, err := r.Recycle(context.Background(), vm.ID) + if !errors.Is(err, wantErr) { + t.Fatalf("error = %v, want %v", err, wantErr) + } + if result.Status != "" || hostRepo.updateCalls != 0 || len(queues.removed) != 0 { + t.Fatalf("result = %+v, update calls = %d, cleanup = %v", result, hostRepo.updateCalls, queues.removed) + } +} + +func TestRecyclerRetriesCleanupWithoutDeletingRemoteAgain(t *testing.T) { + vm := &db.VirtualMachine{ID: "vm-cleanup-retry", UserID: uuid.New()} + r, hostRepo, _, vmClient := newStubRecycler(t, vm) + sleepQueue := r.sleepQueue.(*recycleQueueStub) + sleepQueue.err = errors.New("redis unavailable") + + result, err := r.Recycle(context.Background(), vm.ID) + if err == nil || result.Status != StatusRecycled { + t.Fatalf("first result = %+v, error = %v", result, err) + } + if vmClient.deleteCalls != 1 || hostRepo.updateCalls != 1 || !vm.IsRecycled { + t.Fatalf("delete calls = %d, update calls = %d, recycled = %v", vmClient.deleteCalls, hostRepo.updateCalls, vm.IsRecycled) + } + if r.recycleQueue.(*recycleQueueStub).removedContains(RecycleQueueKey, vm.ID) { + t.Fatal("recycle job must remain when cleanup fails") + } + + sleepQueue.err = nil + result, err = r.Recycle(context.Background(), vm.ID) + if err != nil || result.Status != StatusAlreadyRecycled { + t.Fatalf("retry result = %+v, error = %v", result, err) + } + if vmClient.deleteCalls != 1 { + t.Fatalf("delete calls = %d, want 1", vmClient.deleteCalls) + } + if !r.recycleQueue.(*recycleQueueStub).removedContains(RecycleQueueKey, vm.ID) { + t.Fatal("recycle job should be removed after cleanup succeeds") + } +} + +func TestRecyclerReturnsInProgressWhenLockExists(t *testing.T) { + vm := &db.VirtualMachine{ID: "vm-locked"} + r, _, _, vmClient := newStubRecycler(t, vm) + if err := r.redis.Set(context.Background(), recycleLockKey(vm.ID), "other", time.Minute).Err(); err != nil { + t.Fatal(err) + } + + _, err := r.Recycle(context.Background(), vm.ID) + if !errors.Is(err, ErrInProgress) { + t.Fatalf("error = %v, want ErrInProgress", err) + } + if vmClient.deleteCalls != 0 { + t.Fatalf("delete calls = %d, want 0", vmClient.deleteCalls) + } +} + +func TestRecyclerReturnsNotFoundResult(t *testing.T) { + r, _, _, _ := newStubRecycler(t, nil) + r.hostRepo.(*recycleHostRepoStub).err = &db.NotFoundError{} + + result, err := r.Recycle(context.Background(), "missing-vm") + if err != nil { + t.Fatal(err) + } + if result.Status != StatusNotFound { + t.Fatalf("status = %s, want %s", result.Status, StatusNotFound) + } +} + +func newStubRecycler(t *testing.T, vm *db.VirtualMachine) (*recycler, *recycleHostRepoStub, *recycleTaskRepoStub, *recycleVMStub) { + t.Helper() + rdb := newTestRedis(t) + hostRepo := &recycleHostRepoStub{vm: vm} + taskRepo := &recycleTaskRepoStub{} + vmClient := &recycleVMStub{} + return &recycler{ + redis: rdb, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + hostRepo: hostRepo, + taskRepo: taskRepo, + taskflow: &recycleTaskflowStub{vm: vmClient}, + sleepQueue: &recycleQueueStub{}, + notifyQueue: &recycleQueueStub{}, + recycleQueue: &recycleQueueStub{}, + expireQueue: &recycleQueueStub{}, + now: time.Now, + }, hostRepo, taskRepo, vmClient +} + +func newTestRedis(t *testing.T) *redis.Client { + t.Helper() + srv := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + return client +} + +func mustEnqueue(t *testing.T, queue interface { + Enqueue(context.Context, string, *domain.VmIdleInfo, time.Time, string) (string, error) +}, ctx context.Context, queueKey string, payload *domain.VmIdleInfo, runAt time.Time, id string) { + t.Helper() + if _, err := queue.Enqueue(ctx, queueKey, payload, runAt, id); err != nil { + t.Fatal(err) + } +} + +func assertJobMissing(t *testing.T, queue *delayqueue.RedisDelayQueue[*domain.VmIdleInfo], ctx context.Context, queueKey, id string) { + t.Helper() + if _, _, ok, err := queue.GetJobInfo(ctx, queueKey, id); err != nil || ok { + t.Fatalf("job %q in %q ok = %v, err = %v", id, queueKey, ok, err) + } +} + +type recycleHostRepoStub struct { + domain.HostRepo + vm *db.VirtualMachine + err error + updateErr error + updateCalls int +} + +func (s *recycleHostRepoStub) GetVirtualMachine(context.Context, string) (*db.VirtualMachine, error) { + return s.vm, s.err +} + +func (s *recycleHostRepoStub) UpdateVirtualMachine(context.Context, string, func(*db.VirtualMachineUpdateOne) error) error { + s.updateCalls++ + if s.updateErr == nil && s.vm != nil { + s.vm.IsRecycled = true + } + return s.updateErr +} + +type recycleTaskRepoStub struct { + domain.TaskRepo + updated []uuid.UUID + err error +} + +func (s *recycleTaskRepoStub) Update(_ context.Context, _ *domain.User, id uuid.UUID, _ func(*db.TaskUpdateOne) error) error { + s.updated = append(s.updated, id) + return s.err +} + +type recycleTaskflowStub struct { + taskflow.Clienter + vm taskflow.VirtualMachiner +} + +func (s *recycleTaskflowStub) VirtualMachiner() taskflow.VirtualMachiner { return s.vm } + +type recycleVMStub struct { + taskflow.VirtualMachiner + deleteCalls int + lastDelete *taskflow.DeleteVirtualMachineReq + err error +} + +func (s *recycleVMStub) Delete(_ context.Context, req *taskflow.DeleteVirtualMachineReq) error { + s.deleteCalls++ + s.lastDelete = req + return s.err +} + +type recycleQueueStub struct { + removed []queueRemoval + err error +} + +type queueRemoval struct { + queue string + id string +} + +func (s *recycleQueueStub) Remove(_ context.Context, queue, id string) error { + s.removed = append(s.removed, queueRemoval{queue: queue, id: id}) + return s.err +} + +func (s *recycleQueueStub) RemoveByPrefix(_ context.Context, queue, prefix string) (int, error) { + s.removed = append(s.removed, queueRemoval{queue: queue, id: prefix}) + return 0, s.err +} + +func (s *recycleQueueStub) removedContains(queue, id string) bool { + for _, item := range s.removed { + if item.queue == queue && item.id == id { + return true + } + } + return false +} From b92662404964be98bb56948dbe7aad6ffa83a34f Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 12:05:56 +0800 Subject: [PATCH 03/10] =?UTF-8?q?=E6=8C=81=E4=B9=85=E5=8C=96=E8=99=9A?= =?UTF-8?q?=E6=8B=9F=E6=9C=BA=E6=B4=BB=E5=8A=A8=E5=B9=B6=E7=A7=BB=E9=99=A4?= =?UTF-8?q?=E4=BA=8B=E4=BB=B6=E6=97=A5=E5=BF=97=E4=BE=9D=E8=B5=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/biz/host/handler/v1/internal.go | 37 +++++++- .../handler/v1/internal_vm_activity_test.go | 94 ++++++++++++++++++- backend/biz/task/handler/v1/task_control.go | 19 +--- .../biz/task/handler/v1/task_control_test.go | 36 +------ backend/biz/task/repo/task.go | 13 +++ backend/biz/task/repo/task_activity_test.go | 91 ++++++++++++++++++ backend/pkg/vmrecycle/analyzer.go | 46 +++------ backend/pkg/vmrecycle/analyzer_test.go | 64 ++++--------- 8 files changed, 266 insertions(+), 134 deletions(-) create mode 100644 backend/biz/task/repo/task_activity_test.go diff --git a/backend/biz/host/handler/v1/internal.go b/backend/biz/host/handler/v1/internal.go index f76449c67..115182148 100644 --- a/backend/biz/host/handler/v1/internal.go +++ b/backend/biz/host/handler/v1/internal.go @@ -34,6 +34,7 @@ type InternalHostHandler struct { logger *slog.Logger repo domain.HostRepo taskRepo taskLogStoreRepo + taskActivity vmActivityTaskRepo teamRepo domain.TeamHostRepo redis *redis.Client getAgentToken agentTokenGetter @@ -53,15 +54,27 @@ type taskLogStoreRepo interface { GetLogStore(ctx context.Context, id uuid.UUID) (consts.LogStore, error) } +type vmActivityTaskRepo interface { + RefreshLastActiveAtByVMID(ctx context.Context, vmID string, at time.Time, minInterval time.Duration) error +} + +const vmActivityTaskRefreshInterval = 5 * time.Minute + func NewInternalHostHandler(i *do.Injector) (*InternalHostHandler, error) { w := do.MustInvoke[*web.Web](i) tf := do.MustInvoke[taskflow.Clienter](i) rdb := do.MustInvoke[*redis.Client](i) + taskRepo := do.MustInvoke[domain.TaskRepo](i) + taskActivity, ok := taskRepo.(vmActivityTaskRepo) + if !ok { + return nil, errors.New("task repo does not support VM activity") + } h := &InternalHostHandler{ logger: do.MustInvoke[*slog.Logger](i).With("module", "InternalHostHandler"), repo: do.MustInvoke[domain.HostRepo](i), - taskRepo: do.MustInvoke[domain.TaskRepo](i), + taskRepo: taskRepo, + taskActivity: taskActivity, teamRepo: do.MustInvoke[domain.TeamHostRepo](i), redis: rdb, getAgentToken: defaultAgentTokenGetter(rdb), @@ -122,8 +135,26 @@ func (h *InternalHostHandler) VMActivity(c *web.Context, req VMActivityReq) erro if strings.TrimSpace(req.VMID) == "" { return errors.New("vm_id is required") } - if err := h.idleRefresher.RecordActivity(c.Request().Context(), req.VMID); err != nil { - h.logger.WarnContext(c.Request().Context(), "failed to refresh vm idle timers on activity", "vm_id", req.VMID, "error", err) + if req.LastActiveAt <= 0 { + return errors.New("last_active_at is required") + } + + ctx := c.Request().Context() + activeAt := time.Unix(req.LastActiveAt, 0) + if now := time.Now(); activeAt.After(now) { + h.logger.WarnContext(ctx, "vm activity timestamp is in the future", "vm_id", req.VMID, "last_active_at", req.LastActiveAt) + activeAt = now + } + var errs []error + if err := h.taskActivity.RefreshLastActiveAtByVMID(ctx, req.VMID, activeAt, vmActivityTaskRefreshInterval); err != nil { + h.logger.WarnContext(ctx, "failed to persist vm activity", "vm_id", req.VMID, "error", err) + errs = append(errs, err) + } + if err := h.idleRefresher.RecordActivity(ctx, req.VMID); err != nil { + h.logger.WarnContext(ctx, "failed to refresh vm idle timers on activity", "vm_id", req.VMID, "error", err) + errs = append(errs, err) + } + if err := errors.Join(errs...); err != nil { return err } return c.Success(nil) diff --git a/backend/biz/host/handler/v1/internal_vm_activity_test.go b/backend/biz/host/handler/v1/internal_vm_activity_test.go index 469e988b0..b2df70056 100644 --- a/backend/biz/host/handler/v1/internal_vm_activity_test.go +++ b/backend/biz/host/handler/v1/internal_vm_activity_test.go @@ -3,21 +3,25 @@ package v1 import ( "context" "encoding/json" + "errors" "io" "log/slog" "net/http" "net/http/httptest" "strings" "testing" + "time" "github.com/GoYoko/web" ) -func TestInternalHostHandler_VMActivityRefreshesIdleTimer(t *testing.T) { +func TestInternalHostHandler_VMActivityPersistsAndRefreshesIdleTimer(t *testing.T) { refresher := &internalVMIdleRefresherStub{ch: make(chan string, 1)} + activity := &internalVMActivityTaskRepoStub{ch: make(chan internalVMActivityCall, 1)} h := &InternalHostHandler{ logger: slog.New(slog.NewTextHandler(io.Discard, nil)), idleRefresher: refresher, + taskActivity: activity, } body := `{"vm_id":"vm-activity-1","last_active_at":1710000000}` @@ -40,12 +44,57 @@ func TestInternalHostHandler_VMActivityRefreshesIdleTimer(t *testing.T) { default: t.Fatal("expected idle refresher to be called") } + select { + case got := <-activity.ch: + if got.vmID != "vm-activity-1" || !got.at.Equal(time.Unix(1710000000, 0)) || got.minInterval != vmActivityTaskRefreshInterval { + t.Fatalf("activity call = %+v", got) + } + default: + t.Fatal("expected task activity to be persisted") + } +} + +func TestInternalHostHandler_VMActivityRefreshesIdleTimerWhenPersistFails(t *testing.T) { + refresher := &internalVMIdleRefresherStub{ch: make(chan string, 1)} + h := &InternalHostHandler{ + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + idleRefresher: refresher, + taskActivity: &internalVMActivityTaskRepoStub{ + ch: make(chan internalVMActivityCall, 1), + err: errors.New("persist failed"), + }, + } + + w := web.New() + w.POST("/internal/vm/activity", web.BindHandler(h.VMActivity)) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/internal/vm/activity", strings.NewReader(`{"vm_id":"vm-activity-1","last_active_at":1710000000}`)) + req.Header.Set("Content-Type", "application/json") + w.Echo().ServeHTTP(rec, req) + + select { + case got := <-refresher.ch: + if got != "vm-activity-1" { + t.Fatalf("refreshed vm id = %q", got) + } + default: + t.Fatal("expected idle refresher to run after persist failure") + } + + var resp web.Resp + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("unmarshal resp: %v, body = %s", err, rec.Body.String()) + } + if resp.Code == 0 { + t.Fatalf("response = %+v, want error", resp) + } } func TestInternalHostHandler_VMActivityRejectsEmptyVMID(t *testing.T) { h := &InternalHostHandler{ logger: slog.New(slog.NewTextHandler(io.Discard, nil)), idleRefresher: &internalVMIdleRefresherStub{ch: make(chan string, 1)}, + taskActivity: &internalVMActivityTaskRepoStub{ch: make(chan internalVMActivityCall, 1)}, } w := web.New() @@ -65,6 +114,30 @@ func TestInternalHostHandler_VMActivityRejectsEmptyVMID(t *testing.T) { } } +func TestInternalHostHandler_VMActivityRejectsMissingTimestamp(t *testing.T) { + h := &InternalHostHandler{ + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + idleRefresher: &internalVMIdleRefresherStub{ch: make(chan string, 1)}, + taskActivity: &internalVMActivityTaskRepoStub{ch: make(chan internalVMActivityCall, 1)}, + } + + w := web.New() + w.POST("/internal/vm/activity", web.BindHandler(h.VMActivity)) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/internal/vm/activity", strings.NewReader(`{"vm_id":"vm-activity-1"}`)) + req.Header.Set("Content-Type", "application/json") + w.Echo().ServeHTTP(rec, req) + + var resp web.Resp + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("unmarshal resp: %v, body = %s", err, rec.Body.String()) + } + if resp.Code == 0 { + t.Fatalf("response = %+v, want error", resp) + } +} + type internalVMIdleRefresherStub struct { ch chan string } @@ -78,3 +151,22 @@ func (s *internalVMIdleRefresherStub) RecordActivity(_ context.Context, vmID str } return nil } + +type internalVMActivityCall struct { + vmID string + at time.Time + minInterval time.Duration +} + +type internalVMActivityTaskRepoStub struct { + ch chan internalVMActivityCall + err error +} + +func (s *internalVMActivityTaskRepoStub) RefreshLastActiveAtByVMID(_ context.Context, vmID string, at time.Time, minInterval time.Duration) error { + select { + case s.ch <- internalVMActivityCall{vmID: vmID, at: at, minInterval: minInterval}: + default: + } + return s.err +} diff --git a/backend/biz/task/handler/v1/task_control.go b/backend/biz/task/handler/v1/task_control.go index 19cce13ab..866ef0a87 100644 --- a/backend/biz/task/handler/v1/task_control.go +++ b/backend/biz/task/handler/v1/task_control.go @@ -10,9 +10,7 @@ import ( "golang.org/x/sync/errgroup" "github.com/GoYoko/web" - "github.com/google/uuid" - "github.com/chaitin/MonkeyCode/backend/biz/task/service" "github.com/chaitin/MonkeyCode/backend/consts" "github.com/chaitin/MonkeyCode/backend/domain" "github.com/chaitin/MonkeyCode/backend/middleware" @@ -142,10 +140,6 @@ func (h *TaskHandler) Control(c *web.Context, req domain.TaskControlReq) error { logger := h.logger.With("task_id", task.ID, "fn", "task.control") taskID := task.ID.String() - if err := h.taskActivity.Refresh(c.Request().Context(), task.ID); err != nil { - logger.WarnContext(c.Request().Context(), "failed to refresh task last active on control connect", "error", err) - } - // 连接建立：刷新空闲计时器 if vm := task.VirtualMachine; vm != nil { if err := h.idleRefresher.KeepAwake(c.Request().Context(), vm.ID); err != nil { @@ -197,7 +191,7 @@ func (h *TaskHandler) Control(c *web.Context, req domain.TaskControlReq) error { // 定期刷新空闲计时器，保持 VM 活跃 if vm := task.VirtualMachine; vm != nil { g.Go(func() error { - return h.controlKeepAlive(ctx, task.ID, vm.ID) + return h.controlKeepAlive(ctx, vm.ID) }) } @@ -226,18 +220,13 @@ func (h *TaskHandler) controlPing(ctx context.Context, wsConn *ws.WebsocketManag } // controlKeepAlive 定期刷新空闲计时器，防止 VM 被误判空闲 -func (h *TaskHandler) controlKeepAlive(ctx context.Context, taskID uuid.UUID, vmID string) error { +func (h *TaskHandler) controlKeepAlive(ctx context.Context, vmID string) error { if err := h.idleRefresher.KeepAwake(ctx, vmID); err != nil { h.logger.WarnContext(ctx, "keepalive refresh failed", "vmID", vmID, "error", err) } - if err := h.taskActivity.Refresh(ctx, taskID); err != nil { - h.logger.WarnContext(ctx, "task activity refresh failed", "taskID", taskID, "error", err) - } idleTicker := time.NewTicker(1 * time.Minute) - activityTicker := time.NewTicker(service.TaskActivityRefreshInterval) defer idleTicker.Stop() - defer activityTicker.Stop() for { select { case <-ctx.Done(): @@ -246,10 +235,6 @@ func (h *TaskHandler) controlKeepAlive(ctx context.Context, taskID uuid.UUID, vm if err := h.idleRefresher.KeepAwake(ctx, vmID); err != nil { h.logger.WarnContext(ctx, "keepalive refresh failed", "vmID", vmID, "error", err) } - case <-activityTicker.C: - if err := h.taskActivity.Refresh(ctx, taskID); err != nil { - h.logger.WarnContext(ctx, "task activity refresh failed", "taskID", taskID, "error", err) - } } } } diff --git a/backend/biz/task/handler/v1/task_control_test.go b/backend/biz/task/handler/v1/task_control_test.go index 67f7720aa..ded3761b6 100644 --- a/backend/biz/task/handler/v1/task_control_test.go +++ b/backend/biz/task/handler/v1/task_control_test.go @@ -6,22 +6,15 @@ import ( "log/slog" "testing" "time" - - "github.com/google/uuid" - - "github.com/chaitin/MonkeyCode/backend/biz/task/service" ) func TestControlKeepAliveRefreshesImmediately(t *testing.T) { - taskID := uuid.MustParse("11111111-1111-1111-1111-111111111111") vmID := "vm-1" idleRefresher := &testVMIdleRefresher{ch: make(chan string, 1)} - taskActivity := &testTaskActivityRefresher{ch: make(chan uuid.UUID, 1)} handler := &TaskHandler{ logger: slog.New(slog.NewTextHandler(io.Discard, nil)), idleRefresher: idleRefresher, - taskActivity: taskActivity, } ctx, cancel := context.WithCancel(context.Background()) @@ -29,7 +22,7 @@ func TestControlKeepAliveRefreshesImmediately(t *testing.T) { done := make(chan error, 1) go func() { - done <- handler.controlKeepAlive(ctx, taskID, vmID) + done <- handler.controlKeepAlive(ctx, vmID) }() select { @@ -41,15 +34,6 @@ func TestControlKeepAliveRefreshesImmediately(t *testing.T) { t.Fatal("expected idle refresher to run immediately") } - select { - case got := <-taskActivity.ch: - if got != taskID { - t.Fatalf("task activity task id = %s, want %s", got, taskID) - } - case <-time.After(100 * time.Millisecond): - t.Fatal("expected task activity refresher to run immediately") - } - cancel() select { @@ -75,21 +59,3 @@ func (r *testVMIdleRefresher) KeepAwake(_ context.Context, vmID string) error { } func (r *testVMIdleRefresher) RecordActivity(context.Context, string) error { return nil } - -type testTaskActivityRefresher struct { - ch chan uuid.UUID -} - -func (r *testTaskActivityRefresher) Refresh(_ context.Context, taskID uuid.UUID) error { - select { - case r.ch <- taskID: - default: - } - return nil -} - -func (r *testTaskActivityRefresher) ForceRefresh(context.Context, uuid.UUID) error { - return nil -} - -var _ service.TaskActivityRefresher = (*testTaskActivityRefresher)(nil) diff --git a/backend/biz/task/repo/task.go b/backend/biz/task/repo/task.go index 63b155ffc..e3ba5d469 100644 --- a/backend/biz/task/repo/task.go +++ b/backend/biz/task/repo/task.go @@ -286,6 +286,19 @@ func (t *TaskRepo) RefreshLastActiveAt(ctx context.Context, id uuid.UUID, at tim return err } +func (t *TaskRepo) RefreshLastActiveAtByVMID(ctx context.Context, vmID string, at time.Time, minInterval time.Duration) error { + up := t.db.Task.Update().Where( + task.HasVmsWith(virtualmachine.ID(vmID)), + task.StatusNotIn(consts.TaskStatusFinished, consts.TaskStatusError), + task.LastActiveAtLT(at), + ) + if minInterval > 0 { + up = up.Where(task.LastActiveAtLT(at.Add(-minInterval))) + } + _, err := up.SetLastActiveAt(at).Save(ctx) + return err +} + // Delete implements domain.TaskRepo. func (t *TaskRepo) Delete(ctx context.Context, user *domain.User, id uuid.UUID) error { _, err := t.db.Task.Delete(). diff --git a/backend/biz/task/repo/task_activity_test.go b/backend/biz/task/repo/task_activity_test.go new file mode 100644 index 000000000..7919ddd3e --- /dev/null +++ b/backend/biz/task/repo/task_activity_test.go @@ -0,0 +1,91 @@ +package repo + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + _ "github.com/mattn/go-sqlite3" + + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/db/enttest" +) + +func TestRefreshLastActiveAtByVMIDUpdatesOnlyActiveTasks(t *testing.T) { + ctx := context.Background() + client := enttest.Open(t, "sqlite3", fmt.Sprintf("file:task-vm-activity-%s?mode=memory&cache=shared&_fk=1", uuid.NewString())) + t.Cleanup(func() { _ = client.Close() }) + + userID := uuid.New() + if _, err := client.User.Create().SetID(userID).SetName("tester").SetRole(consts.UserRoleIndividual).SetStatus(consts.UserStatusActive).Save(ctx); err != nil { + t.Fatal(err) + } + if _, err := client.Host.Create().SetID("host-1").SetUserID(userID).SetHostname("host").Save(ctx); err != nil { + t.Fatal(err) + } + for _, vmID := range []string{"vm-active", "vm-other"} { + if _, err := client.VirtualMachine.Create().SetID(vmID).SetHostID("host-1").SetUserID(userID).SetName(vmID).Save(ctx); err != nil { + t.Fatal(err) + } + } + + old := time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC) + activeTaskID := createVMActivityTask(t, client, userID, "vm-active", consts.TaskStatusProcessing, old) + finishedTaskID := createVMActivityTask(t, client, userID, "vm-active", consts.TaskStatusFinished, old) + otherTaskID := createVMActivityTask(t, client, userID, "vm-other", consts.TaskStatusProcessing, old) + + repo := &TaskRepo{db: client} + activityAt := old.Add(10 * time.Minute) + if err := repo.RefreshLastActiveAtByVMID(ctx, "vm-active", activityAt, 5*time.Minute); err != nil { + t.Fatal(err) + } + + assertTaskLastActiveAt(t, client, activeTaskID, activityAt) + assertTaskLastActiveAt(t, client, finishedTaskID, old) + assertTaskLastActiveAt(t, client, otherTaskID, old) + + if err := repo.RefreshLastActiveAtByVMID(ctx, "vm-active", activityAt.Add(2*time.Minute), 5*time.Minute); err != nil { + t.Fatal(err) + } + assertTaskLastActiveAt(t, client, activeTaskID, activityAt) + + if err := repo.RefreshLastActiveAtByVMID(ctx, "vm-active", old, 0); err != nil { + t.Fatal(err) + } + assertTaskLastActiveAt(t, client, activeTaskID, activityAt) +} + +func createVMActivityTask(t *testing.T, client *db.Client, userID uuid.UUID, vmID string, status consts.TaskStatus, lastActiveAt time.Time) uuid.UUID { + t.Helper() + ctx := context.Background() + taskID := uuid.New() + if _, err := client.Task.Create(). + SetID(taskID). + SetUserID(userID). + SetKind(consts.TaskTypeDevelop). + SetContent("content"). + SetStatus(status). + SetCreatedAt(lastActiveAt.Add(-time.Hour)). + SetLastActiveAt(lastActiveAt). + Save(ctx); err != nil { + t.Fatal(err) + } + if _, err := client.TaskVirtualMachine.Create().SetID(uuid.New()).SetTaskID(taskID).SetVirtualmachineID(vmID).Save(ctx); err != nil { + t.Fatal(err) + } + return taskID +} + +func assertTaskLastActiveAt(t *testing.T, client *db.Client, taskID uuid.UUID, want time.Time) { + t.Helper() + task, err := client.Task.Get(context.Background(), taskID) + if err != nil { + t.Fatal(err) + } + if !task.LastActiveAt.Equal(want) { + t.Fatalf("task %s last active at = %v, want %v", taskID, task.LastActiveAt, want) + } +} diff --git a/backend/pkg/vmrecycle/analyzer.go b/backend/pkg/vmrecycle/analyzer.go index ffe385a97..3f223f301 100644 --- a/backend/pkg/vmrecycle/analyzer.go +++ b/backend/pkg/vmrecycle/analyzer.go @@ -19,7 +19,6 @@ import ( "github.com/chaitin/MonkeyCode/backend/domain" "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" "github.com/chaitin/MonkeyCode/backend/pkg/entx" - "github.com/chaitin/MonkeyCode/backend/pkg/tasklog" ) const ( @@ -27,15 +26,6 @@ const ( analyzeWorkerCount = 8 ) -var activityEvents = []string{ - "user-input", - "task-started", - "task-running", - "task-ended", - "task-error", - "task-event", -} - type Decision string const ( @@ -49,10 +39,11 @@ const ( ) type TaskInfo struct { - ID uuid.UUID - Status consts.TaskStatus - LogStore consts.LogStore - CreatedAt time.Time + ID uuid.UUID + Status consts.TaskStatus + LogStore consts.LogStore + CreatedAt time.Time + LastActiveAt time.Time } type Analysis struct { @@ -79,10 +70,6 @@ type Analyzer interface { AnalyzeTargets(ctx context.Context, taskIDs []uuid.UUID, vmIDs []string) ([]Analysis, error) } -type activityReader interface { - LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string, store consts.LogStore) (time.Time, bool, error) -} - type deadlineReader interface { GetRunAt(ctx context.Context, queue, id string) (time.Time, bool, error) } @@ -91,7 +78,6 @@ type analyzer struct { cfg *config.Config db *db.Client teamPolicyRepo domain.TeamPolicyRepo - activity activityReader deadlines deadlineReader now func() time.Time } @@ -101,7 +87,6 @@ func NewAnalyzer(i *do.Injector) (Analyzer, error) { cfg: do.MustInvoke[*config.Config](i), db: do.MustInvoke[*db.Client](i), teamPolicyRepo: do.MustInvoke[domain.TeamPolicyRepo](i), - activity: do.MustInvoke[*tasklog.Gateway](i), deadlines: do.MustInvoke[*delayqueue.VMRecycleQueue](i), now: time.Now, }, nil @@ -233,7 +218,13 @@ func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysi if tk.LogStore != nil && *tk.LogStore != "" { store = *tk.LogStore } - item.Tasks = append(item.Tasks, TaskInfo{ID: tk.ID, Status: tk.Status, LogStore: store, CreatedAt: tk.CreatedAt}) + item.Tasks = append(item.Tasks, TaskInfo{ + ID: tk.ID, + Status: tk.Status, + LogStore: store, + CreatedAt: tk.CreatedAt, + LastActiveAt: tk.LastActiveAt, + }) } if vm.IsRecycled { item.Decision = DecisionAlreadyRecycled @@ -260,15 +251,9 @@ func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysi item.RecycleSeconds = policy.EffectiveRecycleSeconds now := a.now() for _, tk := range item.Tasks { - lastActivity := tk.CreatedAt - latest, ok, err := a.activity.LatestEventTime(ctx, tk.ID, tk.CreatedAt, now, activityEvents, tk.LogStore) - if err != nil { - item.Decision = DecisionUnavailable - item.Reason = fmt.Sprintf("query task %s activity: %v", tk.ID, err) - return item - } - if ok && latest.After(lastActivity) { - lastActivity = latest + lastActivity := tk.LastActiveAt + if lastActivity.IsZero() || lastActivity.Before(tk.CreatedAt) { + lastActivity = tk.CreatedAt } if lastActivity.After(item.LastActivityAt) { item.LastActivityAt = lastActivity @@ -321,7 +306,6 @@ func (a *analyzer) loadVM(ctx context.Context, vmID string) (*db.VirtualMachine, } var _ Analyzer = (*analyzer)(nil) -var _ activityReader = (*tasklog.Gateway)(nil) var _ deadlineReader = (*delayqueue.VMRecycleQueue)(nil) func IsExecutable(decision Decision) bool { diff --git a/backend/pkg/vmrecycle/analyzer_test.go b/backend/pkg/vmrecycle/analyzer_test.go index e0e0ed0ba..9804102d9 100644 --- a/backend/pkg/vmrecycle/analyzer_test.go +++ b/backend/pkg/vmrecycle/analyzer_test.go @@ -2,7 +2,6 @@ package vmrecycle import ( "context" - "errors" "fmt" "testing" "time" @@ -21,19 +20,20 @@ func TestAnalyzerCandidateRequiresBothDeadlinesOverdue(t *testing.T) { now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) taskID := uuid.New() logStore := consts.LogStoreClickHouse + activityAt := now.Add(-90 * time.Minute) vm := &db.VirtualMachine{ ID: "vm-candidate", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ - ID: taskID, - Status: consts.TaskStatusProcessing, - CreatedAt: now.Add(-2 * time.Hour), - LogStore: &logStore, + ID: taskID, + Status: consts.TaskStatusProcessing, + CreatedAt: now.Add(-2 * time.Hour), + LastActiveAt: activityAt, + LogStore: &logStore, }}}, } - activityAt := now.Add(-90 * time.Minute) redisAt := now.Add(-time.Minute) - analyzer := newAnalyzerForTest(now, &analyzerActivityStub{latest: map[uuid.UUID]time.Time{taskID: activityAt}}, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: redisAt}}) + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: redisAt}}) got := analyzer.analyzeVM(context.Background(), vm) if got.Decision != DecisionCandidate { @@ -64,7 +64,7 @@ func TestAnalyzerMarksHistoricalUncertainWhenRedisDeadlineCannotConfirm(t *testi } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - analyzer := newAnalyzerForTest(now, &analyzerActivityStub{}, tt.deadlines) + analyzer := newAnalyzerForTest(now, tt.deadlines) got := analyzer.analyzeVM(context.Background(), vm) if got.Decision != DecisionHistoricalUncertain { t.Fatalf("decision = %s, reason = %s", got.Decision, got.Reason) @@ -78,14 +78,10 @@ func TestAnalyzerUsesLatestActivityAcrossTasks(t *testing.T) { oldTaskID := uuid.New() recentTaskID := uuid.New() vm := &db.VirtualMachine{ID: "vm-multi-task", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{ - {ID: oldTaskID, CreatedAt: now.Add(-3 * time.Hour)}, - {ID: recentTaskID, CreatedAt: now.Add(-2 * time.Hour)}, + {ID: oldTaskID, CreatedAt: now.Add(-3 * time.Hour), LastActiveAt: now.Add(-2 * time.Hour)}, + {ID: recentTaskID, CreatedAt: now.Add(-2 * time.Hour), LastActiveAt: now.Add(-30 * time.Minute)}, }}} - activity := &analyzerActivityStub{latest: map[uuid.UUID]time.Time{ - oldTaskID: now.Add(-2 * time.Hour), - recentTaskID: now.Add(-30 * time.Minute), - }} - analyzer := newAnalyzerForTest(now, activity, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(-time.Hour)}}) + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(-time.Hour)}}) got := analyzer.analyzeVM(context.Background(), vm) if got.Decision != DecisionNotDue { @@ -94,29 +90,22 @@ func TestAnalyzerUsesLatestActivityAcrossTasks(t *testing.T) { if !got.LastActivityAt.Equal(now.Add(-30 * time.Minute)) { t.Fatalf("last activity = %v", got.LastActivityAt) } - if len(activity.events) != len(activityEvents) { - t.Fatalf("activity events = %v", activity.events) - } - if _, ok := activity.events["task-event"]; !ok { - t.Fatalf("activity events = %v, missing task-event", activity.events) - } } -func TestAnalyzerReturnsUnavailableOnLogFailure(t *testing.T) { +func TestAnalyzerFallsBackToCreatedAtWhenLastActiveAtIsZero(t *testing.T) { now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) - wantErr := errors.New("log query failed") taskID := uuid.New() vm := &db.VirtualMachine{ID: "vm-log-error", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ID: taskID, CreatedAt: now.Add(-2 * time.Hour)}}}} - analyzer := newAnalyzerForTest(now, &analyzerActivityStub{err: wantErr}, &analyzerDeadlineStub{}) + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(-time.Minute)}}) got := analyzer.analyzeVM(context.Background(), vm) - if got.Decision != DecisionUnavailable || got.Reason == "" { + if got.Decision != DecisionCandidate || !got.LastActivityAt.Equal(now.Add(-2*time.Hour)) { t.Fatalf("analysis = %+v", got) } } func TestAnalyzerAlreadyRecycledIsSuccessfulNoop(t *testing.T) { - analyzer := newAnalyzerForTest(time.Now(), &analyzerActivityStub{}, &analyzerDeadlineStub{}) + analyzer := newAnalyzerForTest(time.Now(), &analyzerDeadlineStub{}) got := analyzer.analyzeVM(context.Background(), &db.VirtualMachine{ID: "vm-recycled", IsRecycled: true}) if got.Decision != DecisionAlreadyRecycled || !IsSuccessfulNoop(got.Decision) { t.Fatalf("analysis = %+v", got) @@ -161,7 +150,7 @@ func TestAnalyzerDeduplicatesTaskAndVMTargets(t *testing.T) { t.Fatal(err) } - analyzer := newAnalyzerForTest(now, &analyzerActivityStub{}, &analyzerDeadlineStub{runAt: map[string]time.Time{vmID: now.Add(-time.Minute)}}) + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{runAt: map[string]time.Time{vmID: now.Add(-time.Minute)}}) analyzer.db = client items, err := analyzer.AnalyzeTargets(ctx, []uuid.UUID{taskID, taskID}, []string{vmID, vmID}) if err != nil { @@ -172,7 +161,7 @@ func TestAnalyzerDeduplicatesTaskAndVMTargets(t *testing.T) { } } -func newAnalyzerForTest(now time.Time, activity activityReader, deadlines deadlineReader) *analyzer { +func newAnalyzerForTest(now time.Time, deadlines deadlineReader) *analyzer { return &analyzer{ cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}}, teamPolicyRepo: &analyzerTeamPolicyRepoStub{team: &db.Team{ @@ -183,30 +172,11 @@ func newAnalyzerForTest(now time.Time, activity activityReader, deadlines deadli TaskVMRecycleEnabled: true, TaskVMRecycleSeconds: 0, }}, - activity: activity, deadlines: deadlines, now: func() time.Time { return now }, } } -type analyzerActivityStub struct { - latest map[uuid.UUID]time.Time - err error - events map[string]struct{} -} - -func (s *analyzerActivityStub) LatestEventTime(_ context.Context, taskID uuid.UUID, _, _ time.Time, events []string, _ consts.LogStore) (time.Time, bool, error) { - s.events = make(map[string]struct{}, len(events)) - for _, event := range events { - s.events[event] = struct{}{} - } - if s.err != nil { - return time.Time{}, false, s.err - } - latest, ok := s.latest[taskID] - return latest, ok, nil -} - type analyzerDeadlineStub struct { runAt map[string]time.Time err error From d2b2fba6299f641cd4f2de2f5967ce7f34117bbe Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 12:10:51 +0800 Subject: [PATCH 04/10] =?UTF-8?q?=E7=A7=BB=E9=99=A4=E6=9C=AA=E4=BD=BF?= =?UTF-8?q?=E7=94=A8=E7=9A=84=E4=BB=BB=E5=8A=A1=E6=97=A5=E5=BF=97=E6=B4=BB?= =?UTF-8?q?=E5=8A=A8=E6=9F=A5=E8=AF=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/pkg/loki/client.go | 62 +++++++---------- backend/pkg/loki/client_test.go | 67 +------------------ backend/pkg/tasklog/clickhouse_provider.go | 25 ------- .../pkg/tasklog/clickhouse_provider_test.go | 47 ------------- backend/pkg/tasklog/gateway.go | 8 --- backend/pkg/tasklog/gateway_test.go | 20 ------ backend/pkg/tasklog/loki_provider.go | 7 -- backend/pkg/tasklog/provider.go | 1 - 8 files changed, 24 insertions(+), 213 deletions(-) diff --git a/backend/pkg/loki/client.go b/backend/pkg/loki/client.go index 4067e1559..0dd5a6bfa 100644 --- a/backend/pkg/loki/client.go +++ b/backend/pkg/loki/client.go @@ -597,19 +597,7 @@ func (c *Client) tailWebSocketSession( // 优先从 Loki structured metadata (labels) 中读取 event 字段，回退到解析 JSON body。 // end 为搜索的结束时间上界，零值表示 time.Now()。 func (c *Client) FindLastEvent(ctx context.Context, taskID string, event string, start, end time.Time) (time.Time, error) { - latest, _, err := c.FindLastEventIn(ctx, taskID, []string{event}, start, end) - return latest, err -} - -func (c *Client) FindLastEventIn(ctx context.Context, taskID string, events []string, start, end time.Time) (time.Time, bool, error) { const pageSize = 200 - if len(events) == 0 { - return time.Time{}, false, nil - } - eventSet := make(map[string]struct{}, len(events)) - for _, event := range events { - eventSet[event] = struct{}{} - } if end.IsZero() { end = time.Now() @@ -621,40 +609,36 @@ func (c *Client) FindLastEventIn(ctx context.Context, taskID string, events []st for { entries, err := c.QueryByTaskID(ctx, taskID, start, end, pageSize, "backward") if err != nil { - return time.Time{}, false, fmt.Errorf("FindLastEventIn query failed: %w", err) + return time.Time{}, fmt.Errorf("FindLastEvent query failed: %w", err) } - if latest, ok := latestMatchingEvent(entries, eventSet); ok { - return latest, true, nil + + for _, entry := range entries { + // 优先从 labels（structured metadata）读取 + if ev, ok := entry.Labels["event"]; ok { + if ev == event { + return entry.Timestamp, nil + } + continue + } + // 回退：解析 JSON body + var chunk struct { + Event string `json:"event"` + } + if err := json.Unmarshal([]byte(entry.Line), &chunk); err != nil { + continue + } + if chunk.Event == event { + return entry.Timestamp, nil + } } if len(entries) < pageSize { - return time.Time{}, false, nil + return time.Time{}, nil } - end = entries[len(entries)-1].Timestamp.Add(-time.Nanosecond) - } -} - -func latestMatchingEvent(entries []LogEntry, events map[string]struct{}) (time.Time, bool) { - for _, entry := range entries { - if event, ok := entry.Labels["event"]; ok { - _, matched := events[event] - if matched { - return entry.Timestamp, true - } - continue - } - var chunk struct { - Event string `json:"event"` - } - if err := json.Unmarshal([]byte(entry.Line), &chunk); err != nil { - continue - } - if _, ok := events[chunk.Event]; ok { - return entry.Timestamp, true - } + // 继续往前扫描 + end = entries[len(entries)-1].Timestamp } - return time.Time{}, false } // FindLatestRoundStart 定位 attach 模式下最新论次的起点。 diff --git a/backend/pkg/loki/client_test.go b/backend/pkg/loki/client_test.go index f5e8b85ed..cb9121544 100644 --- a/backend/pkg/loki/client_test.go +++ b/backend/pkg/loki/client_test.go @@ -1,11 +1,7 @@ package loki import ( - "context" "encoding/json" - "net/http" - "net/http/httptest" - "strconv" "testing" "time" ) @@ -58,7 +54,7 @@ func TestFindLatestRoundStart(t *testing.T) { {Timestamp: base.Add(2 * time.Second), Line: marshalChunk(t, "task-started")}, {Timestamp: base.Add(3 * time.Second), Line: marshalChunk(t, "task-running")}, }, - want: taskCreatedAt, + want: taskCreatedAt, }, } @@ -102,64 +98,3 @@ func TestFilterEntriesByTimeWindow(t *testing.T) { } } } - -func TestLatestMatchingEventUsesLabelsAndJSONFallback(t *testing.T) { - base := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) - events := map[string]struct{}{"user-input": {}, "task-event": {}} - entries := []LogEntry{ - {Timestamp: base.Add(3 * time.Minute), Labels: map[string]string{"event": "task-running"}, Line: marshalChunk(t, "task-event")}, - {Timestamp: base.Add(2 * time.Minute), Labels: map[string]string{"event": "task-event"}}, - {Timestamp: base.Add(time.Minute), Line: marshalChunk(t, "user-input")}, - } - - got, ok := latestMatchingEvent(entries, events) - if !ok || !got.Equal(base.Add(2*time.Minute)) { - t.Fatalf("latest = %v, ok = %v", got, ok) - } -} - -func TestLatestMatchingEventReturnsMissing(t *testing.T) { - entries := []LogEntry{{Timestamp: time.Now(), Line: marshalChunk(t, "task-running")}} - if _, ok := latestMatchingEvent(entries, map[string]struct{}{"task-event": {}}); ok { - t.Fatal("expected no matching event") - } -} - -func TestFindLastEventInPagesBackward(t *testing.T) { - base := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) - requests := 0 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - requests++ - values := make([][]string, 0, 200) - event := "task-running" - if requests == 1 { - for i := range 200 { - ts := base.Add(-time.Duration(i) * time.Second) - values = append(values, []string{strconv.FormatInt(ts.UnixNano(), 10), marshalChunk(t, "task-running")}) - } - } else { - event = "task-event" - values = append(values, []string{strconv.FormatInt(base.Add(-201*time.Second).UnixNano(), 10), marshalChunk(t, "task-event")}) - } - _ = json.NewEncoder(w).Encode(map[string]any{ - "status": "success", - "data": map[string]any{ - "resultType": "streams", - "result": []any{map[string]any{ - "stream": map[string]string{"task_id": "task-1", "event": event}, - "values": values, - }}, - }, - }) - })) - defer server.Close() - - client := NewClient(server.URL) - got, ok, err := client.FindLastEventIn(context.Background(), "task-1", []string{"user-input", "task-event"}, base.Add(-time.Hour), base.Add(time.Second)) - if err != nil || !ok { - t.Fatalf("FindLastEventIn() ok = %v, err = %v", ok, err) - } - if !got.Equal(base.Add(-201*time.Second)) || requests != 2 { - t.Fatalf("latest = %v, requests = %d", got, requests) - } -} diff --git a/backend/pkg/tasklog/clickhouse_provider.go b/backend/pkg/tasklog/clickhouse_provider.go index 1a5b34216..3b76704ac 100644 --- a/backend/pkg/tasklog/clickhouse_provider.go +++ b/backend/pkg/tasklog/clickhouse_provider.go @@ -5,7 +5,6 @@ import ( "database/sql" "fmt" "strconv" - "strings" "time" "github.com/google/uuid" @@ -25,30 +24,6 @@ func (p *ClickHouseProvider) Name() string { return "clickhouse" } -func (p *ClickHouseProvider) LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string) (time.Time, bool, error) { - if p.client == nil { - return time.Time{}, false, ErrProviderUnavailable - } - if len(events) == 0 { - return time.Time{}, false, nil - } - placeholders := strings.TrimSuffix(strings.Repeat("?,", len(events)), ",") - query := fmt.Sprintf(`SELECT maxOrNull(ts) FROM %s WHERE task_id = ? AND ts >= ? AND ts <= ? AND event IN (%s)`, p.client.Table(), placeholders) - args := make([]any, 0, len(events)+3) - args = append(args, taskID, start, end) - for _, event := range events { - args = append(args, event) - } - var latest sql.NullTime - if err := p.client.QueryRowContext(ctx, query, args...).Scan(&latest); err != nil { - return time.Time{}, false, err - } - if !latest.Valid { - return time.Time{}, false, nil - } - return latest.Time.UTC(), true, nil -} - func (p *ClickHouseProvider) QueryLatestTurn(ctx context.Context, taskID uuid.UUID, taskCreatedAt, end time.Time) (*QueryLatestTurnResp, error) { if p.client == nil { return nil, ErrProviderUnavailable diff --git a/backend/pkg/tasklog/clickhouse_provider_test.go b/backend/pkg/tasklog/clickhouse_provider_test.go index c507e296c..8107730f1 100644 --- a/backend/pkg/tasklog/clickhouse_provider_test.go +++ b/backend/pkg/tasklog/clickhouse_provider_test.go @@ -13,53 +13,6 @@ import ( "github.com/chaitin/MonkeyCode/backend/pkg/tasklog" ) -func TestClickHouseProviderLatestEventTimeFiltersEvents(t *testing.T) { - db, mock, err := sqlmock.New() - if err != nil { - t.Fatal(err) - } - defer db.Close() - taskID := uuid.New() - start := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC) - end := time.Date(2026, 7, 13, 0, 0, 0, 0, time.UTC) - want := end.Add(-time.Minute) - mock.ExpectQuery("SELECT maxOrNull\\(ts\\)[\\s\\S]*event IN \\(\\?,\\?\\)"). - WithArgs(taskID, start, end, "user-input", "task-event"). - WillReturnRows(sqlmock.NewRows([]string{"max"}).AddRow(want)) - - provider := tasklog.NewClickHouseProvider(clickhouse.NewWithDBAndTable(db, "task_logs_test")) - got, ok, err := provider.LatestEventTime(context.Background(), taskID, start, end, []string{"user-input", "task-event"}) - if err != nil || !ok { - t.Fatalf("LatestEventTime() ok = %v, err = %v", ok, err) - } - if !got.Equal(want) { - t.Fatalf("latest = %v, want %v", got, want) - } - if err := mock.ExpectationsWereMet(); err != nil { - t.Fatal(err) - } -} - -func TestClickHouseProviderLatestEventTimeHandlesNoMatch(t *testing.T) { - db, mock, err := sqlmock.New() - if err != nil { - t.Fatal(err) - } - defer db.Close() - taskID := uuid.New() - start := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC) - end := time.Date(2026, 7, 13, 0, 0, 0, 0, time.UTC) - mock.ExpectQuery("SELECT maxOrNull\\(ts\\)"). - WithArgs(taskID, start, end, "task-event"). - WillReturnRows(sqlmock.NewRows([]string{"max"}).AddRow(nil)) - - provider := tasklog.NewClickHouseProvider(clickhouse.NewWithDBAndTable(db, "task_logs_test")) - _, ok, err := provider.LatestEventTime(context.Background(), taskID, start, end, []string{"task-event"}) - if err != nil || ok { - t.Fatalf("LatestEventTime() ok = %v, err = %v", ok, err) - } -} - func TestClickHouseProviderQueryLatestTurnUsesTurnSeqCursor(t *testing.T) { db, mock, err := sqlmock.New() if err != nil { diff --git a/backend/pkg/tasklog/gateway.go b/backend/pkg/tasklog/gateway.go index 5a34013fe..d452dbfe2 100644 --- a/backend/pkg/tasklog/gateway.go +++ b/backend/pkg/tasklog/gateway.go @@ -16,14 +16,6 @@ type Gateway struct { ClickHouse Provider } -func (g *Gateway) LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string, store consts.LogStore) (time.Time, bool, error) { - p, err := g.providerByStore(store) - if err != nil { - return time.Time{}, false, err - } - return p.LatestEventTime(ctx, taskID, start, end, events) -} - func (g *Gateway) QueryLatestTurn(ctx context.Context, taskID uuid.UUID, taskCreatedAt, end time.Time, store consts.LogStore) (*QueryLatestTurnResp, error) { p, err := g.providerByStore(store) if err != nil { diff --git a/backend/pkg/tasklog/gateway_test.go b/backend/pkg/tasklog/gateway_test.go index fa7def7fa..ef2ca7518 100644 --- a/backend/pkg/tasklog/gateway_test.go +++ b/backend/pkg/tasklog/gateway_test.go @@ -16,12 +16,6 @@ type gatewayProviderStub struct { name string queryLatestTurnCalled bool queryTurnsCalled bool - latestEventTimeCalled bool -} - -func (s *gatewayProviderStub) LatestEventTime(context.Context, uuid.UUID, time.Time, time.Time, []string) (time.Time, bool, error) { - s.latestEventTimeCalled = true - return time.Now(), true, nil } func (s *gatewayProviderStub) Name() string { @@ -76,20 +70,6 @@ func TestGatewayClickHouseStoreUsesClickHouse(t *testing.T) { } } -func TestGatewayLatestEventTimeUsesConfiguredStore(t *testing.T) { - loki := &gatewayProviderStub{name: "loki"} - clickHouse := &gatewayProviderStub{name: "clickhouse"} - gateway := &Gateway{Loki: loki, ClickHouse: clickHouse} - - _, _, err := gateway.LatestEventTime(context.Background(), uuid.New(), time.Now().Add(-time.Hour), time.Now(), []string{"task-event"}, consts.LogStoreClickHouse) - if err != nil { - t.Fatal(err) - } - if !clickHouse.latestEventTimeCalled || loki.latestEventTimeCalled { - t.Fatalf("loki called = %v, clickhouse called = %v", loki.latestEventTimeCalled, clickHouse.latestEventTimeCalled) - } -} - func TestGatewayUnknownStoreReturnsError(t *testing.T) { loki := &gatewayProviderStub{name: "loki"} clickHouse := &gatewayProviderStub{name: "clickhouse"} diff --git a/backend/pkg/tasklog/loki_provider.go b/backend/pkg/tasklog/loki_provider.go index 94a8336dc..1691331ca 100644 --- a/backend/pkg/tasklog/loki_provider.go +++ b/backend/pkg/tasklog/loki_provider.go @@ -26,13 +26,6 @@ func (p *LokiProvider) Name() string { return "loki" } -func (p *LokiProvider) LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string) (time.Time, bool, error) { - if p.client == nil { - return time.Time{}, false, ErrProviderUnavailable - } - return p.client.FindLastEventIn(ctx, taskID.String(), events, start, end) -} - func (p *LokiProvider) QueryWindow(ctx context.Context, taskID uuid.UUID, start, end time.Time) ([]Entry, error) { if p.client == nil { return nil, ErrProviderUnavailable diff --git a/backend/pkg/tasklog/provider.go b/backend/pkg/tasklog/provider.go index f35c3a95c..8f7ca4a80 100644 --- a/backend/pkg/tasklog/provider.go +++ b/backend/pkg/tasklog/provider.go @@ -9,7 +9,6 @@ import ( type Provider interface { Name() string - LatestEventTime(ctx context.Context, taskID uuid.UUID, start, end time.Time, events []string) (time.Time, bool, error) QueryLatestTurn(ctx context.Context, taskID uuid.UUID, taskCreatedAt, end time.Time) (*QueryLatestTurnResp, error) QueryTurns(ctx context.Context, taskID uuid.UUID, taskCreatedAt time.Time, opts QueryTurnsOpts) (*QueryTurnsResp, error) QueryUserInputs(ctx context.Context, taskID uuid.UUID, taskCreatedAt time.Time, cursor string, limit int) (*QueryUserInputsResp, error) From b550624cbe78a0a51ae97f773fb0caaaca1aac6d Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 12:17:29 +0800 Subject: [PATCH 05/10] =?UTF-8?q?=E5=9C=A8=E7=94=A8=E6=88=B7=E8=BE=93?= =?UTF-8?q?=E5=85=A5=E5=90=8E=E6=8C=81=E4=B9=85=E5=8C=96=E4=BB=BB=E5=8A=A1?= =?UTF-8?q?=E6=B4=BB=E5=8A=A8=E6=97=B6=E9=97=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/biz/task/usecase/task.go | 3 ++ .../biz/task/usecase/task_activity_test.go | 29 ++++++++++++------- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/backend/biz/task/usecase/task.go b/backend/biz/task/usecase/task.go index c4741abde..9008a5256 100644 --- a/backend/biz/task/usecase/task.go +++ b/backend/biz/task/usecase/task.go @@ -432,6 +432,9 @@ func (a *TaskUsecase) Continue(ctx context.Context, user *domain.User, id uuid.U }); err != nil { return err } + if err := a.taskActivityRefresher.ForceRefresh(ctx, id); err != nil { + a.logger.WarnContext(ctx, "failed to refresh task last active on user input", "task_id", id, "error", err) + } if err := a.idleRefresher.RecordActivity(ctx, tk.VirtualMachine.ID); err != nil { a.logger.WarnContext(ctx, "failed to refresh vm idle timers on user input", "task_id", id, "vm_id", tk.VirtualMachine.ID, "error", err) } diff --git a/backend/biz/task/usecase/task_activity_test.go b/backend/biz/task/usecase/task_activity_test.go index 8c5b278f7..01e1c184a 100644 --- a/backend/biz/task/usecase/task_activity_test.go +++ b/backend/biz/task/usecase/task_activity_test.go @@ -51,7 +51,8 @@ func TestRefreshCreatedTaskStateAlwaysRefreshesIdleTimer(t *testing.T) { func TestContinueRecordsVMActivityAfterTaskflowAcceptsInput(t *testing.T) { manager := &continueTaskManagerStub{} idleRefresher := &vmIdleRefresherStub{} - u, user, taskID := newContinueTaskUsecase(t, manager, idleRefresher) + u, user, taskID, taskRefresher := newContinueTaskUsecase(t, manager, idleRefresher) + taskRefresher.err = errors.New("db write failed") if err := u.Continue(context.Background(), user, taskID, domain.ContinueTaskReq{Content: []byte("继续执行")}); err != nil { t.Fatal(err) @@ -59,6 +60,9 @@ func TestContinueRecordsVMActivityAfterTaskflowAcceptsInput(t *testing.T) { if !manager.called { t.Fatal("expected taskflow continue to be called") } + if !taskRefresher.forceCalled || taskRefresher.taskID != taskID { + t.Fatalf("task activity called = %v, task id = %s", taskRefresher.forceCalled, taskRefresher.taskID) + } if !idleRefresher.recordActivityCalled || idleRefresher.vmID != "vm-continue" { t.Fatalf("record activity called = %v, vm id = %q", idleRefresher.recordActivityCalled, idleRefresher.vmID) } @@ -71,7 +75,7 @@ func TestContinueDoesNotRecordVMActivityWhenTaskflowRejectsInput(t *testing.T) { wantErr := errors.New("taskflow unavailable") manager := &continueTaskManagerStub{err: wantErr} idleRefresher := &vmIdleRefresherStub{} - u, user, taskID := newContinueTaskUsecase(t, manager, idleRefresher) + u, user, taskID, taskRefresher := newContinueTaskUsecase(t, manager, idleRefresher) err := u.Continue(context.Background(), user, taskID, domain.ContinueTaskReq{Content: []byte("继续执行")}) if !errors.Is(err, wantErr) { @@ -80,9 +84,12 @@ func TestContinueDoesNotRecordVMActivityWhenTaskflowRejectsInput(t *testing.T) { if idleRefresher.recordActivityCalled || idleRefresher.keepAwakeCalled { t.Fatal("did not expect rejected input to refresh vm idle state") } + if taskRefresher.forceCalled { + t.Fatal("did not expect rejected input to refresh task activity") + } } -func newContinueTaskUsecase(t *testing.T, manager *continueTaskManagerStub, idleRefresher *vmIdleRefresherStub) (*TaskUsecase, *domain.User, uuid.UUID) { +func newContinueTaskUsecase(t *testing.T, manager *continueTaskManagerStub, idleRefresher *vmIdleRefresherStub) (*TaskUsecase, *domain.User, uuid.UUID, *taskActivityRefresherStub) { t.Helper() user := &domain.User{ID: uuid.New()} taskID := uuid.New() @@ -101,14 +108,16 @@ func newContinueTaskUsecase(t *testing.T, manager *continueTaskManagerStub, idle Edges: db.TaskEdges{Vms: []*db.VirtualMachine{vm}}, }} redisClient := newTaskActivityRedis(t) + taskRefresher := &taskActivityRefresherStub{} return &TaskUsecase{ - cfg: &config.Config{}, - repo: repo, - logger: slog.New(slog.NewTextHandler(io.Discard, nil)), - taskflow: &continueTaskflowStub{manager: manager}, - redis: redisClient, - idleRefresher: idleRefresher, - }, user, taskID + cfg: &config.Config{}, + repo: repo, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + taskflow: &continueTaskflowStub{manager: manager}, + redis: redisClient, + idleRefresher: idleRefresher, + taskActivityRefresher: taskRefresher, + }, user, taskID, taskRefresher } func newTaskActivityRedis(t *testing.T) *redis.Client { From a598d4712fc6a5b6a9051f1e6710159ec261b757 Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 15:59:48 +0800 Subject: [PATCH 06/10] =?UTF-8?q?=E6=94=AF=E6=8C=81=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E7=BC=BA=E5=A4=B1=E7=9A=84=E8=99=9A=E6=8B=9F=E6=9C=BA=E5=9B=9E?= =?UTF-8?q?=E6=94=B6=E6=97=B6=E9=97=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/pkg/delayqueue/delayqueue.go | 29 ++++ backend/pkg/delayqueue/delayqueue_test.go | 51 +++++++ backend/pkg/register.go | 1 + backend/pkg/vmrecycle/analyzer.go | 68 ++++++--- backend/pkg/vmrecycle/analyzer_test.go | 66 +++++--- backend/pkg/vmrecycle/deadline_repair.go | 122 +++++++++++++++ backend/pkg/vmrecycle/deadline_repair_test.go | 141 ++++++++++++++++++ 7 files changed, 443 insertions(+), 35 deletions(-) create mode 100644 backend/pkg/vmrecycle/deadline_repair.go create mode 100644 backend/pkg/vmrecycle/deadline_repair_test.go diff --git a/backend/pkg/delayqueue/delayqueue.go b/backend/pkg/delayqueue/delayqueue.go index dabfc02ba..34216dd85 100644 --- a/backend/pkg/delayqueue/delayqueue.go +++ b/backend/pkg/delayqueue/delayqueue.go @@ -115,6 +115,26 @@ func (q *RedisDelayQueue[T]) Enqueue(ctx context.Context, queue string, payload return id, nil } +// EnqueueIfMissing 仅在延迟队列中不存在指定 ID 时入队。 +func (q *RedisDelayQueue[T]) EnqueueIfMissing(ctx context.Context, queue string, payload T, runAt time.Time, id string) (string, bool, error) { + if id == "" { + id = uuid.NewString() + } + + b, err := json.Marshal(&jobData[T]{Payload: payload, Attempts: 0}) + if err != nil { + return "", false, err + } + added, err := enqueueIfMissingScript.Run(ctx, q.rdb, []string{ + q.jobKey(queue, id), + q.zsetKey(queue), + }, id, b, q.jobTTL.Milliseconds(), runAt.UnixMilli()).Int() + if err != nil { + return "", false, err + } + return id, added == 1, nil +} + // GetJobInfo 查询任务信息 func (q *RedisDelayQueue[T]) GetJobInfo(ctx context.Context, queue, id string) (*Job[T], time.Time, bool, error) { runAt, ok, err := q.GetRunAt(ctx, queue, id) @@ -325,3 +345,12 @@ if #items > 0 then end return items `) + +var enqueueIfMissingScript = redis.NewScript(` +if redis.call('ZSCORE', KEYS[2], ARGV[1]) then + return 0 +end +redis.call('SET', KEYS[1], ARGV[2], 'PX', ARGV[3]) +redis.call('ZADD', KEYS[2], ARGV[4], ARGV[1]) +return 1 +`) diff --git a/backend/pkg/delayqueue/delayqueue_test.go b/backend/pkg/delayqueue/delayqueue_test.go index f57f87bd1..fa0406388 100644 --- a/backend/pkg/delayqueue/delayqueue_test.go +++ b/backend/pkg/delayqueue/delayqueue_test.go @@ -45,3 +45,54 @@ func TestGetRunAtReturnsMissing(t *testing.T) { t.Fatalf("GetRunAt() ok = %v, err = %v", ok, err) } } + +func TestEnqueueIfMissingAddsNewJob(t *testing.T) { + ctx := context.Background() + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + queue := NewRedisDelayQueue[string](rdb, slog.Default()) + runAt := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + + id, added, err := queue.EnqueueIfMissing(ctx, "recycle", "payload", runAt, "vm-1") + if err != nil { + t.Fatal(err) + } + if id != "vm-1" || !added { + t.Fatalf("id = %q, added = %v", id, added) + } + job, gotRunAt, ok, err := queue.GetJobInfo(ctx, "recycle", "vm-1") + if err != nil || !ok { + t.Fatalf("GetJobInfo() ok = %v, err = %v", ok, err) + } + if job.Payload != "payload" || !gotRunAt.Equal(runAt) { + t.Fatalf("job = %+v, run at = %v", job, gotRunAt) + } +} + +func TestEnqueueIfMissingDoesNotOverwriteExistingJob(t *testing.T) { + ctx := context.Background() + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + queue := NewRedisDelayQueue[string](rdb, slog.Default()) + originalRunAt := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + if _, err := queue.Enqueue(ctx, "recycle", "original", originalRunAt, "vm-1"); err != nil { + t.Fatal(err) + } + + _, added, err := queue.EnqueueIfMissing(ctx, "recycle", "replacement", originalRunAt.Add(time.Hour), "vm-1") + if err != nil { + t.Fatal(err) + } + if added { + t.Fatal("existing job must not be overwritten") + } + job, gotRunAt, ok, err := queue.GetJobInfo(ctx, "recycle", "vm-1") + if err != nil || !ok { + t.Fatalf("GetJobInfo() ok = %v, err = %v", ok, err) + } + if job.Payload != "original" || !gotRunAt.Equal(originalRunAt) { + t.Fatalf("job = %+v, run at = %v", job, gotRunAt) + } +} diff --git a/backend/pkg/register.go b/backend/pkg/register.go index 94bb2d854..c1cbf92be 100644 --- a/backend/pkg/register.go +++ b/backend/pkg/register.go @@ -194,6 +194,7 @@ func RegisterInfra(i *do.Injector, w ...*web.Web) error { do.Provide(i, vmrecycle.NewRecycler) do.Provide(i, vmrecycle.NewAnalyzer) + do.Provide(i, vmrecycle.NewDeadlineRepairer) // Channel Registry（通知渠道） do.Provide(i, func(i *do.Injector) (*msgpush.WechatClient, error) { diff --git a/backend/pkg/vmrecycle/analyzer.go b/backend/pkg/vmrecycle/analyzer.go index 3f223f301..51d399953 100644 --- a/backend/pkg/vmrecycle/analyzer.go +++ b/backend/pkg/vmrecycle/analyzer.go @@ -32,6 +32,8 @@ const ( DecisionCandidate Decision = "candidate" DecisionUnavailable Decision = "unavailable" DecisionHistoricalUncertain Decision = "historical_uncertain" + DecisionDeadlineMissing Decision = "deadline_missing" + DecisionOrphan Decision = "orphan" DecisionNotDue Decision = "not_due" DecisionDisabled Decision = "disabled" DecisionAlreadyRecycled Decision = "already_recycled" @@ -116,7 +118,7 @@ func (a *analyzer) Scan(ctx context.Context) (ScanReport, error) { items := a.analyzeBatch(ctx, vms) for _, item := range items { report.Counts[item.Decision]++ - if item.Decision == DecisionCandidate || item.Decision == DecisionUnavailable || item.Decision == DecisionHistoricalUncertain { + if shouldReport(item.Decision) { report.Items = append(report.Items, item) } } @@ -209,7 +211,7 @@ func (a *analyzer) analyzeBatch(ctx context.Context, vms []*db.VirtualMachine) [ } func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysis { - item := Analysis{VMID: vm.ID, AlreadyRecycled: vm.IsRecycled} + item := Analysis{Target: "vm:" + vm.ID, VMID: vm.ID, AlreadyRecycled: vm.IsRecycled} for _, tk := range vm.Edges.Tasks { if tk == nil { continue @@ -232,7 +234,7 @@ func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysi return item } if len(item.Tasks) == 0 { - item.Decision = DecisionUnavailable + item.Decision = DecisionOrphan item.Reason = "VM has no task" return item } @@ -269,17 +271,31 @@ func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysi if ok { item.RedisRecycleAt = &redisAt } + if item.RedisRecycleAt == nil { + if allTasksTerminal(item.Tasks) { + if now.Before(item.DueAt) { + item.Decision = DecisionNotDue + item.Reason = "activity deadline not reached" + return item + } + item.Overdue = now.Sub(item.DueAt) + item.Decision = DecisionCandidate + item.Reason = "terminal tasks overdue and redis recycle deadline missing" + return item + } + if now.After(item.DueAt) { + item.Overdue = now.Sub(item.DueAt) + } + item.Decision = DecisionDeadlineMissing + item.Reason = "active task redis recycle deadline missing" + return item + } if now.Before(item.DueAt) { item.Decision = DecisionNotDue item.Reason = "activity deadline not reached" return item } item.Overdue = now.Sub(item.DueAt) - if item.RedisRecycleAt == nil { - item.Decision = DecisionHistoricalUncertain - item.Reason = "redis recycle deadline missing" - return item - } if now.Before(*item.RedisRecycleAt) { item.Decision = DecisionHistoricalUncertain item.Reason = "redis recycle deadline is still in the future" @@ -291,25 +307,43 @@ func (a *analyzer) analyzeVM(ctx context.Context, vm *db.VirtualMachine) Analysi } func (a *analyzer) resolvePolicy(ctx context.Context, userID uuid.UUID) (*domain.TeamTaskVMIdlePolicy, error) { - team, err := a.teamPolicyRepo.GetTeamByUserID(ctx, userID) - if err != nil && !db.IsNotFound(err) { - return nil, err - } - if db.IsNotFound(err) { - team = nil - } - return domain.ResolveTeamTaskVMIdlePolicy(team, a.cfg.VMIdle) + return resolvePolicy(ctx, a.teamPolicyRepo, a.cfg.VMIdle, userID) } func (a *analyzer) loadVM(ctx context.Context, vmID string) (*db.VirtualMachine, error) { return a.db.VirtualMachine.Query().Where(virtualmachine.ID(vmID)).WithTasks().Only(ctx) } +func allTasksTerminal(tasks []TaskInfo) bool { + if len(tasks) == 0 { + return false + } + for _, task := range tasks { + if task.Status != consts.TaskStatusFinished && task.Status != consts.TaskStatusError { + return false + } + } + return true +} + +func shouldReport(decision Decision) bool { + switch decision { + case DecisionCandidate, DecisionUnavailable, DecisionHistoricalUncertain, DecisionDeadlineMissing, DecisionOrphan: + return true + default: + return false + } +} + var _ Analyzer = (*analyzer)(nil) var _ deadlineReader = (*delayqueue.VMRecycleQueue)(nil) func IsExecutable(decision Decision) bool { - return decision == DecisionCandidate + return decision == DecisionCandidate || decision == DecisionOrphan +} + +func IsDeadlineRepairable(decision Decision) bool { + return decision == DecisionDeadlineMissing } func IsSuccessfulNoop(decision Decision) bool { diff --git a/backend/pkg/vmrecycle/analyzer_test.go b/backend/pkg/vmrecycle/analyzer_test.go index 9804102d9..77845b8ac 100644 --- a/backend/pkg/vmrecycle/analyzer_test.go +++ b/backend/pkg/vmrecycle/analyzer_test.go @@ -50,26 +50,56 @@ func TestAnalyzerCandidateRequiresBothDeadlinesOverdue(t *testing.T) { } } -func TestAnalyzerMarksHistoricalUncertainWhenRedisDeadlineCannotConfirm(t *testing.T) { +func TestAnalyzerMarksActiveTaskForDeadlineRepairWhenRedisDeadlineMissing(t *testing.T) { now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) taskID := uuid.New() - vm := &db.VirtualMachine{ID: "vm-uncertain", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ID: taskID, CreatedAt: now.Add(-2 * time.Hour)}}}} - - tests := []struct { - name string - deadlines *analyzerDeadlineStub - }{ - {name: "missing", deadlines: &analyzerDeadlineStub{}}, - {name: "future", deadlines: &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(time.Hour)}}}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - analyzer := newAnalyzerForTest(now, tt.deadlines) - got := analyzer.analyzeVM(context.Background(), vm) - if got.Decision != DecisionHistoricalUncertain { - t.Fatalf("decision = %s, reason = %s", got.Decision, got.Reason) - } - }) + vm := &db.VirtualMachine{ID: "vm-missing-deadline", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ + ID: taskID, Status: consts.TaskStatusProcessing, CreatedAt: now.Add(-2 * time.Hour), + }}}} + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionDeadlineMissing || !IsDeadlineRepairable(got.Decision) { + t.Fatalf("analysis = %+v", got) + } + if got.Target != "vm:"+vm.ID { + t.Fatalf("target = %q", got.Target) + } +} + +func TestAnalyzerTerminalTaskCandidateWhenRedisDeadlineMissing(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + vm := &db.VirtualMachine{ID: "vm-terminal", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ + ID: uuid.New(), Status: consts.TaskStatusFinished, CreatedAt: now.Add(-2 * time.Hour), LastActiveAt: now.Add(-2 * time.Hour), + }}}} + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionCandidate || !IsExecutable(got.Decision) { + t.Fatalf("analysis = %+v", got) + } +} + +func TestAnalyzerKeepsFutureRedisDeadlineHistoricalUncertain(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + vm := &db.VirtualMachine{ID: "vm-future", UserID: uuid.New(), Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ + ID: uuid.New(), Status: consts.TaskStatusProcessing, CreatedAt: now.Add(-2 * time.Hour), + }}}} + analyzer := newAnalyzerForTest(now, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(time.Hour)}}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionHistoricalUncertain { + t.Fatalf("analysis = %+v", got) + } +} + +func TestAnalyzerOrphanVMIsExplicitlyExecutable(t *testing.T) { + vm := &db.VirtualMachine{ID: "vm-orphan", UserID: uuid.New()} + analyzer := newAnalyzerForTest(time.Now(), &analyzerDeadlineStub{}) + + got := analyzer.analyzeVM(context.Background(), vm) + if got.Decision != DecisionOrphan || !IsExecutable(got.Decision) { + t.Fatalf("analysis = %+v", got) } } diff --git a/backend/pkg/vmrecycle/deadline_repair.go b/backend/pkg/vmrecycle/deadline_repair.go new file mode 100644 index 000000000..be684ef16 --- /dev/null +++ b/backend/pkg/vmrecycle/deadline_repair.go @@ -0,0 +1,122 @@ +package vmrecycle + +import ( + "context" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/samber/do" + + "github.com/chaitin/MonkeyCode/backend/config" + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/db/virtualmachine" + "github.com/chaitin/MonkeyCode/backend/domain" + "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" +) + +type DeadlineRepairResult struct { + VMID string + RecycleAt time.Time + Repaired bool +} + +type DeadlineRepairer interface { + Repair(ctx context.Context, vmID string) (DeadlineRepairResult, error) +} + +type deadlineQueue interface { + EnqueueIfMissing(ctx context.Context, queue string, payload *domain.VmIdleInfo, runAt time.Time, id string) (string, bool, error) + GetRunAt(ctx context.Context, queue, id string) (time.Time, bool, error) +} + +type deadlineRepairer struct { + cfg *config.Config + db *db.Client + teamPolicyRepo domain.TeamPolicyRepo + queue deadlineQueue + now func() time.Time +} + +func NewDeadlineRepairer(i *do.Injector) (DeadlineRepairer, error) { + return &deadlineRepairer{ + cfg: do.MustInvoke[*config.Config](i), + db: do.MustInvoke[*db.Client](i), + teamPolicyRepo: do.MustInvoke[domain.TeamPolicyRepo](i), + queue: do.MustInvoke[*delayqueue.VMRecycleQueue](i), + now: time.Now, + }, nil +} + +func (r *deadlineRepairer) Repair(ctx context.Context, vmID string) (DeadlineRepairResult, error) { + result := DeadlineRepairResult{VMID: vmID} + vm, err := r.db.VirtualMachine.Query().Where(virtualmachine.ID(vmID)).WithTasks().Only(ctx) + if err != nil { + return result, fmt.Errorf("query VM: %w", err) + } + if vm.IsRecycled { + return result, fmt.Errorf("VM already recycled") + } + if !hasNonTerminalTask(vm.Edges.Tasks) { + return result, fmt.Errorf("VM has no non-terminal task") + } + + policy, err := resolvePolicy(ctx, r.teamPolicyRepo, r.cfg.VMIdle, vm.UserID) + if err != nil { + return result, fmt.Errorf("resolve recycle policy: %w", err) + } + if !policy.RecycleEnabled { + return result, fmt.Errorf("recycle disabled") + } + + recycleAt := r.now().Add(time.Duration(policy.EffectiveRecycleSeconds) * time.Second) + payload := &domain.VmIdleInfo{ + UID: vm.UserID, + VmID: vm.ID, + HostID: vm.HostID, + EnvID: vm.EnvironmentID, + RecycleAt: recycleAt, + } + _, added, err := r.queue.EnqueueIfMissing(ctx, RecycleQueueKey, payload, recycleAt, vm.ID) + if err != nil { + return result, fmt.Errorf("enqueue recycle deadline: %w", err) + } + result.Repaired = added + if added { + result.RecycleAt = recycleAt + return result, nil + } + + existingAt, ok, err := r.queue.GetRunAt(ctx, RecycleQueueKey, vm.ID) + if err != nil { + return result, fmt.Errorf("query existing recycle deadline: %w", err) + } + if ok { + result.RecycleAt = existingAt + } + return result, nil +} + +func hasNonTerminalTask(tasks []*db.Task) bool { + for _, task := range tasks { + if task != nil && task.Status != consts.TaskStatusFinished && task.Status != consts.TaskStatusError { + return true + } + } + return false +} + +func resolvePolicy(ctx context.Context, repo domain.TeamPolicyRepo, cfg config.VMIdle, userID uuid.UUID) (*domain.TeamTaskVMIdlePolicy, error) { + team, err := repo.GetTeamByUserID(ctx, userID) + if err != nil && !db.IsNotFound(err) { + return nil, err + } + if db.IsNotFound(err) { + team = nil + } + return domain.ResolveTeamTaskVMIdlePolicy(team, cfg) +} + +var _ DeadlineRepairer = (*deadlineRepairer)(nil) +var _ deadlineQueue = (*delayqueue.VMRecycleQueue)(nil) diff --git a/backend/pkg/vmrecycle/deadline_repair_test.go b/backend/pkg/vmrecycle/deadline_repair_test.go new file mode 100644 index 000000000..f9052752a --- /dev/null +++ b/backend/pkg/vmrecycle/deadline_repair_test.go @@ -0,0 +1,141 @@ +package vmrecycle + +import ( + "context" + "fmt" + "io" + "log/slog" + "strings" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/chaitin/MonkeyCode/backend/config" + "github.com/chaitin/MonkeyCode/backend/consts" + "github.com/chaitin/MonkeyCode/backend/db" + "github.com/chaitin/MonkeyCode/backend/db/enttest" + "github.com/chaitin/MonkeyCode/backend/domain" + "github.com/chaitin/MonkeyCode/backend/pkg/delayqueue" + _ "github.com/mattn/go-sqlite3" +) + +func TestDeadlineRepairerCreatesDeadlineFromCurrentPolicy(t *testing.T) { + ctx := context.Background() + now := time.Date(2026, 7, 14, 12, 0, 0, 0, time.UTC) + client, vm := createDeadlineRepairVM(t, consts.TaskStatusProcessing, false) + queue := delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.New(slog.NewTextHandler(io.Discard, nil))) + repairer := newDeadlineRepairerForTest(client, queue, now, true) + + result, err := repairer.Repair(ctx, vm.ID) + if err != nil { + t.Fatal(err) + } + wantRecycleAt := now.Add(24 * time.Hour) + if !result.Repaired || !result.RecycleAt.Equal(wantRecycleAt) { + t.Fatalf("result = %+v", result) + } + job, runAt, ok, err := queue.GetJobInfo(ctx, RecycleQueueKey, vm.ID) + if err != nil || !ok { + t.Fatalf("GetJobInfo() ok = %v, err = %v", ok, err) + } + if job.Payload.VmID != vm.ID || job.Payload.UID != vm.UserID || !job.Payload.RecycleAt.Equal(wantRecycleAt) || !runAt.Equal(wantRecycleAt) { + t.Fatalf("job = %+v, run at = %v", job, runAt) + } +} + +func TestDeadlineRepairerDoesNotOverwriteConcurrentDeadline(t *testing.T) { + ctx := context.Background() + now := time.Date(2026, 7, 14, 12, 0, 0, 0, time.UTC) + client, vm := createDeadlineRepairVM(t, consts.TaskStatusPending, false) + queue := delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.New(slog.NewTextHandler(io.Discard, nil))) + existingAt := now.Add(2 * time.Hour) + if _, err := queue.Enqueue(ctx, RecycleQueueKey, &domain.VmIdleInfo{VmID: vm.ID, RecycleAt: existingAt}, existingAt, vm.ID); err != nil { + t.Fatal(err) + } + repairer := newDeadlineRepairerForTest(client, queue, now, true) + + result, err := repairer.Repair(ctx, vm.ID) + if err != nil { + t.Fatal(err) + } + if result.Repaired || !result.RecycleAt.Equal(existingAt) { + t.Fatalf("result = %+v", result) + } + job, runAt, ok, err := queue.GetJobInfo(ctx, RecycleQueueKey, vm.ID) + if err != nil || !ok || !runAt.Equal(existingAt) || !job.Payload.RecycleAt.Equal(existingAt) { + t.Fatalf("job = %+v, run at = %v, ok = %v, err = %v", job, runAt, ok, err) + } +} + +func TestDeadlineRepairerRevalidatesVMAndPolicy(t *testing.T) { + tests := []struct { + name string + status consts.TaskStatus + recycled bool + recycleEnabled bool + want string + }{ + {name: "terminal task", status: consts.TaskStatusFinished, recycleEnabled: true, want: "no non-terminal task"}, + {name: "recycled VM", status: consts.TaskStatusProcessing, recycled: true, recycleEnabled: true, want: "already recycled"}, + {name: "disabled policy", status: consts.TaskStatusProcessing, recycleEnabled: false, want: "recycle disabled"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + client, vm := createDeadlineRepairVM(t, tt.status, tt.recycled) + queue := delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.New(slog.NewTextHandler(io.Discard, nil))) + repairer := newDeadlineRepairerForTest(client, queue, time.Now(), tt.recycleEnabled) + + _, err := repairer.Repair(context.Background(), vm.ID) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("error = %v, want containing %q", err, tt.want) + } + if _, ok, queueErr := queue.GetRunAt(context.Background(), RecycleQueueKey, vm.ID); queueErr != nil || ok { + t.Fatalf("deadline ok = %v, error = %v", ok, queueErr) + } + }) + } +} + +func createDeadlineRepairVM(t *testing.T, status consts.TaskStatus, recycled bool) (*db.Client, *db.VirtualMachine) { + t.Helper() + ctx := context.Background() + client := enttest.Open(t, "sqlite3", fmt.Sprintf("file:deadline-repair-%s?mode=memory&cache=shared&_fk=1", uuid.NewString())) + t.Cleanup(func() { _ = client.Close() }) + userID := uuid.New() + if _, err := client.User.Create().SetID(userID).SetName("tester").SetRole(consts.UserRoleIndividual).SetStatus(consts.UserStatusActive).Save(ctx); err != nil { + t.Fatal(err) + } + if _, err := client.Host.Create().SetID("host-1").SetUserID(userID).SetHostname("host").Save(ctx); err != nil { + t.Fatal(err) + } + vm, err := client.VirtualMachine.Create().SetID("vm-1").SetHostID("host-1").SetUserID(userID).SetName("vm").SetIsRecycled(recycled).Save(ctx) + if err != nil { + t.Fatal(err) + } + task, err := client.Task.Create().SetID(uuid.New()).SetUserID(userID).SetKind(consts.TaskTypeDevelop).SetContent("content").SetStatus(status).Save(ctx) + if err != nil { + t.Fatal(err) + } + if _, err := client.TaskVirtualMachine.Create().SetID(uuid.New()).SetTaskID(task.ID).SetVirtualmachineID(vm.ID).Save(ctx); err != nil { + t.Fatal(err) + } + return client, vm +} + +func newDeadlineRepairerForTest(client *db.Client, queue deadlineQueue, now time.Time, recycleEnabled bool) *deadlineRepairer { + return &deadlineRepairer{ + cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 86400}}, + db: client, + teamPolicyRepo: &analyzerTeamPolicyRepoStub{team: &db.Team{ + ID: uuid.New(), + TaskConcurrencyLimit: 3, + TaskVMSleepEnabled: true, + TaskVMSleepSeconds: 600, + TaskVMRecycleEnabled: recycleEnabled, + TaskVMRecycleSeconds: 0, + }}, + queue: queue, + now: func() time.Time { return now }, + } +} From ed05bd787eadea0da18ebd8d07129991866465f6 Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 16:19:32 +0800 Subject: [PATCH 07/10] =?UTF-8?q?=E6=89=AB=E6=8F=8F=E6=97=A0=E5=85=B3?= =?UTF-8?q?=E8=81=94=E4=BB=BB=E5=8A=A1=E7=9A=84=E5=AD=A4=E5=84=BF=E8=99=9A?= =?UTF-8?q?=E6=8B=9F=E6=9C=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/pkg/vmrecycle/analyzer.go | 2 +- backend/pkg/vmrecycle/analyzer_test.go | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/backend/pkg/vmrecycle/analyzer.go b/backend/pkg/vmrecycle/analyzer.go index 51d399953..c51b67b8a 100644 --- a/backend/pkg/vmrecycle/analyzer.go +++ b/backend/pkg/vmrecycle/analyzer.go @@ -99,7 +99,7 @@ func (a *analyzer) Scan(ctx context.Context) (ScanReport, error) { cursor := "" for { query := a.db.VirtualMachine.Query(). - Where(virtualmachine.IsRecycled(false), virtualmachine.HasTasks()). + Where(virtualmachine.IsRecycled(false)). WithTasks(). Order(db.Asc(virtualmachine.FieldID)). Limit(analyzePageSize) diff --git a/backend/pkg/vmrecycle/analyzer_test.go b/backend/pkg/vmrecycle/analyzer_test.go index 77845b8ac..dd5b44826 100644 --- a/backend/pkg/vmrecycle/analyzer_test.go +++ b/backend/pkg/vmrecycle/analyzer_test.go @@ -103,6 +103,32 @@ func TestAnalyzerOrphanVMIsExplicitlyExecutable(t *testing.T) { } } +func TestAnalyzerScanReportsOrphanVM(t *testing.T) { + ctx := context.Background() + client := enttest.Open(t, "sqlite3", fmt.Sprintf("file:vm-recycle-orphan-%s?mode=memory&cache=shared&_fk=1", uuid.NewString())) + t.Cleanup(func() { _ = client.Close() }) + userID := uuid.New() + if _, err := client.User.Create().SetID(userID).SetName("tester").SetRole(consts.UserRoleIndividual).SetStatus(consts.UserStatusActive).Save(ctx); err != nil { + t.Fatal(err) + } + if _, err := client.Host.Create().SetID("host-orphan").SetUserID(userID).SetHostname("host").Save(ctx); err != nil { + t.Fatal(err) + } + if _, err := client.VirtualMachine.Create().SetID("vm-orphan").SetHostID("host-orphan").SetUserID(userID).SetName("vm").SetIsRecycled(false).Save(ctx); err != nil { + t.Fatal(err) + } + analyzer := newAnalyzerForTest(time.Now(), &analyzerDeadlineStub{}) + analyzer.db = client + + report, err := analyzer.Scan(ctx) + if err != nil { + t.Fatal(err) + } + if report.Counts[DecisionOrphan] != 1 || len(report.Items) != 1 || report.Items[0].Decision != DecisionOrphan { + t.Fatalf("report = %+v", report) + } +} + func TestAnalyzerUsesLatestActivityAcrossTasks(t *testing.T) { now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) oldTaskID := uuid.New() From 83d4afe494aa954a7f90f64558692315411bd225 Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 18:23:24 +0800 Subject: [PATCH 08/10] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E7=A9=BA=E9=97=B2?= =?UTF-8?q?=E5=9B=9E=E6=94=B6=E5=A4=B1=E8=B4=A5=E5=90=8E=E7=9A=84=E9=87=8D?= =?UTF-8?q?=E8=BF=9E=E8=A1=A5=E5=81=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/biz/vmidle/usecase/policy_test.go | 135 +++++++++++++++++++++- backend/biz/vmidle/usecase/vmidle.go | 48 ++++++-- backend/domain/host.go | 2 + backend/pkg/delayqueue/delayqueue.go | 14 ++- backend/pkg/delayqueue/delayqueue_test.go | 46 ++++++++ backend/pkg/vmrecycle/vmrecycle.go | 7 +- backend/pkg/vmrecycle/vmrecycle_test.go | 3 + 7 files changed, 239 insertions(+), 16 deletions(-) diff --git a/backend/biz/vmidle/usecase/policy_test.go b/backend/biz/vmidle/usecase/policy_test.go index 61afcf102..50964b793 100644 --- a/backend/biz/vmidle/usecase/policy_test.go +++ b/backend/biz/vmidle/usecase/policy_test.go @@ -230,6 +230,109 @@ func TestKeepAwakeWithSleepDisabledOnlyRemovesSleepSchedule(t *testing.T) { assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, vmrecycle.RecycleQueueKey, vmID, oldRecycleAt) } +func TestRecycleJobMarksVMRecycledAfterFinalRemoteDeleteFailure(t *testing.T) { + repo := &refreshHostRepoStub{} + r := &vmIdleRefresher{ + logger: slog.Default(), + hostRepo: repo, + recycleQueue: delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.Default()), + recycler: &idleRecyclerStub{err: fmt.Errorf("%w: taskflow unavailable", vmrecycle.ErrRemoteDelete)}, + } + job := &delayqueue.Job[*domain.VmIdleInfo]{ + Payload: &domain.VmIdleInfo{VmID: "vm-final"}, + Attempts: 4, + } + + if err := r.handleRecycleJob(context.Background(), job); err != nil { + t.Fatal(err) + } + if repo.updateVirtualMachineCalls != 1 || !repo.isRecycled { + t.Fatalf("update calls = %d, is recycled = %v", repo.updateVirtualMachineCalls, repo.isRecycled) + } +} + +func TestRecycleJobDoesNotMarkVMBeforeFinalAttempt(t *testing.T) { + wantErr := fmt.Errorf("%w: taskflow unavailable", vmrecycle.ErrRemoteDelete) + repo := &refreshHostRepoStub{} + r := &vmIdleRefresher{ + logger: slog.Default(), + hostRepo: repo, + recycleQueue: delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.Default()), + recycler: &idleRecyclerStub{err: wantErr}, + } + job := &delayqueue.Job[*domain.VmIdleInfo]{ + Payload: &domain.VmIdleInfo{VmID: "vm-retry"}, + Attempts: 3, + } + + err := r.handleRecycleJob(context.Background(), job) + if !errors.Is(err, wantErr) { + t.Fatalf("error = %v, want %v", err, wantErr) + } + if repo.updateVirtualMachineCalls != 0 { + t.Fatalf("update calls = %d, want 0", repo.updateVirtualMachineCalls) + } +} + +func TestRecycleJobDoesNotMarkVMForNonRemoteFailure(t *testing.T) { + wantErr := errors.New("database unavailable") + repo := &refreshHostRepoStub{} + r := &vmIdleRefresher{ + logger: slog.Default(), + hostRepo: repo, + recycleQueue: delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.Default()), + recycler: &idleRecyclerStub{err: wantErr}, + } + job := &delayqueue.Job[*domain.VmIdleInfo]{ + Payload: &domain.VmIdleInfo{VmID: "vm-database-error"}, + Attempts: 4, + } + + err := r.handleRecycleJob(context.Background(), job) + if !errors.Is(err, wantErr) { + t.Fatalf("error = %v, want %v", err, wantErr) + } + if repo.updateVirtualMachineCalls != 0 { + t.Fatalf("update calls = %d, want 0", repo.updateVirtualMachineCalls) + } +} + +func TestRecycleJobRetriesMarkWithoutDeletingRemotelyAgain(t *testing.T) { + remoteErr := fmt.Errorf("%w: taskflow unavailable", vmrecycle.ErrRemoteDelete) + markErr := errors.New("database unavailable") + repo := &refreshHostRepoStub{updateVirtualMachineErr: markErr} + recycler := &idleRecyclerStub{err: remoteErr} + r := &vmIdleRefresher{ + logger: slog.Default(), + hostRepo: repo, + recycleQueue: delayqueue.NewVMRecycleQueue(newTestRedis(t), slog.Default()), + recycler: recycler, + } + job := &delayqueue.Job[*domain.VmIdleInfo]{ + Payload: &domain.VmIdleInfo{VmID: "vm-mark-retry"}, + Attempts: 4, + } + + err := r.handleRecycleJob(context.Background(), job) + if !errors.Is(err, remoteErr) || !errors.Is(err, markErr) || !errors.Is(err, delayqueue.ErrRetryAfterMaxAttempts) { + t.Fatalf("error = %v, want remote, database and persistent retry errors", err) + } + if !job.Payload.RecycleDeleteExhausted || recycler.calls != 1 { + t.Fatalf("delete exhausted = %v, recycle calls = %d", job.Payload.RecycleDeleteExhausted, recycler.calls) + } + + repo.updateVirtualMachineErr = nil + if err := r.handleRecycleJob(context.Background(), job); err != nil { + t.Fatal(err) + } + if recycler.calls != 1 { + t.Fatalf("recycle calls = %d, want 1", recycler.calls) + } + if repo.updateVirtualMachineCalls != 2 || !repo.isRecycled { + t.Fatalf("update calls = %d, is recycled = %v", repo.updateVirtualMachineCalls, repo.isRecycled) + } +} + func newQueueTestVMIdleRefresher(redisClient *redis.Client, repo domain.HostRepo) *vmIdleRefresher { logger := slog.Default() return &vmIdleRefresher{ @@ -293,9 +396,22 @@ func newVirtualMachineNotFoundErr(t *testing.T) error { } type refreshHostRepoStub struct { - vm *db.VirtualMachine - err error - getVirtualMachineCalls int + vm *db.VirtualMachine + err error + updateVirtualMachineErr error + getVirtualMachineCalls int + updateVirtualMachineCalls int + isRecycled bool +} + +type idleRecyclerStub struct { + err error + calls int +} + +func (s *idleRecyclerStub) Recycle(_ context.Context, vmID string) (vmrecycle.Result, error) { + s.calls++ + return vmrecycle.Result{VMID: vmID}, s.err } type refreshTeamPolicyRepoStub struct { @@ -383,8 +499,17 @@ func (s *refreshHostRepoStub) UpsertVirtualMachine(context.Context, *taskflow.Vi return errors.New("not implemented") } -func (s *refreshHostRepoStub) UpdateVirtualMachine(context.Context, string, func(*db.VirtualMachineUpdateOne) error) error { - return errors.New("not implemented") +func (s *refreshHostRepoStub) UpdateVirtualMachine(_ context.Context, id string, fn func(*db.VirtualMachineUpdateOne) error) error { + s.updateVirtualMachineCalls++ + if s.updateVirtualMachineErr != nil { + return s.updateVirtualMachineErr + } + up := db.NewClient().VirtualMachine.UpdateOneID(id) + if err := fn(up); err != nil { + return err + } + s.isRecycled, _ = up.Mutation().IsRecycled() + return nil } func (s *refreshHostRepoStub) UpsertHost(context.Context, *taskflow.Host) error { diff --git a/backend/biz/vmidle/usecase/vmidle.go b/backend/biz/vmidle/usecase/vmidle.go index 5da275928..9e5c2a84f 100644 --- a/backend/biz/vmidle/usecase/vmidle.go +++ b/backend/biz/vmidle/usecase/vmidle.go @@ -457,19 +457,51 @@ func (r *vmIdleRefresher) lookupSchedule(jobID string) (notifySchedule, bool) { } func (r *vmIdleRefresher) recycleConsumer() { - logger := r.logger.With("fn", "recycleConsumer") for { - err := r.recycleQueue.StartConsumer(context.Background(), vmrecycle.RecycleQueueKey, - func(ctx context.Context, job *delayqueue.Job[*domain.VmIdleInfo]) error { - logger.InfoContext(ctx, "vm recycle triggered", "vmID", job.Payload.VmID) - _, err := r.recycler.Recycle(ctx, job.Payload.VmID) - return err - }) - logger.Warn("recycle consumer error, retrying...", "error", err) + err := r.recycleQueue.StartConsumer(context.Background(), vmrecycle.RecycleQueueKey, r.handleRecycleJob) + r.logger.With("fn", "recycleConsumer").Warn("recycle consumer error, retrying...", "error", err) time.Sleep(10 * time.Second) } } +func (r *vmIdleRefresher) handleRecycleJob(ctx context.Context, job *delayqueue.Job[*domain.VmIdleInfo]) error { + vmID := job.Payload.VmID + logger := r.logger.With("fn", "recycleConsumer", "vm_id", vmID) + if job.Payload.RecycleDeleteExhausted { + return r.markRecycledForReconnectCompensation(ctx, job, nil) + } + logger.InfoContext(ctx, "vm recycle triggered", "attempt", job.Attempts+1) + _, err := r.recycler.Recycle(ctx, vmID) + if err == nil { + return nil + } + if !errors.Is(err, vmrecycle.ErrRemoteDelete) || !r.recycleQueue.IsFinalAttempt(job) { + return err + } + job.Payload.RecycleDeleteExhausted = true + return r.markRecycledForReconnectCompensation(ctx, job, err) +} + +func (r *vmIdleRefresher) markRecycledForReconnectCompensation(ctx context.Context, job *delayqueue.Job[*domain.VmIdleInfo], remoteErr error) error { + vmID := job.Payload.VmID + if err := r.hostRepo.UpdateVirtualMachine(ctx, vmID, func(up *db.VirtualMachineUpdateOne) error { + up.SetIsRecycled(true) + return nil + }); err != nil { + return errors.Join( + delayqueue.ErrRetryAfterMaxAttempts, + remoteErr, + fmt.Errorf("mark vm %s recycled after remote delete retries exhausted: %w", vmID, err), + ) + } + r.logger.With("fn", "recycleConsumer", "vm_id", vmID).WarnContext(ctx, + "vm remote delete retries exhausted, marked recycled for agent reconnect compensation", + "attempts", job.Attempts+1, + "error", remoteErr, + ) + return nil +} + func (r *vmIdleRefresher) buildRecycleNotifyEvent(ctx context.Context, vm *db.VirtualMachine, expiresAt time.Time) (*domain.NotifyEvent, error) { // 直接按 virtualmachine_id 查 task_virtualmachines 拿 task_id， // 不再依赖 vm.Edges.Tasks 的 eager load 链。 diff --git a/backend/domain/host.go b/backend/domain/host.go index 663f66fe5..8438e5f9d 100644 --- a/backend/domain/host.go +++ b/backend/domain/host.go @@ -89,6 +89,8 @@ type VmIdleInfo struct { EnvID string `json:"env_id"` TaskID string `json:"task_id,omitempty"` // 关联的任务 ID，用于通知 Name string `json:"name,omitempty"` // 任务名称，用于通知内容 + // RecycleDeleteExhausted 表示远端删除已重试耗尽，后续消费只补写回收状态。 + RecycleDeleteExhausted bool `json:"recycle_delete_exhausted,omitempty"` // RecycleAt 是本次 RecordActivity 算出的预计回收时间。每次用户活动都会延长这个值， // consumer 把它编进 RefID，让每个回收窗口都能产生不同的 dedup key（否则 // dispatcher 会按 (subID, eventType, RefID) 把同一 task 的后续推送全部静默）。 diff --git a/backend/pkg/delayqueue/delayqueue.go b/backend/pkg/delayqueue/delayqueue.go index 34216dd85..c3f010fc5 100644 --- a/backend/pkg/delayqueue/delayqueue.go +++ b/backend/pkg/delayqueue/delayqueue.go @@ -47,6 +47,9 @@ const ( defaultJobTTL = 7 * 24 * time.Hour ) +// ErrRetryAfterMaxAttempts 让已耗尽普通尝试次数的任务继续保留并重试。 +var ErrRetryAfterMaxAttempts = errors.New("retry delay queue job after max attempts") + // NewRedisDelayQueue 创建延迟队列（泛型） func NewRedisDelayQueue[T any](rdb *redis.Client, logger *slog.Logger, opts ...Option[T]) *RedisDelayQueue[T] { q := &RedisDelayQueue[T]{ @@ -89,6 +92,11 @@ func WithJobTTL[T any](d time.Duration) Option[T] { return func(q *RedisDelayQueue[T]) { q.jobTTL = d } } +// IsFinalAttempt 判断当前处理是否为队列允许的最后一次尝试。 +func (q *RedisDelayQueue[T]) IsFinalAttempt(job *Job[T]) bool { + return job != nil && job.Attempts+1 >= q.maxAttempts +} + // Enqueue 入队：始终写入/覆盖 payload 并更新到期时间。 // 这样即使 pollOnce 在 handleJob 完成后删除了旧 job key， // 新的 Enqueue 调用也能重建 job 数据，避免 ZSet 中存在孤立条目。 @@ -233,10 +241,14 @@ func (q *RedisDelayQueue[T]) pollOnce(ctx context.Context, queue string, handler if err := handler(ctx, job); err != nil { q.logger.Warn("handle job failed, will requeue", "queue", queue, "id", id, "attempts", job.Attempts+1, "err", err) job.Attempts++ - if job.Attempts >= q.maxAttempts { + retryAfterMaxAttempts := errors.Is(err, ErrRetryAfterMaxAttempts) + if job.Attempts >= q.maxAttempts && !retryAfterMaxAttempts { _ = q.deleteJob(ctx, queue, id) continue } + if retryAfterMaxAttempts && job.Attempts >= q.maxAttempts { + job.Attempts = max(q.maxAttempts-1, 0) + } if err := q.saveJob(ctx, queue, id, job.Attempts, job.Payload); err != nil { q.logger.Error("save job failed when requeue", "queue", queue, "id", id, "err", err) continue diff --git a/backend/pkg/delayqueue/delayqueue_test.go b/backend/pkg/delayqueue/delayqueue_test.go index fa0406388..494c2bf9c 100644 --- a/backend/pkg/delayqueue/delayqueue_test.go +++ b/backend/pkg/delayqueue/delayqueue_test.go @@ -96,3 +96,49 @@ func TestEnqueueIfMissingDoesNotOverwriteExistingJob(t *testing.T) { t.Fatalf("job = %+v, run at = %v", job, gotRunAt) } } + +func TestIsFinalAttemptUsesConfiguredMaxAttempts(t *testing.T) { + queue := NewRedisDelayQueue[string](nil, slog.Default(), WithMaxAttempts[string](3)) + if queue.IsFinalAttempt(&Job[string]{Attempts: 1}) { + t.Fatal("second attempt must not be final") + } + if !queue.IsFinalAttempt(&Job[string]{Attempts: 2}) { + t.Fatal("third attempt must be final") + } +} + +func TestPollOncePreservesJobForRetryAfterMaxAttempts(t *testing.T) { + ctx := context.Background() + srv := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: srv.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + queue := NewRedisDelayQueue[string]( + rdb, + slog.Default(), + WithMaxAttempts[string](1), + WithRequeueDelay[string](0), + ) + if _, err := queue.Enqueue(ctx, "recycle", "payload", time.Now(), "vm-1"); err != nil { + t.Fatal(err) + } + + if err := queue.pollOnce(ctx, "recycle", func(context.Context, *Job[string]) error { + return ErrRetryAfterMaxAttempts + }); err != nil { + t.Fatal(err) + } + job, _, ok, err := queue.GetJobInfo(ctx, "recycle", "vm-1") + if err != nil || !ok { + t.Fatalf("job preserved = %v, err = %v", ok, err) + } + if job.Attempts != 0 { + t.Fatalf("attempts = %d, want 0", job.Attempts) + } + + if err := queue.pollOnce(ctx, "recycle", func(context.Context, *Job[string]) error { return nil }); err != nil { + t.Fatal(err) + } + if _, _, ok, err := queue.GetJobInfo(ctx, "recycle", "vm-1"); err != nil || ok { + t.Fatalf("job preserved = %v, err = %v, want removed", ok, err) + } +} diff --git a/backend/pkg/vmrecycle/vmrecycle.go b/backend/pkg/vmrecycle/vmrecycle.go index 461548a3d..b88fed43d 100644 --- a/backend/pkg/vmrecycle/vmrecycle.go +++ b/backend/pkg/vmrecycle/vmrecycle.go @@ -37,7 +37,10 @@ const ( StatusNotFound Status = "not_found" ) -var ErrInProgress = errors.New("vm recycle in progress") +var ( + ErrInProgress = errors.New("vm recycle in progress") + ErrRemoteDelete = errors.New("vm remote delete failed") +) type Result struct { VMID string @@ -113,7 +116,7 @@ func (r *recycler) Recycle(ctx context.Context, vmID string) (Result, error) { HostID: vm.HostID, ID: vm.EnvironmentID, }); err != nil { - return result, fmt.Errorf("delete vm %s: %w", vmID, err) + return result, fmt.Errorf("%w: delete vm %s: %w", ErrRemoteDelete, vmID, err) } if err := r.hostRepo.UpdateVirtualMachine(ctx, vmID, func(up *db.VirtualMachineUpdateOne) error { up.SetIsRecycled(true) diff --git a/backend/pkg/vmrecycle/vmrecycle_test.go b/backend/pkg/vmrecycle/vmrecycle_test.go index c9cf475b2..d2ca07bc6 100644 --- a/backend/pkg/vmrecycle/vmrecycle_test.go +++ b/backend/pkg/vmrecycle/vmrecycle_test.go @@ -142,6 +142,9 @@ func TestRecyclerRemoteFailureDoesNotMarkOrClean(t *testing.T) { if !errors.Is(err, wantErr) { t.Fatalf("error = %v, want %v", err, wantErr) } + if !errors.Is(err, ErrRemoteDelete) { + t.Fatalf("error = %v, want ErrRemoteDelete", err) + } if result.Status != "" || hostRepo.updateCalls != 0 || len(queues.removed) != 0 { t.Fatalf("result = %+v, update calls = %d, cleanup = %v", result, hostRepo.updateCalls, queues.removed) } From cbdfc78947bc8d4cc53fd635b6e511a3929f18ec Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 18:57:46 +0800 Subject: [PATCH 09/10] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E8=BF=9C=E7=AB=AF?= =?UTF-8?q?=E8=99=9A=E6=8B=9F=E6=9C=BA=E4=B8=8D=E5=AD=98=E5=9C=A8=E6=97=B6?= =?UTF-8?q?=E5=9B=9E=E6=94=B6=E4=B8=AD=E6=96=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/pkg/taskflow/vm.go | 6 ++++++ backend/pkg/taskflow/vm_test.go | 20 ++++++++++++++++++++ backend/pkg/vmrecycle/vmrecycle.go | 2 +- backend/pkg/vmrecycle/vmrecycle_test.go | 20 ++++++++++++++++++++ 4 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 backend/pkg/taskflow/vm_test.go diff --git a/backend/pkg/taskflow/vm.go b/backend/pkg/taskflow/vm.go index bbc6526d4..6f077a2bb 100644 --- a/backend/pkg/taskflow/vm.go +++ b/backend/pkg/taskflow/vm.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/url" + "strings" "github.com/coder/websocket" @@ -38,6 +39,11 @@ func (v *virtualMachineClient) Delete(ctx context.Context, req *DeleteVirtualMac return err } +// IsVirtualMachineNotFound 判断删除失败是否表示远端环境已经不存在。 +func IsVirtualMachineNotFound(err error) bool { + return err != nil && strings.Contains(err.Error(), "environment not found:") +} + // List implements VirtualMachiner. func (v *virtualMachineClient) List(ctx context.Context, id string) ([]*VirtualMachine, error) { q := request.Query{ diff --git a/backend/pkg/taskflow/vm_test.go b/backend/pkg/taskflow/vm_test.go new file mode 100644 index 000000000..249ecb4a5 --- /dev/null +++ b/backend/pkg/taskflow/vm_test.go @@ -0,0 +1,20 @@ +package taskflow + +import ( + "errors" + "fmt" + "testing" +) + +func TestIsVirtualMachineNotFound(t *testing.T) { + notFound := errors.New("recv err failed to stop environment: environment not found: env-1") + if !IsVirtualMachineNotFound(fmt.Errorf("delete vm: %w", notFound)) { + t.Fatal("wrapped environment-not-found error must be recognized") + } + if IsVirtualMachineNotFound(errors.New("failed to stop environment: connection refused")) { + t.Fatal("unrelated delete error must remain a failure") + } + if IsVirtualMachineNotFound(nil) { + t.Fatal("nil must not be recognized as not found") + } +} diff --git a/backend/pkg/vmrecycle/vmrecycle.go b/backend/pkg/vmrecycle/vmrecycle.go index b88fed43d..bb5a0f53c 100644 --- a/backend/pkg/vmrecycle/vmrecycle.go +++ b/backend/pkg/vmrecycle/vmrecycle.go @@ -115,7 +115,7 @@ func (r *recycler) Recycle(ctx context.Context, vmID string) (Result, error) { UserID: vm.UserID.String(), HostID: vm.HostID, ID: vm.EnvironmentID, - }); err != nil { + }); err != nil && !taskflow.IsVirtualMachineNotFound(err) { return result, fmt.Errorf("%w: delete vm %s: %w", ErrRemoteDelete, vmID, err) } if err := r.hostRepo.UpdateVirtualMachine(ctx, vmID, func(up *db.VirtualMachineUpdateOne) error { diff --git a/backend/pkg/vmrecycle/vmrecycle_test.go b/backend/pkg/vmrecycle/vmrecycle_test.go index d2ca07bc6..9c3a7d379 100644 --- a/backend/pkg/vmrecycle/vmrecycle_test.go +++ b/backend/pkg/vmrecycle/vmrecycle_test.go @@ -150,6 +150,26 @@ func TestRecyclerRemoteFailureDoesNotMarkOrClean(t *testing.T) { } } +func TestRecyclerTreatsMissingRemoteEnvironmentAsDeleted(t *testing.T) { + vm := &db.VirtualMachine{ID: "vm-remote-missing", UserID: uuid.New()} + r, hostRepo, taskRepo, vmClient := newStubRecycler(t, vm) + vmClient.err = errors.New("recv err failed to stop environment: environment not found: env-1") + + result, err := r.Recycle(context.Background(), vm.ID) + if err != nil { + t.Fatal(err) + } + if result.Status != StatusRecycled || hostRepo.updateCalls != 1 || !vm.IsRecycled { + t.Fatalf("result = %+v, update calls = %d, recycled = %v", result, hostRepo.updateCalls, vm.IsRecycled) + } + if vmClient.deleteCalls != 1 || len(taskRepo.updated) != 0 { + t.Fatalf("delete calls = %d, updated tasks = %v", vmClient.deleteCalls, taskRepo.updated) + } + if !r.recycleQueue.(*recycleQueueStub).removedContains(RecycleQueueKey, vm.ID) { + t.Fatal("recycle cleanup must continue when the remote environment is already absent") + } +} + func TestRecyclerRetriesCleanupWithoutDeletingRemoteAgain(t *testing.T) { vm := &db.VirtualMachine{ID: "vm-cleanup-retry", UserID: uuid.New()} r, hostRepo, _, vmClient := newStubRecycler(t, vm) From af969a3d86dd36deb44d4c2610e356dec7f382c4 Mon Sep 17 00:00:00 2001 From: yokowu <18836617@qq.com> Date: Tue, 14 Jul 2026 19:56:19 +0800 Subject: [PATCH 10/10] =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=BC=BA=E5=88=B6?= =?UTF-8?q?=E5=9B=9E=E6=94=B6=E5=BF=BD=E7=95=A5=E8=BF=9C=E7=AB=AF=E5=88=A0?= =?UTF-8?q?=E9=99=A4=E5=A4=B1=E8=B4=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/pkg/vmrecycle/vmrecycle.go | 23 ++++++++++++++--- backend/pkg/vmrecycle/vmrecycle_test.go | 33 +++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/backend/pkg/vmrecycle/vmrecycle.go b/backend/pkg/vmrecycle/vmrecycle.go index bb5a0f53c..a9d80adee 100644 --- a/backend/pkg/vmrecycle/vmrecycle.go +++ b/backend/pkg/vmrecycle/vmrecycle.go @@ -52,6 +52,11 @@ type Recycler interface { Recycle(ctx context.Context, vmID string) (Result, error) } +type ForceRecycler interface { + Recycler + ForceRecycle(ctx context.Context, vmID string) (Result, error) +} + type removableQueue interface { Remove(ctx context.Context, queue, id string) error } @@ -90,6 +95,14 @@ func NewRecycler(i *do.Injector) (Recycler, error) { } func (r *recycler) Recycle(ctx context.Context, vmID string) (Result, error) { + return r.recycle(ctx, vmID, false) +} + +func (r *recycler) ForceRecycle(ctx context.Context, vmID string) (Result, error) { + return r.recycle(ctx, vmID, true) +} + +func (r *recycler) recycle(ctx context.Context, vmID string, ignoreRemoteDeleteFailure bool) (Result, error) { result := Result{VMID: vmID} token, err := r.acquire(ctx, vmID) if err != nil { @@ -111,12 +124,16 @@ func (r *recycler) Recycle(ctx context.Context, vmID string) (Result, error) { if vm.IsRecycled { result.Status = StatusAlreadyRecycled } else { - if err := r.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{ + deleteErr := r.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{ UserID: vm.UserID.String(), HostID: vm.HostID, ID: vm.EnvironmentID, - }); err != nil && !taskflow.IsVirtualMachineNotFound(err) { - return result, fmt.Errorf("%w: delete vm %s: %w", ErrRemoteDelete, vmID, err) + }) + if deleteErr != nil && !taskflow.IsVirtualMachineNotFound(deleteErr) { + if !ignoreRemoteDeleteFailure { + return result, fmt.Errorf("%w: delete vm %s: %w", ErrRemoteDelete, vmID, deleteErr) + } + r.logger.WarnContext(ctx, "force recycle continuing after remote delete failure", "vm_id", vmID, "error", deleteErr) } if err := r.hostRepo.UpdateVirtualMachine(ctx, vmID, func(up *db.VirtualMachineUpdateOne) error { up.SetIsRecycled(true) diff --git a/backend/pkg/vmrecycle/vmrecycle_test.go b/backend/pkg/vmrecycle/vmrecycle_test.go index 9c3a7d379..1c87103fe 100644 --- a/backend/pkg/vmrecycle/vmrecycle_test.go +++ b/backend/pkg/vmrecycle/vmrecycle_test.go @@ -170,6 +170,39 @@ func TestRecyclerTreatsMissingRemoteEnvironmentAsDeleted(t *testing.T) { } } +func TestForceRecyclerContinuesAfterOpaqueRemoteDeleteFailure(t *testing.T) { + taskID := uuid.New() + vm := &db.VirtualMachine{ + ID: "vm-force-remote-fail", + UserID: uuid.New(), + Edges: db.VirtualMachineEdges{Tasks: []*db.Task{{ + ID: taskID, Status: consts.TaskStatusProcessing, + }}}, + } + r, hostRepo, taskRepo, vmClient := newStubRecycler(t, vm) + vmClient.err = errors.New(`HTTP 500: {"code":500,"message":"服务器错误 [trace_id: test]"}`) + forceRecycler, ok := any(r).(interface { + ForceRecycle(context.Context, string) (Result, error) + }) + if !ok { + t.Fatal("recycler must support force recycle") + } + + result, err := forceRecycler.ForceRecycle(context.Background(), vm.ID) + if err != nil { + t.Fatal(err) + } + if result.Status != StatusRecycled || hostRepo.updateCalls != 1 || !vm.IsRecycled { + t.Fatalf("result = %+v, update calls = %d, recycled = %v", result, hostRepo.updateCalls, vm.IsRecycled) + } + if vmClient.deleteCalls != 1 || len(taskRepo.updated) != 1 || taskRepo.updated[0] != taskID { + t.Fatalf("delete calls = %d, updated tasks = %v", vmClient.deleteCalls, taskRepo.updated) + } + if !r.recycleQueue.(*recycleQueueStub).removedContains(RecycleQueueKey, vm.ID) { + t.Fatal("force recycle must continue local cleanup after remote delete failure") + } +} + func TestRecyclerRetriesCleanupWithoutDeletingRemoteAgain(t *testing.T) { vm := &db.VirtualMachine{ID: "vm-cleanup-retry", UserID: uuid.New()} r, hostRepo, _, vmClient := newStubRecycler(t, vm)