Merge pull request #522 from chaitin/feat/backend-vm-recycle-check-token-20260416

feat(host): 增加回收态虚拟机回收补偿逻辑

Author: yokowu

backend/biz/host/handler/v1/internal.go

@@ -22,6 +22,7 @@ import (
 	"github.com/chaitin/MonkeyCode/backend/domain"
 	etypes "github.com/chaitin/MonkeyCode/backend/ent/types"
 	"github.com/chaitin/MonkeyCode/backend/pkg/cvt"
+	"github.com/chaitin/MonkeyCode/backend/pkg/entx"
 	"github.com/chaitin/MonkeyCode/backend/pkg/lifecycle"
 	"github.com/chaitin/MonkeyCode/backend/pkg/taskflow"
 	"github.com/chaitin/MonkeyCode/backend/pkg/ws"
@@ -33,6 +34,10 @@ type InternalHostHandler struct {
 	repo           domain.HostRepo
 	teamRepo       domain.TeamHostRepo
 	redis          *redis.Client
+	getAgentToken  agentTokenGetter
+	limiter        *redis.Client
+	vmDeleter      taskflow.VirtualMachiner
+	skipSoftDelete func(context.Context) context.Context
 	cache          *cache.Cache
 	taskLifecycle  *lifecycle.Manager[uuid.UUID, consts.TaskStatus, lifecycle.TaskMetadata]
 	hostUsecase    domain.HostUsecase
@@ -43,16 +48,22 @@ type InternalHostHandler struct {
 
 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)
 
 	h := &InternalHostHandler{
 		logger:         do.MustInvoke[*slog.Logger](i).With("module", "InternalHostHandler"),
 		repo:           do.MustInvoke[domain.HostRepo](i),
 		teamRepo:       do.MustInvoke[domain.TeamHostRepo](i),
-		redis:          do.MustInvoke[*redis.Client](i),
+		redis:          rdb,
+		getAgentToken:  defaultAgentTokenGetter(rdb),
+		limiter:        rdb,
+		vmDeleter:      tf.VirtualMachiner(),
+		skipSoftDelete: entx.SkipSoftDelete,
 		cache:          cache.New(15*time.Minute, 10*time.Minute),
+		taskLifecycle:  do.MustInvoke[*lifecycle.Manager[uuid.UUID, consts.TaskStatus, lifecycle.TaskMetadata]](i),
 		hostUsecase:    do.MustInvoke[domain.HostUsecase](i),
 		taskConns:      do.MustInvoke[*ws.TaskConn](i),
-		taskLifecycle:  do.MustInvoke[*lifecycle.Manager[uuid.UUID, consts.TaskStatus, lifecycle.TaskMetadata]](i),
 		projectUsecase: do.MustInvoke[domain.ProjectUsecase](i),
 		tokenProvider:  do.MustInvoke[*gituc.TokenProvider](i),
 	}
@@ -178,55 +189,46 @@ func (h *InternalHostHandler) CheckToken(c *web.Context, req taskflow.CheckToken
 func (h *InternalHostHandler) agentAuth(ctx context.Context, token, mid string) (*taskflow.Token, error) {
 	// 1) 优先从 Redis 读取一次性 agent token，并清除
 	key := fmt.Sprintf("agent:token:%s", token)
-	luaGetDel := `
-local v = redis.call('GET', KEYS[1])
-if v then
-	 redis.call('DEL', KEYS[1])
-	 return v
-end
-return nil
-`
-	res, err := h.redis.Eval(ctx, luaGetDel, []string{key}).Result()
-	h.logger.With("mid", mid, "key", key, "res", res, "error", err).DebugContext(ctx, "agent auth...")
+	res, err := h.getAgentToken(ctx, key)
+	h.logger.With("mid", mid, "key", key, "hit", err == nil, "error", err).DebugContext(ctx, "agent auth...")
 	if err == nil {
-		if b, ok := res.(string); ok && b != "" {
-			var t taskflow.Token
-			if uerr := json.Unmarshal([]byte(b), &t); uerr != nil {
-				h.logger.With("error", uerr, "token", token).ErrorContext(ctx, "failed to unmarshal token from redis")
-				return nil, uerr
-			}
+		var t taskflow.Token
+		if uerr := json.Unmarshal([]byte(res), &t); uerr != nil {
+			h.logger.With("error", uerr, "token", token).ErrorContext(ctx, "failed to unmarshal token from redis")
+			return nil, uerr
+		}
 
-			if mid != "" {
-				if err := h.repo.UpdateVirtualMachine(ctx, token, func(up *db.VirtualMachineUpdateOne) error {
-					up.SetMachineID(mid)
-					return nil
-				}); err != nil {
-					h.logger.With("error", err, "token", token).ErrorContext(ctx, "failed to update virtual machine machine id")
-					return nil, err
-				}
+		if mid != "" {
+			if err := h.repo.UpdateVirtualMachine(ctx, token, func(up *db.VirtualMachineUpdateOne) error {
+				up.SetMachineID(mid)
+				return nil
+			}); err != nil {
+				h.logger.With("error", err, "token", token).ErrorContext(ctx, "failed to update virtual machine machine id")
+				return nil, err
 			}
-
-			return &t, nil
 		}
+
+		return &t, nil
 	} else if !errors.Is(err, redis.Nil) {
 		h.logger.With("error", err, "token", token).ErrorContext(ctx, "failed to get redis token via lua, fallback to db")
 	}
 
 	// 2) Redis 没值时根据数据库校验 token
-	vm, err := h.repo.GetVirtualMachine(ctx, token)
+	vm, err := h.repo.GetVirtualMachine(h.skipSoftDelete(ctx), token)
 	if err != nil {
 		return nil, err
 	}
 
+	if vm.IsRecycled {
+		h.tryRecycledVMDelete(ctx, vm, mid)
+		return nil, errAgentVMRecycled
+	}
+
 	// 机器码绑定校验
 	if mid != "" && vm.MachineID != "" && vm.MachineID != mid {
 		return nil, fmt.Errorf("mismatch machine id")
 	}
 
-	if vm.IsRecycled {
-		return nil, fmt.Errorf("vm is recycled")
-	}
-
 	if vm.Edges.Host == nil {
 		return nil, fmt.Errorf("no host found for vm")
 	}

backend/biz/host/handler/v1/internal_auth.go

@@ -0,0 +1,75 @@
+package v1
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+
+	"github.com/chaitin/MonkeyCode/backend/db"
+	"github.com/chaitin/MonkeyCode/backend/pkg/taskflow"
+)
+
+const (
+	recycledDeleteRetryTTL = 30 * time.Second
+	recycledDeleteTimeout  = 5 * time.Second
+)
+
+var errAgentVMRecycled = errors.New("agent vm is recycled")
+
+type agentTokenGetter func(ctx context.Context, key string) (string, error)
+
+func defaultAgentTokenGetter(rdb *redis.Client) agentTokenGetter {
+	const luaGetDel = `
+local v = redis.call('GET', KEYS[1])
+if v then
+	 redis.call('DEL', KEYS[1])
+	 return v
+end
+return nil
+`
+	return func(ctx context.Context, key string) (string, error) {
+		res, err := rdb.Eval(ctx, luaGetDel, []string{key}).Result()
+		if err != nil {
+			return "", err
+		}
+
+		b, ok := res.(string)
+		if !ok || b == "" {
+			return "", redis.Nil
+		}
+		return b, nil
+	}
+}
+
+func (h *InternalHostHandler) tryRecycledVMDelete(ctx context.Context, vm *db.VirtualMachine, machineID string) {
+	if h.limiter == nil || h.vmDeleter == nil {
+		h.logger.WarnContext(ctx, "skip recycled vm delete retry", "vm_id", vm.ID, "machine_id", machineID, "error", "missing dependency")
+		return
+	}
+
+	key := fmt.Sprintf("vm:recycle:retry:%s", vm.ID)
+	ok, err := h.limiter.SetNX(ctx, key, "1", recycledDeleteRetryTTL).Result()
+	if err != nil || !ok {
+		h.logger.WarnContext(ctx, "skip recycled vm delete retry", "vm_id", vm.ID, "machine_id", machineID, "rate_limited", !ok, "error", err)
+		return
+	}
+
+	go func() {
+		deleteCtx, cancel := context.WithTimeout(context.Background(), recycledDeleteTimeout)
+		defer cancel()
+
+		err := h.vmDeleter.Delete(deleteCtx, &taskflow.DeleteVirtualMachineReq{
+			UserID: vm.UserID.String(),
+			HostID: vm.HostID,
+			ID:     vm.EnvironmentID,
+		})
+		if err != nil {
+			h.logger.ErrorContext(deleteCtx, "reissue recycled vm delete failed", "vm_id", vm.ID, "machine_id", machineID, "error", err)
+			return
+		}
+		h.logger.InfoContext(deleteCtx, "reissue recycled vm delete success", "vm_id", vm.ID, "machine_id", machineID)
+	}()
+}