Skip to content
37 changes: 34 additions & 3 deletions backend/biz/host/handler/v1/internal.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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),
Expand Down Expand Up @@ -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)
Expand Down
98 changes: 96 additions & 2 deletions backend/biz/host/handler/v1/internal_vm_activity_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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}`
Expand All @@ -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()
Expand All @@ -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
}
27 changes: 6 additions & 21 deletions backend/biz/task/handler/v1/task_control.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
}

Expand All @@ -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)
}
}
Expand All @@ -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)
})
}

Expand Down Expand Up @@ -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)
}
}
}
}
Expand Down
38 changes: 3 additions & 35 deletions backend/biz/task/handler/v1/task_control_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 }
13 changes: 13 additions & 0 deletions backend/biz/task/repo/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -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().
Expand Down
Loading