aiforge/services/reward/cloubrain_deduct.go

package reward

import (
	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/services/reward/point/account"
	"encoding/json"
	"fmt"
	"time"
)

var (
	ResourceSpecs      *models.ResourceSpecs
	TrainResourceSpecs *models.ResourceSpecs
)

//IsPointBalanceEnough check whether the user's point balance is bigger than task unit price
func IsPointBalanceEnough(targetUserId int64, jobType string, resourceSpecId int) bool {
	if !setting.CloudBrainTaskPointPaySwitch {
		return true
	}
	spec := getResourceSpec(jobType, resourceSpecId)
	if spec == nil {
		return true
	}
	a, error := account.GetAccount(targetUserId)
	if error != nil {
		return false
	}
	return a.Balance >= spec.UnitPrice

}

func StartCloudBrainPointDeductTask(task models.Cloudbrain) {
	if !setting.CloudBrainTaskPointPaySwitch {
		return
	}

	spec := getResourceSpec(task.JobType, task.ResourceSpecId)
	if spec == nil || spec.UnitPrice == 0 {
		return
	}

	StartPeriodicTask(&models.StartPeriodicTaskOpts{
		SourceType:   models.SourceTypeRunCloudbrainTask,
		SourceId:     getCloudBrainPointTaskSourceId(task),
		TargetUserId: task.UserID,
		RequestId:    getCloudBrainPointTaskSourceId(task),
		OperateType:  models.OperateTypeDecrease,
		Delay:        30 * time.Minute,
		Interval:     60 * time.Minute,
		UnitAmount:   spec.UnitPrice,
		RewardType:   models.RewardTypePoint,
		StartTime:    time.Unix(int64(task.StartTime), 0),
	})
}

func StopCloudBrainPointDeductTask(task models.Cloudbrain) {
	StopPeriodicTask(models.SourceTypeRunCloudbrainTask, getCloudBrainPointTaskSourceId(task), models.OperateTypeDecrease)
}

func getCloudBrainPointTaskSourceId(task models.Cloudbrain) string {
	return models.SourceTypeRunCloudbrainTask.Name() + "_" + task.JobType + "_" + fmt.Sprint(task.Type) + "_" + fmt.Sprint(task.ID)
}

func getResourceSpec(jobType string, resourceSpecId int) *models.ResourceSpec {
	if jobType == string(models.JobTypeTrain) {
		if TrainResourceSpecs == nil {
			json.Unmarshal([]byte(setting.TrainResourceSpecs), &TrainResourceSpecs)
		}
		for _, spec := range TrainResourceSpecs.ResourceSpec {
			if resourceSpecId == spec.Id {
				return spec
			}
		}
	} else {
		if ResourceSpecs == nil {
			json.Unmarshal([]byte(setting.ResourceSpecs), &ResourceSpecs)
		}
		for _, spec := range ResourceSpecs.ResourceSpec {
			if resourceSpecId == spec.Id {
				return spec
			}
		}

	}
	return nil

}

var firstTimeFlag = true

func StartCloudbrainPointDeductTask() {
	defer func() {
		if err := recover(); err != nil {
			combinedErr := fmt.Errorf("%s\n%s", err, log.Stack(2))
			log.Error("PANIC:%v", combinedErr)
		}
	}()
	log.Debug("try to run CloudbrainPointDeductTask")
	end := time.Now()
	start := end.Add(5 * time.Minute)
	if firstTimeFlag {
		//When it is executed for the first time, it needs to process the tasks of the last 1 hours.
		//This is done to prevent the application from hanging for a long time
		start = end.Add(1 * time.Hour)
		firstTimeFlag = false
	}

	taskList, err := models.GetStartedCloudbrainTaskByUpdatedUnix(start, end)
	if err != nil {
		log.Error("GetStartedCloudbrainTaskByUpdatedUnix error. %v", err)
		return
	}
	if taskList == nil || len(taskList) == 0 {
		log.Debug("No cloudbrain task need handled")
		return
	}
	for _, t := range taskList {
		if int64(t.StartTime) <= end.Unix() && int64(t.StartTime) >= start.Unix() {
			StartCloudBrainPointDeductTask(t)
		}
		if int64(t.EndTime) <= end.Unix() && int64(t.EndTime) >= start.Unix() {
			StopCloudBrainPointDeductTask(t)
		}
	}
}