From 5a7904fe3b978856678efe16f3425654f70790f2 Mon Sep 17 00:00:00 2001
From: ychao_1983 <ychao_1983@sina.com>
Date: Sat, 9 Oct 2021 09:43:24 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D#333?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 models/cloudbrain.go       | 12 +++++++++++
 routers/repo/cloudbrain.go | 54 ++++++++++++++++++++++++++++++++++++++++++++++
 routers/repo/setting.go    |  1 +
 routers/user/auth.go       |  3 +++
 4 files changed, 70 insertions(+)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index edd5f102a..5be07f741 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -620,6 +620,18 @@ func GetCloudbrainByJobID(jobID string) (*Cloudbrain, error) {
 	return getRepoCloudBrain(cb)
 }
 
+func GetCloudbrainsNeededStopByUserID(userID int64) ([]*Cloudbrain, error) {
+	cloudBrains := make([]*Cloudbrain, 0)
+	err := x.Cols("job_id", "status", "type").Where("user_id=? AND (status =? OR status=?)", userID, string(JobRunning), string(JobWaiting)).Find(&cloudBrains)
+	return cloudBrains, err
+}
+
+func GetCloudbrainsNeededStopByRepoID(repoID int64) ([]*Cloudbrain, error) {
+	cloudBrains := make([]*Cloudbrain, 0)
+	err := x.Cols("job_id", "status", "type").Where("repo_id=? AND (status =? OR status=?)", repoID, string(JobRunning), string(JobWaiting)).Find(&cloudBrains)
+	return cloudBrains, err
+}
+
 func SetCloudbrainStatusByJobID(jobID string, status CloudbrainStatus) (err error) {
 	cb := &Cloudbrain{JobID: jobID, Status: string(status)}
 	_, err = x.Cols("status").Where("cloudbrain.job_id=?", jobID).Update(cb)
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index deac7b312..c30906af4 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -13,6 +13,8 @@ import (
 	"strings"
 	"time"
 
+	"code.gitea.io/gitea/modules/modelarts"
+
 	"code.gitea.io/gitea/modules/git"
 	"code.gitea.io/gitea/modules/storage"
 
@@ -361,6 +363,58 @@ func CloudBrainStop(ctx *context.Context) {
 	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain")
 }
 
+func StopJobsByUserID(userID int64) {
+	cloudBrains, err := models.GetCloudbrainsNeededStopByUserID(userID)
+	if err != nil {
+		log.Warn("Failed to get cloudBrain info", err)
+		return
+	}
+	StopJobs(cloudBrains)
+
+}
+
+func StopJobsByRepoID(repoID int64) {
+	cloudBrains, err := models.GetCloudbrainsNeededStopByRepoID(repoID)
+	if err != nil {
+		log.Warn("Failed to get cloudBrain info", err)
+		return
+	}
+	StopJobs(cloudBrains)
+}
+
+/**
+
+ */
+func StopJobs(cloudBrains []*models.Cloudbrain) {
+
+	for _, taskInfo := range cloudBrains {
+
+		if taskInfo.Type == models.TypeCloudBrainOne {
+			err := cloudbrain.StopJob(taskInfo.JobID)
+			logErrorAndUpdateJobStatus(err, taskInfo)
+		} else {
+			param := models.NotebookAction{
+				Action: models.ActionStop,
+			}
+			_, err := modelarts.StopJob(taskInfo.JobID, param)
+			logErrorAndUpdateJobStatus(err, taskInfo)
+		}
+
+	}
+}
+
+func logErrorAndUpdateJobStatus(err error, taskInfo *models.Cloudbrain) {
+	if err != nil {
+		log.Warn("Failed to stop cloudBrain job:"+taskInfo.JobID, err)
+	} else {
+		taskInfo.Status = string(models.JobStopped)
+		err = models.UpdateJob(taskInfo)
+		if err != nil {
+			log.Warn("UpdateJob failed", err)
+		}
+	}
+}
+
 func CloudBrainDel(ctx *context.Context) {
 	var jobID = ctx.Params(":jobid")
 	task, err := models.GetCloudbrainByJobID(jobID)
diff --git a/routers/repo/setting.go b/routers/repo/setting.go
index b2ca042de..bf11f9e5a 100644
--- a/routers/repo/setting.go
+++ b/routers/repo/setting.go
@@ -440,6 +440,7 @@ func SettingsPost(ctx *context.Context, form auth.RepoSettingForm) {
 			return
 		}
 		log.Trace("Repository deleted: %s/%s", ctx.Repo.Owner.Name, repo.Name)
+		go StopJobsByRepoID(repo.ID)
 
 		ctx.Flash.Success(ctx.Tr("repo.settings.deletion_success"))
 		ctx.Redirect(ctx.Repo.Owner.DashboardLink())
diff --git a/routers/user/auth.go b/routers/user/auth.go
index 13e338565..126d0a4c8 100755
--- a/routers/user/auth.go
+++ b/routers/user/auth.go
@@ -11,6 +11,8 @@ import (
 	"net/http"
 	"strings"
 
+	"code.gitea.io/gitea/routers/repo"
+
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/auth"
 	"code.gitea.io/gitea/modules/auth/oauth2"
@@ -1056,6 +1058,7 @@ func SignOut(ctx *context.Context) {
 		})
 	}
 	HandleSignOut(ctx)
+	go repo.StopJobsByUserID(ctx.User.ID)
 	ctx.Redirect(setting.AppSubURL + "/")
 }