支持同一个用户可以覆盖自己提交的镜像

models/cloudbrain_image.go

@@ -12,6 +12,9 @@ import (
 
 const RECOMMOND_TYPE = 5
 const NORMAL_TYPE = 0
+const IMAGE_STATUS_COMMIT = 0
+const IMAGE_STATUS_SUCCESS = 1
+const IMAGE_STATUS_Failed = 2
 
 type Image struct {
 	ID             int64              `xorm:"pk autoincr" json:"id"`
@@ -27,6 +30,7 @@ type Image struct {
 	IsStar         bool               `xorm:"-" json:"isStar"`
 	UserName       string             `xorm:"-" json:"userName"`
 	RelAvatarLink  string             `xorm:"-" json:"relAvatarLink"`
+	Status         int                `xorm:"NOT NULL DEFAULT 0" json:"status"` //0代表正在提交，1提交完成，2提交失败
 	CreatedUnix    timeutil.TimeStamp `xorm:"INDEX created" json:"createdUnix"`
 	UpdatedUnix    timeutil.TimeStamp `xorm:"INDEX updated" json:"updatedUnix"`
 }
@@ -70,6 +74,10 @@ type ErrorImageTagExist struct {
 	Tag string
 }
 
+type ErrorImageCommitting struct {
+	Tag string
+}
+
 type ImagesPageResult struct {
 	Count  int64    `json:"count"`
 	Images []*Image `json:"images"`
@@ -79,12 +87,27 @@ func (err ErrorImageTagExist) Error() string {
 	return fmt.Sprintf("Image already exists [tag: %s]", err.Tag)
 }
 
+func (err ErrorImageCommitting) Error() string {
+	return fmt.Sprintf("Image already exists [tag: %s]", err.Tag)
+}
+
 type ErrImageNotExist struct {
-	ID int64
+	ID  int64
+	Tag string
 }
 
 func (err ErrImageNotExist) Error() string {
-	return fmt.Sprintf("Image does not exist [id: %d]", err.ID)
+	return fmt.Sprintf("Image does not exist [id: %d] [tag: %s]", err.ID, err.Tag)
+}
+
+func IsErrorImageCommitting(err error) bool {
+	_, ok := err.(ErrorImageCommitting)
+	return ok
+}
+
+func IsErrImageNotExist(err error) bool {
+	_, ok := err.(ErrImageNotExist)
+	return ok
 }
 
 func IsErrImageTagExist(err error) bool {
@@ -98,6 +121,14 @@ func IsImageExist(tag string) (bool, error) {
 	})
 }
 
+func IsImageExistByUser(tag string, uid int64) (bool, error) {
+	return x.Exist(&Image{
+		Tag:    tag,
+		UID:    uid,
+		Status: IMAGE_STATUS_SUCCESS,
+	})
+}
+
 type FindImageTopicOptions struct {
 	ListOptions
 	ImageID int64
@@ -125,12 +156,26 @@ func GetImageByID(id int64) (*Image, error) {
 	if err != nil {
 		return nil, err
 	} else if !has {
-		return nil, ErrImageNotExist{id}
+		return nil, ErrImageNotExist{ID: id}
 	}
 
 	return rel, nil
 }
 
+func GetImageByTag(tag string) (*Image, error) {
+
+	image := &Image{Tag: tag}
+	has, err := x.
+		Get(image)
+	if err != nil {
+		return nil, err
+	} else if !has {
+		return nil, ErrImageNotExist{Tag: tag}
+	}
+
+	return image, nil
+}
+
 func SanitizeAndValidateImageTopics(topics []string) (validTopics []string, invalidTopics []string) {
 	validTopics = make([]string, 0)
 	mValidTopics := make(map[string]struct{})
@@ -431,6 +476,16 @@ func (images ImageList) loadAttributes(e Engine, uid int64) error {
 	return nil
 }
 
+func GetCommittingImageCount() int {
+
+	total, err := x.Where("status =?", 0).Count(new(Image))
+
+	if err != nil {
+		return 0
+	}
+	return int(total)
+}
+
 func CreateLocalImage(image *Image) error {
 
 	_, err := x.Insert(image)
@@ -439,7 +494,13 @@ func CreateLocalImage(image *Image) error {
 
 func UpdateLocalImage(image *Image) error {
 
-	_, err := x.ID(image.ID).Cols("description", "is_private").Update(image)
+	_, err := x.ID(image.ID).Cols("description", "is_private", "status").Update(image)
+	return err
+}
+
+func UpdateLocalImageStatus(image *Image) error {
+
+	_, err := x.ID(image.ID).Cols("status").Update(image)
 	return err
 }
 

modules/cloudbrain/resty.go

@@ -4,9 +4,11 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"net/http"
 	"strconv"
 	"strings"
+	"time"
 
 	"code.gitea.io/gitea/modules/log"
 
@@ -211,13 +213,23 @@ func getQueryString(page int, size int, name string) string {
 
 func CommitImage(jobID string, params models.CommitImageParams) error {
 
-	has, err := models.IsImageExist(params.ImageTag)
-	if err != nil {
+	dbImage, err := models.GetImageByTag(params.ImageTag)
+
+	if err != nil && !models.IsErrImageNotExist(err) {
 		return fmt.Errorf("resty CommitImage: %v", err)
 	}
-	if has {
-		return models.ErrorImageTagExist{
-			Tag: params.ImageTag,
+	if dbImage != nil {
+		if dbImage.UID != params.UID {
+			return models.ErrorImageTagExist{
+				Tag: params.ImageTag,
+			}
+		} else {
+			if dbImage.Status == models.IMAGE_STATUS_COMMIT {
+				return models.ErrorImageCommitting{
+					Tag: params.ImageTag,
+				}
+
+			}
 		}
 	}
 
@@ -245,23 +257,36 @@ sendjob:
 		goto sendjob
 	}
 
-	//if result.Code != Success {
-	//return fmt.Errorf("CommitImage err: %s", res.String())
-	//}
+	if result.Code != Success {
+		return fmt.Errorf("CommitImage err: %s", res.String())
+	}
+
+	image := models.Image{
+		Type:           models.NORMAL_TYPE,
+		CloudbrainType: params.CloudBrainType,
+		UID:            params.UID,
+		IsPrivate:      params.IsPrivate,
+		Tag:            params.ImageTag,
+		Description:    params.ImageDescription,
+		Place:          setting.Cloudbrain.ImageURLPrefix + params.ImageTag,
+		Status:         models.IMAGE_STATUS_COMMIT,
+	}
 
 	err = models.WithTx(func(ctx models.DBContext) error {
-		image := models.Image{
-			Type:           models.NORMAL_TYPE,
-			CloudbrainType: params.CloudBrainType,
-			UID:            params.UID,
-			IsPrivate:      params.IsPrivate,
-			Tag:            params.ImageTag,
-			Description:    params.ImageDescription,
-			Place:          setting.Cloudbrain.ImageURLPrefix + params.ImageTag,
-		}
-		if err := models.CreateLocalImage(&image); err != nil {
-			log.Error("Failed to insert image record.", err)
-			return fmt.Errorf("CommitImage err: %s", res.String())
+		if dbImage != nil {
+			dbImage.IsPrivate = params.IsPrivate
+			dbImage.Description = params.ImageDescription
+			dbImage.Status = models.IMAGE_STATUS_COMMIT
+			if err := models.UpdateLocalImage(dbImage); err != nil {
+				log.Error("Failed to update image record.", err)
+				return fmt.Errorf("CommitImage err: %s", res.String())
+			}
+
+		} else {
+			if err := models.CreateLocalImage(&image); err != nil {
+				log.Error("Failed to insert image record.", err)
+				return fmt.Errorf("CommitImage err: %s", res.String())
+			}
 		}
 		if err := models.SaveImageTopics(image.ID, params.Topics...); err != nil {
 			log.Error("Failed to insert image record.", err)
@@ -269,9 +294,45 @@ sendjob:
 		}
 		return nil
 	})
+	if err == nil {
+		go updateImageStatus(image)
+	}
 	return err
 }
 
+func updateImageStatus(image models.Image) {
+	attemps := 5
+	commitSuccess := false
+	time.Sleep(5 * time.Second)
+	for i := 0; i < attemps; i++ {
+		pageSize := models.GetCommittingImageCount()
+
+		result, err := GetImagesPageable(1, pageSize, Custom, "")
+		if err == nil && result.Code == "S000" {
+			for _, v := range result.Payload.ImageInfo {
+				if v.Place == image.Place {
+					image.Status = models.IMAGE_STATUS_SUCCESS
+					models.UpdateLocalImageStatus(&image)
+					commitSuccess = true
+					break
+				}
+
+			}
+
+		}
+		//第一次循环等待4秒，第二次等待4的2次方16秒，...,第5次。。。   ，总共大概是20多分钟内进行5次重试
+		var sleepTime = time.Duration(int(math.Pow(4, (float64(i + 1)))))
+
+		time.Sleep(sleepTime * time.Second)
+
+	}
+	if !commitSuccess {
+		image.Status = models.IMAGE_STATUS_Failed
+		models.UpdateLocalImageStatus(&image)
+	}
+
+}
+
 func StopJob(jobID string) error {
 	checkSetting()
 	client := getRestyClient()

options/locale/locale_en-US.ini

@@ -933,10 +933,12 @@ model_download=Model Download
 submit_image=Submit Image
 modify_image=Modify Image
 image_exist=Image name has been used, please use a new one.
+image_committing=Image is submitting, please try again later.
 image_commit_fail=Failed to submit image, please try again later.
 image_not_exist=Image does not exits.
 image_edit_fail=Failed to edit image, please try again later.
 image_delete_fail=Failed to delete image, please try again later.
+image_overwrite=You had submitted the same name image before, are you sure to overwrite the original image?
 download=Download
 score=Score
 

options/locale/locale_zh-CN.ini

@@ -936,10 +936,12 @@ model_download=结果下载
 submit_image=提交镜像
 modify_image=修改镜像
 image_exist=镜像Tag已被使用，请修改镜像Tag。
+image_committing=镜像正在提交中，请稍后再试。
 image_commit_fail=提交镜像失败，请稍后再试。
 image_not_exist=镜像不存在。
 image_edit_fail=编辑镜像失败，请稍后再试。
 image_delete_fail=删除镜像失败，请稍后再试。
+image_overwrite=您已经提交过相同名称的镜像，您确定要覆盖原来提交的镜像吗？
 download=模型下载
 score=评分
 

routers/repo/cloudbrain.go

@@ -678,6 +678,16 @@ func CloudBrainImageDelete(ctx *context.Context) {
 
 }
 
+func CloudBrainCommitImageCheck(ctx *context.Context, form auth.CommitImageCloudBrainForm) {
+	isExist, _ := models.IsImageExistByUser(form.Tag, ctx.User.ID)
+	if isExist {
+		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("repo.image_overwrite")))
+	} else {
+		ctx.JSON(http.StatusOK, models.BaseOKMessage)
+	}
+
+}
+
 func CloudBrainCommitImage(ctx *context.Context, form auth.CommitImageCloudBrainForm) {
 
 	if !NamePattern.MatchString(form.Tag) {
@@ -713,6 +723,8 @@ func CloudBrainCommitImage(ctx *context.Context, form auth.CommitImageCloudBrain
 		if models.IsErrImageTagExist(err) {
 			ctx.JSON(200, models.BaseErrorMessage(ctx.Tr("repo.image_exist")))
 
+		} else if models.IsErrorImageCommitting(err) {
+			ctx.JSON(200, models.BaseErrorMessage(ctx.Tr("repo.image_committing")))
 		} else {
 			ctx.JSON(200, models.BaseErrorMessage(ctx.Tr("repo.image_commit_fail")))
 		}

routers/routes/routes.go

@@ -1031,6 +1031,7 @@ func RegisterRoutes(m *macaron.Macaron) {
 				m.Get("", reqRepoCloudBrainReader, repo.CloudBrainShow)
 				m.Get("/debug", cloudbrain.AdminOrJobCreaterRight, repo.CloudBrainDebug)
 				m.Get("/commit_image", cloudbrain.AdminOrJobCreaterRight, repo.CloudBrainCommitImageShow)
+				m.Post("/commit_image/check", cloudbrain.AdminOrJobCreaterRight, bindIgnErr(auth.CommitImageCloudBrainForm{}), repo.CloudBrainCommitImageCheck)
 				m.Post("/commit_image", cloudbrain.AdminOrJobCreaterRight, bindIgnErr(auth.CommitImageCloudBrainForm{}), repo.CloudBrainCommitImage)
 				m.Post("/stop", cloudbrain.AdminOrOwnerOrJobCreaterRight, repo.CloudBrainStop)
 				m.Post("/del", cloudbrain.AdminOrOwnerOrJobCreaterRight, repo.CloudBrainDel)