Merge pull request #868 from chaitin/feat/vm-recycle-cli-support

修复 VM 空闲回收活动语义并统一回收流程

Author: yokowu

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)

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
+}

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)
-			}
 		}
 	}
 }

backend/biz/task/handler/v1/task_control_test.go

@@ -6,30 +6,23 @@ 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())
 	defer cancel()
 
 	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,28 +50,12 @@ 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:
 	}
 	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 }

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().