276 lines
9.4 KiB
Go
276 lines
9.4 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"go.uber.org/zap"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
ImageTaskStatusProcessing = "processing"
|
||
|
|
ImageTaskStatusCompleted = "completed"
|
||
|
|
ImageTaskStatusFailed = "failed"
|
||
|
|
|
||
|
|
defaultImageTaskTTL = 24 * time.Hour
|
||
|
|
defaultImageTaskExecutionTimeout = 30 * time.Minute
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
ErrImageTaskNotFound = infraerrors.New(http.StatusNotFound, "IMAGE_TASK_NOT_FOUND", "image task not found")
|
||
|
|
ErrImageTaskForbidden = infraerrors.New(http.StatusForbidden, "IMAGE_TASK_FORBIDDEN", "image task does not belong to this API key")
|
||
|
|
ErrImageTaskUnavailable = infraerrors.New(http.StatusServiceUnavailable, "IMAGE_TASK_UNAVAILABLE", "image task storage is unavailable")
|
||
|
|
)
|
||
|
|
|
||
|
|
// ImageTaskRecord is the private Redis representation of an asynchronous image
|
||
|
|
// request. Ownership fields are intentionally omitted from the public view.
|
||
|
|
type ImageTaskRecord struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
UserID int64 `json:"user_id"`
|
||
|
|
APIKeyID int64 `json:"api_key_id"`
|
||
|
|
Status string `json:"status"`
|
||
|
|
HTTPStatus int `json:"http_status,omitempty"`
|
||
|
|
Result json.RawMessage `json:"result,omitempty"`
|
||
|
|
Error json.RawMessage `json:"error,omitempty"`
|
||
|
|
CreatedAt int64 `json:"created_at"`
|
||
|
|
CompletedAt *int64 `json:"completed_at,omitempty"`
|
||
|
|
ExpiresAt int64 `json:"expires_at"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// ImageTask is the API-safe task representation returned to callers.
|
||
|
|
type ImageTask struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
TaskID string `json:"task_id"`
|
||
|
|
Object string `json:"object"`
|
||
|
|
Status string `json:"status"`
|
||
|
|
HTTPStatus int `json:"http_status,omitempty"`
|
||
|
|
ImageURL string `json:"image_url,omitempty"`
|
||
|
|
Result json.RawMessage `json:"result,omitempty"`
|
||
|
|
Error json.RawMessage `json:"error,omitempty"`
|
||
|
|
CreatedAt int64 `json:"created_at"`
|
||
|
|
CompletedAt *int64 `json:"completed_at,omitempty"`
|
||
|
|
ExpiresAt int64 `json:"expires_at"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type ImageTaskOwner struct {
|
||
|
|
UserID int64
|
||
|
|
APIKeyID int64
|
||
|
|
}
|
||
|
|
|
||
|
|
type ImageTaskStore interface {
|
||
|
|
Save(ctx context.Context, task *ImageTaskRecord, ttl time.Duration) error
|
||
|
|
Get(ctx context.Context, id string) (*ImageTaskRecord, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
// ImageStorageResolver reports the currently effective object-storage binding.
|
||
|
|
// It exists so the async image feature can be switched on and off from the admin
|
||
|
|
// UI without a restart: the wiring below is fixed at startup, but the answer to
|
||
|
|
// "is object storage configured right now" is re-read (and cached) per call.
|
||
|
|
type ImageStorageResolver func() (uploader *ImageResultUploader, enabled bool)
|
||
|
|
|
||
|
|
type ImageTaskService struct {
|
||
|
|
store ImageTaskStore
|
||
|
|
uploader *ImageResultUploader
|
||
|
|
enabled bool
|
||
|
|
resolve ImageStorageResolver
|
||
|
|
ttl time.Duration
|
||
|
|
executionTimeout time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewImageTaskService(store ImageTaskStore) *ImageTaskService {
|
||
|
|
return NewImageTaskServiceWithOptions(store, defaultImageTaskTTL, defaultImageTaskExecutionTimeout)
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewImageTaskServiceWithOptions(store ImageTaskStore, ttl, executionTimeout time.Duration) *ImageTaskService {
|
||
|
|
if ttl <= 0 {
|
||
|
|
ttl = defaultImageTaskTTL
|
||
|
|
}
|
||
|
|
if executionTimeout <= 0 {
|
||
|
|
executionTimeout = defaultImageTaskExecutionTimeout
|
||
|
|
}
|
||
|
|
return &ImageTaskService{store: store, ttl: ttl, executionTimeout: executionTimeout}
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewImageTaskServiceWithUploader 构造一个已启用的图片任务服务：结果会先经 uploader
|
||
|
|
// 转存到对象存储再落 Redis。uploader 为 nil 时不做转存（仅用于测试）。
|
||
|
|
func NewImageTaskServiceWithUploader(store ImageTaskStore, uploader *ImageResultUploader, ttl, executionTimeout time.Duration) *ImageTaskService {
|
||
|
|
s := NewImageTaskServiceWithOptions(store, ttl, executionTimeout)
|
||
|
|
s.uploader = uploader
|
||
|
|
s.enabled = true
|
||
|
|
return s
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewImageTaskServiceWithResolver 构造一个由 resolver 决定启用状态的服务：
|
||
|
|
// 开关与凭证来自后台设置，保存后立即生效，无需重启。
|
||
|
|
func NewImageTaskServiceWithResolver(store ImageTaskStore, resolve ImageStorageResolver, ttl, executionTimeout time.Duration) *ImageTaskService {
|
||
|
|
s := NewImageTaskServiceWithOptions(store, ttl, executionTimeout)
|
||
|
|
s.resolve = resolve
|
||
|
|
return s
|
||
|
|
}
|
||
|
|
|
||
|
|
// current 返回当前生效的 uploader 与启用状态。
|
||
|
|
// 注入了 resolver 时以 resolver 为准（后台设置可热切换），否则回落到构造时固定的值。
|
||
|
|
func (s *ImageTaskService) current() (*ImageResultUploader, bool) {
|
||
|
|
if s == nil {
|
||
|
|
return nil, false
|
||
|
|
}
|
||
|
|
if s.resolve != nil {
|
||
|
|
return s.resolve()
|
||
|
|
}
|
||
|
|
return s.uploader, s.enabled
|
||
|
|
}
|
||
|
|
|
||
|
|
// Enabled 表示异步图片任务功能是否可用（总开关 + 凭证齐全）。
|
||
|
|
// 关闭时 handler 直接返回 404，不创建任务、不写 Redis。
|
||
|
|
func (s *ImageTaskService) Enabled() bool {
|
||
|
|
if s == nil || s.store == nil {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
_, enabled := s.current()
|
||
|
|
return enabled
|
||
|
|
}
|
||
|
|
|
||
|
|
// Pollable 表示已创建的任务能否被查询。
|
||
|
|
// 比 Enabled 弱：只要 store 可用即可，从而在功能被关掉后仍能取回进行中的任务结果。
|
||
|
|
func (s *ImageTaskService) Pollable() bool {
|
||
|
|
return s != nil && s.store != nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ImageTaskService) ExecutionTimeout() time.Duration {
|
||
|
|
if s == nil || s.executionTimeout <= 0 {
|
||
|
|
return defaultImageTaskExecutionTimeout
|
||
|
|
}
|
||
|
|
return s.executionTimeout
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ImageTaskService) Create(ctx context.Context, owner ImageTaskOwner) (*ImageTask, error) {
|
||
|
|
if s == nil || s.store == nil {
|
||
|
|
return nil, ErrImageTaskUnavailable
|
||
|
|
}
|
||
|
|
now := time.Now().UTC()
|
||
|
|
task := &ImageTaskRecord{
|
||
|
|
ID: "imgtask_" + strings.ReplaceAll(uuid.NewString(), "-", ""),
|
||
|
|
UserID: owner.UserID,
|
||
|
|
APIKeyID: owner.APIKeyID,
|
||
|
|
Status: ImageTaskStatusProcessing,
|
||
|
|
CreatedAt: now.Unix(),
|
||
|
|
ExpiresAt: now.Add(s.ttl).Unix(),
|
||
|
|
}
|
||
|
|
if err := s.store.Save(ctx, task, s.ttl); err != nil {
|
||
|
|
return nil, ErrImageTaskUnavailable.WithCause(err)
|
||
|
|
}
|
||
|
|
return imageTaskToPublic(task), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ImageTaskService) Get(ctx context.Context, owner ImageTaskOwner, id string) (*ImageTask, error) {
|
||
|
|
if s == nil || s.store == nil {
|
||
|
|
return nil, ErrImageTaskUnavailable
|
||
|
|
}
|
||
|
|
task, err := s.store.Get(ctx, strings.TrimSpace(id))
|
||
|
|
if err != nil {
|
||
|
|
if errors.Is(err, ErrImageTaskNotFound) {
|
||
|
|
return nil, ErrImageTaskNotFound
|
||
|
|
}
|
||
|
|
return nil, ErrImageTaskUnavailable.WithCause(err)
|
||
|
|
}
|
||
|
|
if task.UserID != owner.UserID || task.APIKeyID != owner.APIKeyID {
|
||
|
|
// Do not reveal whether a random task ID exists for another caller.
|
||
|
|
return nil, ErrImageTaskNotFound
|
||
|
|
}
|
||
|
|
return imageTaskToPublic(task), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ImageTaskService) Complete(ctx context.Context, id string, statusCode int, result json.RawMessage) error {
|
||
|
|
if !json.Valid(result) {
|
||
|
|
return s.Fail(ctx, id, http.StatusBadGateway, imageTaskErrorJSON("api_error", "upstream returned a non-JSON image response"))
|
||
|
|
}
|
||
|
|
if uploader, _ := s.current(); uploader != nil {
|
||
|
|
rewritten, err := uploader.Rewrite(ctx, id, result)
|
||
|
|
if err != nil {
|
||
|
|
// 转存失败不回退存 base64，避免大 blob 撑爆 Redis：直接把任务标记为失败。
|
||
|
|
logger.L().Error("image_task.offload_failed", zap.String("task_id", id), zap.Error(err))
|
||
|
|
return s.Fail(ctx, id, http.StatusBadGateway, imageTaskErrorJSON("api_error", "failed to store generated image to object storage"))
|
||
|
|
}
|
||
|
|
result = rewritten
|
||
|
|
}
|
||
|
|
return s.finish(ctx, id, ImageTaskStatusCompleted, statusCode, result, nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ImageTaskService) Fail(ctx context.Context, id string, statusCode int, taskErr json.RawMessage) error {
|
||
|
|
if !json.Valid(taskErr) {
|
||
|
|
taskErr = imageTaskErrorJSON("api_error", "image generation failed")
|
||
|
|
}
|
||
|
|
return s.finish(ctx, id, ImageTaskStatusFailed, statusCode, nil, taskErr)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ImageTaskService) finish(ctx context.Context, id, status string, statusCode int, result, taskErr json.RawMessage) error {
|
||
|
|
if s == nil || s.store == nil {
|
||
|
|
return ErrImageTaskUnavailable
|
||
|
|
}
|
||
|
|
task, err := s.store.Get(ctx, id)
|
||
|
|
if err != nil {
|
||
|
|
if errors.Is(err, ErrImageTaskNotFound) {
|
||
|
|
return ErrImageTaskNotFound
|
||
|
|
}
|
||
|
|
return ErrImageTaskUnavailable.WithCause(err)
|
||
|
|
}
|
||
|
|
now := time.Now().UTC()
|
||
|
|
completedAt := now.Unix()
|
||
|
|
task.Status = status
|
||
|
|
task.HTTPStatus = statusCode
|
||
|
|
task.Result = result
|
||
|
|
task.Error = taskErr
|
||
|
|
task.CompletedAt = &completedAt
|
||
|
|
task.ExpiresAt = now.Add(s.ttl).Unix()
|
||
|
|
if err := s.store.Save(ctx, task, s.ttl); err != nil {
|
||
|
|
return ErrImageTaskUnavailable.WithCause(err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func imageTaskToPublic(task *ImageTaskRecord) *ImageTask {
|
||
|
|
if task == nil {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return &ImageTask{
|
||
|
|
ID: task.ID,
|
||
|
|
TaskID: task.ID,
|
||
|
|
Object: "image.generation.task",
|
||
|
|
Status: task.Status,
|
||
|
|
HTTPStatus: task.HTTPStatus,
|
||
|
|
ImageURL: firstImageTaskURL(task.Result),
|
||
|
|
Result: task.Result,
|
||
|
|
Error: task.Error,
|
||
|
|
CreatedAt: task.CreatedAt,
|
||
|
|
CompletedAt: task.CompletedAt,
|
||
|
|
ExpiresAt: task.ExpiresAt,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func firstImageTaskURL(result json.RawMessage) string {
|
||
|
|
if len(result) == 0 || !json.Valid(result) {
|
||
|
|
return ""
|
||
|
|
}
|
||
|
|
var response struct {
|
||
|
|
Data []struct {
|
||
|
|
URL string `json:"url"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
if json.Unmarshal(result, &response) != nil || len(response.Data) == 0 {
|
||
|
|
return ""
|
||
|
|
}
|
||
|
|
return strings.TrimSpace(response.Data[0].URL)
|
||
|
|
}
|
||
|
|
|
||
|
|
func imageTaskErrorJSON(errorType, message string) json.RawMessage {
|
||
|
|
data, _ := json.Marshal(map[string]string{"type": errorType, "message": message})
|
||
|
|
return data
|
||
|
|
}
|