修复删除虚拟机后任务未结束

Author: yokowu

backend/biz/host/repo/host.go

@@ -21,6 +21,7 @@ import (
 	"github.com/chaitin/MonkeyCode/backend/db/model"
 	"github.com/chaitin/MonkeyCode/backend/db/predicate"
 	"github.com/chaitin/MonkeyCode/backend/db/projecttask"
+	"github.com/chaitin/MonkeyCode/backend/db/task"
 	"github.com/chaitin/MonkeyCode/backend/db/taskvirtualmachine"
 	"github.com/chaitin/MonkeyCode/backend/db/teamgroup"
 	"github.com/chaitin/MonkeyCode/backend/db/user"
@@ -610,6 +611,7 @@ func (h *HostRepo) CompleteCreateVirtualMachine(ctx context.Context, id string,
 func (h *HostRepo) DeleteVirtualMachine(ctx context.Context, uid uuid.UUID, hostID, id string, fn func(*db.VirtualMachine) error) error {
 	return entx.WithTx2(ctx, h.db, func(tx *db.Tx) error {
 		vm, err := tx.VirtualMachine.Query().
+			WithTasks().
 			Where(virtualmachine.ID(id)).
 			Where(virtualmachine.UserID(uid)).
 			First(ctx)
@@ -627,6 +629,27 @@ func (h *HostRepo) DeleteVirtualMachine(ctx context.Context, uid uuid.UUID, host
 			return err
 		}
 
+		taskIDs := make([]uuid.UUID, 0, len(vm.Edges.Tasks))
+		for _, tk := range vm.Edges.Tasks {
+			if tk == nil {
+				continue
+			}
+			taskIDs = append(taskIDs, tk.ID)
+		}
+		if len(taskIDs) > 0 {
+			_, err := tx.Task.Update().
+				Where(
+					task.IDIn(taskIDs...),
+					task.StatusNotIn(consts.TaskStatusFinished, consts.TaskStatusError),
+				).
+				SetStatus(consts.TaskStatusFinished).
+				SetCompletedAt(time.Now()).
+				Save(ctx)
+			if err != nil {
+				return err
+			}
+		}
+
 		_, err = tx.VirtualMachine.Delete().Where(virtualmachine.ID(id)).Exec(ctx)
 		if err != nil {
 			return err

backend/biz/host/usecase/host.go

@@ -533,7 +533,7 @@ func (h *HostUsecase) CreateVM(ctx context.Context, user *domain.User, req *doma
 
 // DeleteVM 删除虚拟机
 func (h *HostUsecase) DeleteVM(ctx context.Context, uid uuid.UUID, hostID, vmID string) error {
-	h.logger.InfoContext(ctx, "delete vm", "vmID", vmID)
+	h.logger.InfoContext(ctx, "delete vm", "vmID", vmID, "user_id", uid, "host_id", hostID)
 	return h.repo.DeleteVirtualMachine(ctx, uid, hostID, vmID, func(vm *db.VirtualMachine) error {
 		if err := h.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{
 			UserID: uid.String(),

backend/biz/host/usecase/host_test.go

@@ -23,6 +23,7 @@ import (
 	"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/chaitin/MonkeyCode/backend/pkg/taskflow"
 )
 
@@ -337,6 +338,97 @@ func TestHostUsecase_markRecycledTasksFinished(t *testing.T) {
 	}
 }
 
+func TestHostUsecase_DeleteVMFinishesBoundTasks(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	client := enttest.Open(t, "sqlite3", "file:host-usecase-delete-vm-finish-task-test?mode=memory&cache=shared&_fk=1")
+	defer 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.Fatalf("create user: %v", err)
+	}
+
+	hostID := "host-1"
+	if _, err := client.Host.Create().
+		SetID(hostID).
+		SetUserID(userID).
+		SetHostname("host").
+		Save(ctx); err != nil {
+		t.Fatalf("create host: %v", err)
+	}
+
+	vmID := "vm-1"
+	if _, err := client.VirtualMachine.Create().
+		SetID(vmID).
+		SetHostID(hostID).
+		SetUserID(userID).
+		SetName("vm").
+		SetEnvironmentID("env-1").
+		Save(ctx); err != nil {
+		t.Fatalf("create vm: %v", err)
+	}
+
+	taskID := uuid.New()
+	if _, err := client.Task.Create().
+		SetID(taskID).
+		SetUserID(userID).
+		SetKind(consts.TaskTypeDevelop).
+		SetContent("content").
+		SetStatus(consts.TaskStatusProcessing).
+		Save(ctx); err != nil {
+		t.Fatalf("create task: %v", err)
+	}
+	if _, err := client.TaskVirtualMachine.Create().
+		SetID(uuid.New()).
+		SetTaskID(taskID).
+		SetVirtualmachineID(vmID).
+		Save(ctx); err != nil {
+		t.Fatalf("bind task vm: %v", err)
+	}
+
+	i := do.New()
+	do.ProvideValue(i, client)
+	redisServer := miniredis.RunT(t)
+	redisClient := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
+	t.Cleanup(func() { _ = redisClient.Close() })
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	do.ProvideValue(i, &config.Config{})
+	do.ProvideValue(i, logger)
+	do.ProvideValue(i, redisClient)
+	hostRepo, err := repo.NewHostRepo(i)
+	if err != nil {
+		t.Fatalf("new host repo: %v", err)
+	}
+	u := &HostUsecase{
+		repo:          hostRepo,
+		taskflow:      &preinsertTaskflowStub{vm: &preinsertVMCreateStub{db: client}},
+		logger:        logger,
+		vmexpireQueue: delayqueue.NewVMExpireQueue(redisClient, logger),
+	}
+
+	if err := u.DeleteVM(ctx, userID, hostID, vmID); err != nil {
+		t.Fatalf("DeleteVM() error = %v", err)
+	}
+
+	gotTask, err := client.Task.Get(ctx, taskID)
+	if err != nil {
+		t.Fatalf("query task: %v", err)
+	}
+	if gotTask.Status != consts.TaskStatusFinished {
+		t.Fatalf("task status = %s, want %s", gotTask.Status, consts.TaskStatusFinished)
+	}
+	if gotTask.CompletedAt.IsZero() {
+		t.Fatal("expected task completed_at to be set")
+	}
+}
+
 func TestHostUsecase_CreateVMPreinsertsVirtualMachineBeforeTaskflowCreate(t *testing.T) {
 	t.Parallel()