From 77c04405e433d8fef6bc0de864a92a8087d74800 Mon Sep 17 00:00:00 2001
From: zc <zc2638@qq.com>
Date: Sat, 3 Feb 2024 02:03:01 +0800
Subject: [PATCH] fix cancel build to cancel pending stages

---
 core/handler/server/log.go  |  4 +++-
 core/service/build/build.go | 38 +++++++++++++++++++++++++++++++++++++
 go.mod                      |  2 +-
 go.sum                      |  4 ++--
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/core/handler/server/log.go b/core/handler/server/log.go
index dfe4c1e..484f10a 100644
--- a/core/handler/server/log.go
+++ b/core/handler/server/log.go
@@ -16,11 +16,13 @@ package server
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"strconv"
 
 	"github.com/99nil/gopkg/ctr"
 	"github.com/99nil/gopkg/sse"
+	"gorm.io/gorm"
 
 	"github.com/zc2638/ink/core/handler/wrapper"
 	storageV1 "github.com/zc2638/ink/pkg/api/storage/v1"
@@ -62,7 +64,7 @@ func logInfo() http.HandlerFunc {
 		}
 		logS := new(storageV1.Log)
 		logS.SetID(stepS.ID)
-		if err := db.Where(logS).First(logS).Error; err != nil {
+		if err := db.Where(logS).First(logS).Error; err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
 			wrapper.InternalError(w, err)
 			return
 		}
diff --git a/core/service/build/build.go b/core/service/build/build.go
index 577da9d..2e1d35e 100644
--- a/core/service/build/build.go
+++ b/core/service/build/build.go
@@ -20,6 +20,8 @@ import (
 	"fmt"
 	"slices"
 
+	"github.com/99nil/gopkg/sets"
+
 	"github.com/zc2638/ink/core/scheduler"
 	"gorm.io/gorm"
 
@@ -272,6 +274,42 @@ func (s *srv) Cancel(ctx context.Context, namespace, name string, number uint64)
 		return errors.New("already done")
 	}
 
+	var stages []storageV1.Stage
+	if err := db.Where(&storageV1.Stage{
+		BoxID:   buildS.BoxID,
+		BuildID: buildS.ID,
+		Phase:   v1.PhasePending.String(),
+	}).Find(&stages).Error; err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return fmt.Errorf("find pending stages failed: %v", err)
+	}
+
+	stageIDSet := sets.New[uint64]()
+	for _, v := range stages {
+		stageIDSet.Add(v.ID)
+	}
+	stageIds := stageIDSet.List()
+
+	if len(stageIds) > 0 {
+		err := db.Transaction(func(tx *gorm.DB) error {
+			if err := tx.Model(&storageV1.Stage{}).
+				Where("id in (?)", stageIds).
+				Where(&storageV1.Stage{Phase: v1.PhasePending.String()}).
+				Update("phase", v1.PhaseCanceled).Error; err != nil {
+				return fmt.Errorf("cancel pending stages failed: %v", err)
+			}
+			if err := tx.Model(&storageV1.Step{}).
+				Where("stage_id in (?)", stageIds).
+				Where(&storageV1.Step{Phase: v1.PhasePending.String()}).
+				Update("phase", v1.PhaseCanceled).Error; err != nil {
+				return fmt.Errorf("cancel pending steps failed: %v", err)
+			}
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+	}
+
 	sched := scheduler.FromContext(ctx)
 	return sched.Cancel(ctx, int64(build.ID))
 }
diff --git a/go.mod b/go.mod
index 78748fb..f744ede 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/zc2638/ink
 go 1.21.0
 
 require (
-	github.com/99nil/gopkg v0.0.0-20240131094835-25a1aa9ff089
+	github.com/99nil/gopkg v0.0.0-20240202180301-972815b49f29
 	github.com/docker/distribution v2.8.2+incompatible
 	github.com/docker/docker v24.0.5+incompatible
 	github.com/docker/go-units v0.5.0
diff --git a/go.sum b/go.sum
index b6e05b9..e49aafd 100644
--- a/go.sum
+++ b/go.sum
@@ -53,8 +53,8 @@ cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9
 cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3fOKtUw0Xmo=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8=
-github.com/99nil/gopkg v0.0.0-20240131094835-25a1aa9ff089 h1:9xHZRHLMtBq2R8GLL4MqAtKlwrssBDaZpNbp3QLE6pk=
-github.com/99nil/gopkg v0.0.0-20240131094835-25a1aa9ff089/go.mod h1:yQy7Xa2r1VHYfgqfCuDE6VGDmBt8xJgDVTZ2Bw3xTW0=
+github.com/99nil/gopkg v0.0.0-20240202180301-972815b49f29 h1:x6ZHT+3qjWgk3ZwHsAnMc7SXupNlWUNqvrJYVjIi+kQ=
+github.com/99nil/gopkg v0.0.0-20240202180301-972815b49f29/go.mod h1:yQy7Xa2r1VHYfgqfCuDE6VGDmBt8xJgDVTZ2Bw3xTW0=
 github.com/AdaLogics/go-fuzz-headers v0.0.0-20210715213245-6c3934b029d8/go.mod h1:CzsSbkDixRphAF5hS6wbMKq0eI6ccJRb7/A0M6JBnwg=
 github.com/Azure/azure-pipeline-go v0.2.3/go.mod h1:x841ezTBIMG6O3lAcl8ATHnsOPVl2bqk7S3ta6S6u4k=
 github.com/Azure/azure-sdk-for-go v16.2.1+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=