Merge pull request #531 from chaitin/feat/task-last-active-time

feat(task): 增加任务最近活跃时间

Author: yokowu

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

@@ -46,6 +46,7 @@ type TaskHandler struct {
 	taskConns     *ws.TaskConn
 	controlConns  *ws.ControlConn
 	taskSummary   *service.TaskSummaryService
+	taskActivity  service.TaskActivityRefresher
 	idleRefresher vmidle.VMIdleRefresher
 	activeRepo    domain.UserActiveRepo
 }
@@ -64,6 +65,7 @@ func NewTaskHandler(i *do.Injector) (*TaskHandler, error) {
 	tc := do.MustInvoke[*ws.TaskConn](i)
 	cc := do.MustInvoke[*ws.ControlConn](i)
 	ts := do.MustInvoke[*service.TaskSummaryService](i)
+	ta := do.MustInvoke[service.TaskActivityRefresher](i)
 	ir := do.MustInvoke[vmidle.VMIdleRefresher](i)
 
 	// Optional deps
@@ -91,6 +93,7 @@ func NewTaskHandler(i *do.Injector) (*TaskHandler, error) {
 		taskConns:     tc,
 		controlConns:  cc,
 		taskSummary:   ts,
+		taskActivity:  ta,
 		idleRefresher: ir,
 		activeRepo:    activeRepo,
 	}

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

@@ -10,7 +10,9 @@ 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"
@@ -130,6 +132,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 {
@@ -182,7 +187,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, vm.ID)
+			return h.controlKeepAlive(ctx, task.ID, vm.ID)
 		})
 	}
 
@@ -211,17 +216,30 @@ func (h *TaskHandler) controlPing(ctx context.Context, wsConn *ws.WebsocketManag
 }
 
 // controlKeepAlive 定期刷新空闲计时器，防止 VM 被误判空闲
