diff --git a/backend/biz/host/handler/v1/internal.go b/backend/biz/host/handler/v1/internal.go index bb6ca4c27..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.Refresh(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 eca2875a4..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,14 +114,59 @@ 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 } -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: } 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 fbfd3c2db..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,13 +140,9 @@ 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.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 +168,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) } } @@ -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,30 +220,21 @@ 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 { +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(): 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: - 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 7e807df45..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 { @@ -66,7 +50,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,20 +58,4 @@ func (r *testVMIdleRefresher) Refresh(_ context.Context, vmID 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) +func (r *testVMIdleRefresher) RecordActivity(context.Context, string) error { return 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/biz/task/usecase/task.go b/backend/biz/task/usecase/task.go index a3ffe22bf..9008a5256 100644 --- a/backend/biz/task/usecase/task.go +++ b/backend/biz/task/usecase/task.go @@ -432,6 +432,12 @@ 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) + } // 缓存最近一次 user-input，供通知推送使用 a.redis.Set(ctx, fmt.Sprintf("mcai:task:%s:last_input", id.String()), string(req.Content), 24*time.Hour) @@ -721,7 +727,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..01e1c184a 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,97 @@ 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, 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) + } + 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) + } + 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, taskRefresher := 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") + } + 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, *taskActivityRefresherStub) { + 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) + 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, + taskActivityRefresher: taskRefresher, + }, user, taskID, taskRefresher +} + +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 +145,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..50964b793 100644 --- a/backend/biz/vmidle/usecase/policy_test.go +++ b/backend/biz/vmidle/usecase/policy_test.go @@ -14,35 +14,39 @@ 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 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,21 +80,23 @@ 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{err: newVirtualMachineNotFoundErr(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, @@ -98,112 +104,268 @@ 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.KeepAwake(ctx, "vm-activity"); err != nil { + t.Fatalf("keep awake: %v", err) } - if err := r.Refresh(ctx, "missing-vm"); err != nil { - t.Fatalf("second refresh: %v", err) + if err := r.RecordActivity(ctx, "vm-activity"); err != nil { + t.Fatalf("record activity: %v", err) } - if repo.getVirtualMachineCalls != 1 { - t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls) + 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 TestShouldSkipRecycleForRecentTaskLastActiveAtSkipsClickHouse(t *testing.T) { +func TestRecordActivityCachesNotFoundVM(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")} + repo := &refreshHostRepoStub{err: newVirtualMachineNotFoundErr(t)} 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), + cfg: &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 604800}}, + redis: redisClient, + logger: slog.Default(), + hostRepo: repo, } - 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 err := r.RecordActivity(ctx, "missing-vm"); err != nil { + t.Fatalf("first activity: %v", err) } - if activity.calls != 0 { - t.Fatalf("clickhouse calls = %d, want 0", activity.calls) + 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 TestShouldSkipRecycleForRecentTaskLogRefreshesVM(t *testing.T) { +func TestKeepAwakeOnlyUpdatesSleepSchedule(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} + 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, vmrecycle.NotifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { + t.Fatal(err) + } + if _, err := r.recycleQueue.Enqueue(ctx, vmrecycle.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, 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) { + 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, 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) { + 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, vmrecycle.SleepQueueKey, payload, oldSleepAt, vmID); err != nil { + t.Fatal(err) + } + if _, err := r.notifyQueue.Enqueue(ctx, vmrecycle.NotifyQueueKey, payload, oldNotifyAt, vmID+":default"); err != nil { + t.Fatal(err) + } + 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, vmrecycle.SleepQueueKey, vmID); err != nil || ok { + t.Fatalf("sleep job ok = %v, err = %v, want removed", ok, err) + } + assertVMIdleJobAt(t, r.notifyQueue.RedisDelayQueue, ctx, vmrecycle.NotifyQueueKey, vmID+":default", oldNotifyAt) + assertVMIdleJobAt(t, r.recycleQueue.RedisDelayQueue, ctx, vmrecycle.RecycleQueueKey, vmID, oldRecycleAt) +} + +func TestRecycleJobMarksVMRecycledAfterFinalRemoteDeleteFailure(t *testing.T) { + repo := &refreshHostRepoStub{} 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), + 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, } - skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now) - if err != nil { + if err := r.handleRecycleJob(context.Background(), job); err != nil { t.Fatal(err) } - if !skip { - t.Fatal("expected recent task log to skip recycle") + if repo.updateVirtualMachineCalls != 1 || !repo.isRecycled { + t.Fatalf("update calls = %d, is recycled = %v", repo.updateVirtualMachineCalls, repo.isRecycled) } - if repo.getVirtualMachineCalls != 1 { - t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls) +} + +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 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} +func TestRecycleJobDoesNotMarkVMForNonRemoteFailure(t *testing.T) { + wantErr := errors.New("database unavailable") + repo := &refreshHostRepoStub{} 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)}}, + 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, } - skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now) - if err != nil { + 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 skip { - t.Fatal("expected stale task log to allow recycle") + 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{ + 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 repo.getVirtualMachineCalls != 0 { - t.Fatalf("GetVirtualMachine calls = %d, want 0", repo.getVirtualMachineCalls) + 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) } } @@ -234,24 +396,42 @@ 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 taskLogActivityStub struct { - latest map[uuid.UUID]time.Time - err error - calls int +type idleRecyclerStub struct { + err error + calls int } -func (s *taskLogActivityStub) LatestTaskLogTime(_ context.Context, taskID uuid.UUID) (time.Time, bool, error) { +func (s *idleRecyclerStub) Recycle(_ context.Context, vmID string) (vmrecycle.Result, error) { s.calls++ - if s.err != nil { - return time.Time{}, false, s.err - } - latest, ok := s.latest[taskID] - return latest, ok, nil + return vmrecycle.Result{VMID: vmID}, s.err +} + +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") } func (s *refreshHostRepoStub) List(context.Context, uuid.UUID) ([]*db.Host, error) { @@ -319,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 5a4b0bdd8..9e5c2a84f 100644 --- a/backend/biz/vmidle/usecase/vmidle.go +++ b/backend/biz/vmidle/usecase/vmidle.go @@ -17,22 +17,18 @@ 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 { - Refresh(ctx context.Context, vmID string) error + KeepAwake(ctx context.Context, vmID string) error + RecordActivity(ctx context.Context, vmID string) error } const ( - sleepQueueKey = "vm:idle:sleep" - notifyQueueKey = "vm:idle:notify" - recycleQueueKey = "vm:idle:recycle" - vmIdleDebounceTTL = 30 * time.Second vmNotFoundTTL = 30 * time.Second ) @@ -105,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 @@ -129,16 +121,15 @@ 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) - if err != nil { - return nil, err - } - var activityReader taskLogActivityReader - if taskLogActivity != nil { - activityReader = taskLogActivity - } r := &vmIdleRefresher{ cfg: cfg, redis: do.MustInvoke[*redis.Client](i), @@ -147,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), } @@ -212,8 +203,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 +227,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 +243,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 +276,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,21 +288,34 @@ 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 { + 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) - 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 { + } else if err := r.sleepQueue.Remove(ctx, vmrecycle.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) 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)) } @@ -313,7 +324,7 @@ func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { 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 @@ -325,11 +336,11 @@ func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { "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)) } @@ -339,7 +350,7 @@ func (r *vmIdleRefresher) Refresh(ctx context.Context, vmID string) error { 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) @@ -371,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) @@ -446,179 +457,49 @@ 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, - 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 - }) - 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) 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 +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) } - - 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.skipRecycleAndRefresh(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.skipRecycleAndRefresh(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) { - 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) - } - 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) + logger.InfoContext(ctx, "vm recycle triggered", "attempt", job.Attempts+1) + _, err := r.recycler.Recycle(ctx, vmID) + if err == nil { + return nil } - if tk.Status == consts.TaskStatusFinished || tk.Status == consts.TaskStatusError { - return nil, nil + if !errors.Is(err, vmrecycle.ErrRemoteDelete) || !r.recycleQueue.IsFinalAttempt(job) { + return err } - return []*db.Task{tk}, nil + job.Payload.RecycleDeleteExhausted = true + return r.markRecycledForReconnectCompensation(ctx, job, err) } -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) 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) { diff --git a/backend/domain/host.go b/backend/domain/host.go index c25c270ad..8438e5f9d 100644 --- a/backend/domain/host.go +++ b/backend/domain/host.go @@ -89,7 +89,9 @@ type VmIdleInfo struct { EnvID string `json:"env_id"` TaskID string `json:"task_id,omitempty"` // 关联的任务 ID，用于通知 Name string `json:"name,omitempty"` // 任务名称，用于通知内容 - // RecycleAt 是本次 Refresh 算出的预计回收时间。每次用户活动都会延长这个值， + // RecycleDeleteExhausted 表示远端删除已重试耗尽，后续消费只补写回收状态。 + RecycleDeleteExhausted bool `json:"recycle_delete_exhausted,omitempty"` + // RecycleAt 是本次 RecordActivity 算出的预计回收时间。每次用户活动都会延长这个值， // consumer 把它编进 RefID，让每个回收窗口都能产生不同的 dedup key（否则 // dispatcher 会按 (subID, eventType, RefID) 把同一 task 的后续推送全部静默）。 RecycleAt time.Time `json:"recycle_at"` diff --git a/backend/pkg/delayqueue/delayqueue.go b/backend/pkg/delayqueue/delayqueue.go index 92baf2ae7..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 中存在孤立条目。 @@ -115,17 +123,32 @@ func (q *RedisDelayQueue[T]) Enqueue(ctx context.Context, queue string, payload return id, nil } -// 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 +// 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 nil, time.Time{}, false, err + 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) + 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 +160,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 { @@ -207,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 @@ -319,3 +357,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 new file mode 100644 index 000000000..494c2bf9c --- /dev/null +++ b/backend/pkg/delayqueue/delayqueue_test.go @@ -0,0 +1,144 @@ +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) + } +} + +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) + } +} + +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/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/register.go b/backend/pkg/register.go index 2b295e0d3..c1cbf92be 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,10 @@ 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) + do.Provide(i, vmrecycle.NewDeadlineRepairer) + // Channel Registry（通知渠道） do.Provide(i, func(i *do.Injector) (*msgpush.WechatClient, error) { cfg := do.MustInvoke[*config.Config](i) 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/analyzer.go b/backend/pkg/vmrecycle/analyzer.go new file mode 100644 index 000000000..c51b67b8a --- /dev/null +++ b/backend/pkg/vmrecycle/analyzer.go @@ -0,0 +1,351 @@ +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" +) + +const ( + analyzePageSize = 100 + analyzeWorkerCount = 8 +) + +type Decision string + +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" + DecisionNotFound Decision = "not_found" +) + +type TaskInfo struct { + ID uuid.UUID + Status consts.TaskStatus + LogStore consts.LogStore + CreatedAt time.Time + LastActiveAt 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 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 + 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), + 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)). + 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 shouldReport(item.Decision) { + 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{Target: "vm:" + vm.ID, 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, + LastActiveAt: tk.LastActiveAt, + }) + } + if vm.IsRecycled { + item.Decision = DecisionAlreadyRecycled + item.Reason = "VM already recycled" + return item + } + if len(item.Tasks) == 0 { + item.Decision = DecisionOrphan + 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.LastActiveAt + if lastActivity.IsZero() || lastActivity.Before(tk.CreatedAt) { + lastActivity = tk.CreatedAt + } + 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 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 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) { + 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 || decision == DecisionOrphan +} + +func IsDeadlineRepairable(decision Decision) bool { + return decision == DecisionDeadlineMissing +} + +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..dd5b44826 --- /dev/null +++ b/backend/pkg/vmrecycle/analyzer_test.go @@ -0,0 +1,257 @@ +package vmrecycle + +import ( + "context" + "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 + 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), + LastActiveAt: activityAt, + LogStore: &logStore, + }}}, + } + redisAt := now.Add(-time.Minute) + analyzer := newAnalyzerForTest(now, &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 TestAnalyzerMarksActiveTaskForDeadlineRepairWhenRedisDeadlineMissing(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + taskID := uuid.New() + 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) + } +} + +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() + 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), LastActiveAt: now.Add(-2 * time.Hour)}, + {ID: recentTaskID, CreatedAt: now.Add(-2 * time.Hour), LastActiveAt: now.Add(-30 * time.Minute)}, + }}} + 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 { + 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) + } +} + +func TestAnalyzerFallsBackToCreatedAtWhenLastActiveAtIsZero(t *testing.T) { + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + 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, &analyzerDeadlineStub{runAt: map[string]time.Time{vm.ID: now.Add(-time.Minute)}}) + + got := analyzer.analyzeVM(context.Background(), vm) + 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(), &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, &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, 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, + }}, + deadlines: deadlines, + now: func() time.Time { return now }, + } +} + +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/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 }, + } +} diff --git a/backend/pkg/vmrecycle/vmrecycle.go b/backend/pkg/vmrecycle/vmrecycle.go new file mode 100644 index 000000000..a9d80adee --- /dev/null +++ b/backend/pkg/vmrecycle/vmrecycle.go @@ -0,0 +1,252 @@ +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") + ErrRemoteDelete = errors.New("vm remote delete failed") +) + +type Result struct { + VMID string + TaskIDs []uuid.UUID + Status Status +} + +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 +} + +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) { + 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 { + 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 { + deleteErr := r.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{ + UserID: vm.UserID.String(), + HostID: vm.HostID, + ID: vm.EnvironmentID, + }) + 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) + 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..1c87103fe --- /dev/null +++ b/backend/pkg/vmrecycle/vmrecycle_test.go @@ -0,0 +1,387 @@ +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 !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) + } +} + +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 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) + 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 +}