-func (h *TaskHandler) controlKeepAlive(ctx context.Context, vmID string) error {
-	ticker := time.NewTicker(1 * time.Minute)
-	defer ticker.Stop()
+func (h *TaskHandler) controlKeepAlive(ctx context.Context, taskID uuid.UUID, vmID string) error {
+	if err := h.idleRefresher.Refresh(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 <-ticker.C:
+		case <-idleTicker.C:
 			if err := h.idleRefresher.Refresh(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

@@ -0,0 +1,93 @@
+package v1
+
+import (
+	"context"
+	"io"
+	"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)
+	}()
+
+	select {
+	case got := <-idleRefresher.ch:
+		if got != vmID {
+			t.Fatalf("idle refresher vm id = %q, want %q", got, vmID)
+		}
+	case <-time.After(100 * time.Millisecond):
+		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 {
+	case err := <-done:
+		if err == nil {
+			t.Fatal("controlKeepAlive() error = nil, want context canceled")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("controlKeepAlive() did not exit after cancel")
+	}
+}
+
+type testVMIdleRefresher struct {
+	ch chan string
+}
+
+func (r *testVMIdleRefresher) Refresh(_ 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)

backend/biz/task/register.go

@@ -13,6 +13,7 @@ import (
 func ProvideTask(i *do.Injector) {
 	do.Provide(i, usecase.NewTaskUsecase)
 	do.Provide(i, usecase.NewGitTaskUsecase)
+	do.Provide(i, service.NewTaskActivityRefresher)
 	do.Provide(i, service.NewTaskSummaryService)
 	do.Provide(i, v1.NewTaskHandler)
 	do.Provide(i, repo.NewTaskRepo)

backend/biz/task/repo/task.go

@@ -261,6 +261,15 @@ func (t *TaskRepo) Update(ctx context.Context, _ *domain.User, id uuid.UUID, fn
 	})
 }
 
+func (t *TaskRepo) RefreshLastActiveAt(ctx context.Context, id uuid.UUID, at time.Time, minInterval time.Duration) error {
+	up := t.db.Task.Update().Where(task.ID(id))
+	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().

backend/biz/task/service/activity.go

@@ -0,0 +1,42 @@
+package service
+
+import (
+	"context"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/samber/do"
+
+	"github.com/chaitin/MonkeyCode/backend/domain"
+)
+
+const TaskActivityRefreshInterval = 5 * time.Minute
+
+type TaskActivityRefresher interface {
+	Refresh(ctx context.Context, taskID uuid.UUID) error
+	ForceRefresh(ctx context.Context, taskID uuid.UUID) error
+}
+
+type taskActivityRefresher struct {
+	repo  taskActivityRepo
+	clock func() time.Time
+}
+
+type taskActivityRepo interface {
+	RefreshLastActiveAt(ctx context.Context, id uuid.UUID, at time.Time, minInterval time.Duration) error
+}
+
+func NewTaskActivityRefresher(i *do.Injector) (TaskActivityRefresher, error) {
+	return &taskActivityRefresher{
+		repo:  do.MustInvoke[domain.TaskRepo](i),
+		clock: time.Now,
+	}, nil
+}
+
+func (r *taskActivityRefresher) Refresh(ctx context.Context, taskID uuid.UUID) error {
+	return r.repo.RefreshLastActiveAt(ctx, taskID, r.clock(), TaskActivityRefreshInterval)
+}
+
+func (r *taskActivityRefresher) ForceRefresh(ctx context.Context, taskID uuid.UUID) error {
+	return r.repo.RefreshLastActiveAt(ctx, taskID, r.clock(), 0)
+}

backend/biz/task/service/activity_test.go

@@ -0,0 +1,70 @@
+package service
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+func TestTaskActivityRefresherRefreshUsesThrottleInterval(t *testing.T) {
+	taskID := uuid.MustParse("11111111-1111-1111-1111-111111111111")
+	now := time.Unix(1_700_000_000, 0).UTC()
+	repo := &taskActivityRepoStub{}
+	refresher := &taskActivityRefresher{
+		repo:  repo,
+		clock: func() time.Time { return now },
+	}
+
+	if err := refresher.Refresh(context.Background(), taskID); err != nil {
+		t.Fatalf("Refresh() error = %v", err)
+	}
+
+	if repo.taskID != taskID {
+		t.Fatalf("task id = %s, want %s", repo.taskID, taskID)
+	}
+	if !repo.at.Equal(now) {
+		t.Fatalf("refresh time = %s, want %s", repo.at, now)
+	}
+	if repo.minInterval != TaskActivityRefreshInterval {
+		t.Fatalf("min interval = %s, want %s", repo.minInterval, TaskActivityRefreshInterval)
+	}
+}
+
+func TestTaskActivityRefresherForceRefreshBypassesThrottle(t *testing.T) {
+	taskID := uuid.MustParse("22222222-2222-2222-2222-222222222222")
+	now := time.Unix(1_700_000_100, 0).UTC()
+	repo := &taskActivityRepoStub{}
+	refresher := &taskActivityRefresher{
+		repo:  repo,
+		clock: func() time.Time { return now },
+	}
+
+	if err := refresher.ForceRefresh(context.Background(), taskID); err != nil {
+		t.Fatalf("ForceRefresh() error = %v", err)
+	}
+
+	if repo.taskID != taskID {
+		t.Fatalf("task id = %s, want %s", repo.taskID, taskID)
+	}
+	if !repo.at.Equal(now) {
+		t.Fatalf("refresh time = %s, want %s", repo.at, now)
+	}
+	if repo.minInterval != 0 {
+		t.Fatalf("min interval = %s, want 0", repo.minInterval)
+	}
+}
+
+type taskActivityRepoStub struct {
+	taskID      uuid.UUID
+	at          time.Time
+	minInterval time.Duration
+}
+
+func (s *taskActivityRepoStub) RefreshLastActiveAt(_ context.Context, taskID uuid.UUID, at time.Time, minInterval time.Duration) error {
+	s.taskID = taskID
+	s.at = at
+	s.minInterval = minInterval
+	return nil
+}