api/controller/chunk_upload_manager.go

Rama main · réplica de lectura del repositorio original.

package controller

import (
	"fmt"
	"mime/multipart"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/goodrain/rainbond/db"
	"github.com/goodrain/rainbond/db/model"
	"github.com/goodrain/rainbond/pkg/component/storage"
	"github.com/google/uuid"
	"github.com/sirupsen/logrus"
)

// Upload chunk sizing and session lifetime limits.
const (
	DefaultChunkSize     = 5 * 1024 * 1024  // 5MB
	MinChunkSize         = 1 * 1024 * 1024  // 1MB
	MaxChunkSize         = 20 * 1024 * 1024 // 20MB
	SessionExpiryHours   = 24               // 24小时过期
	MaxConcurrentUploads = 10               // 最大并发上传数
)

// ChunkUploadManager 分片上传管理器
type ChunkUploadManager struct {
	admitUse        func(*model.UploadSession, string) (func(bool) error, error)
	cleanupChunks   func(string) error
	mergeChunks     func(string, string, int) error
	sessionCache    map[string]*model.UploadSession
	cacheMutex      sync.RWMutex
	uploadSemaphore chan struct{}
	cleanupTicker   *time.Ticker
}

var chunkUploadManager *ChunkUploadManager
var chunkUploadOnce sync.Once

// GetChunkUploadManager 获取分片上传管理器单例
func GetChunkUploadManager() *ChunkUploadManager {
	chunkUploadOnce.Do(func() {
		chunkUploadManager = &ChunkUploadManager{
			sessionCache:    make(map[string]*model.UploadSession),
			uploadSemaphore: make(chan struct{}, MaxConcurrentUploads),
			cleanupTicker:   time.NewTicker(1 * time.Hour),
		}
		// 启动清理协程
		go chunkUploadManager.startCleanupWorker()
	})
	return chunkUploadManager
}

// InitUploadSession 初始化上传会话
func (m *ChunkUploadManager) InitUploadSession(eventID, fileName string, fileSize int64, fileMD5 string, chunkSize int) (*model.UploadSession, error) {
	// 验证分片大小
	if chunkSize == 0 {
		chunkSize = DefaultChunkSize
	}
	if chunkSize < MinChunkSize || chunkSize > MaxChunkSize {
		return nil, fmt.Errorf("chunk size must be between %d and %d", MinChunkSize, MaxChunkSize)
	}

	// 检查是否已存在未完成的上传会话
	existingSession, err := db.GetManager().UploadSessionDao().GetByEventIDAndFileName(eventID, fileName)
	if err != nil {
		return nil, fmt.Errorf("failed to check existing session: %v", err)
	}

	if existingSession != nil && existingSession.Status == "uploading" && existingSession.ExpiresAt.After(time.Now()) {
		// 返回现有会话（支持断点续传）
		logrus.Infof("Resume existing upload session: %s", existingSession.ID)
		m.cacheSession(existingSession)
		return existingSession, nil
	}

	// 创建新的上传会话
	sessionID := uuid.New().String()
	totalChunks := int((fileSize + int64(chunkSize) - 1) / int64(chunkSize))

	session := &model.UploadSession{
		ID:             sessionID,
		EventID:        eventID,
		FileName:       fileName,
		FileSize:       fileSize,
		FileMD5:        fileMD5,
		ChunkSize:      chunkSize,
		TotalChunks:    totalChunks,
		UploadedChunks: "",
		Status:         "uploading",
		StoragePath:    fmt.Sprintf("/grdata/package_build/temp/events/%s/%s", eventID, fileName),
		CreatedAt:      time.Now(),
		UpdatedAt:      time.Now(),
		ExpiresAt:      time.Now().Add(SessionExpiryHours * time.Hour),
	}

	if err := db.GetManager().UploadSessionDao().AddModel(session); err != nil {
		return nil, fmt.Errorf("failed to create session: %v", err)
	}

	m.cacheSession(session)
	logrus.Infof("Created new upload session: %s, total chunks: %d", sessionID, totalChunks)
	return session, nil
}

// SaveChunk 保存分片
func (m *ChunkUploadManager) SaveChunk(sessionID string, chunkIndex int, reader multipart.File) (resultErr error) {
	// 获取会话
	session, err := m.getSession(sessionID)
	if err != nil {
		return err
	}

	if !session.ExpiresAt.After(time.Now()) {
		return fmt.Errorf("upload session expired")
	}
	// 验证会话状态
	if session.Status != "uploading" {
		return fmt.Errorf("session status is %s, cannot upload", session.Status)
	}

	// 验证分片索引
	if chunkIndex < 0 || chunkIndex >= session.TotalChunks {
		return fmt.Errorf("invalid chunk index: %d, total chunks: %d", chunkIndex, session.TotalChunks)
	}

	finish, err := m.admitSessionUse(session, "chunk")
	if err != nil {
		return err
	}
	confirmed := false
	defer func() {
		if err := finish(confirmed); resultErr == nil {
			resultErr = err
		}
	}()

	// 检查分片是否已存在（幂等性）
	if storage.Default().StorageCli.ChunkExists(sessionID, chunkIndex) {
		logrus.Debugf("Chunk %d already exists for session %s, skipping", chunkIndex, sessionID)
		confirmed = true
		return nil
	}

	// 并发控制
	m.uploadSemaphore <- struct{}{}
	defer func() { <-m.uploadSemaphore }()

	// 保存分片
	_, err = storage.Default().StorageCli.SaveChunk(sessionID, chunkIndex, reader)
	if err != nil {
		return fmt.Errorf("failed to save chunk: %v", err)
	}

	// 更新已上传分片列表
	if err := m.updateUploadedChunks(session, chunkIndex); err != nil {
		return err
	}

	confirmed = true
	logrus.Infof("Saved chunk %d/%d for session %s", chunkIndex+1, session.TotalChunks, sessionID)
	return nil
}

// CompleteUpload 完成上传，合并所有分片
func (m *ChunkUploadManager) CompleteUpload(sessionID string) (resultPath string, resultErr error) {
	session, err := m.getSession(sessionID)
	if err != nil {
		return "", err
	}

	// A completed upload is an immutable result, not permission to merge again.
	if session.Status == "completed" {
		return session.StoragePath, nil
	}
	if session.Status != "uploading" {
		return "", fmt.Errorf("session status is %s, cannot complete", session.Status)
	}

	if !session.ExpiresAt.After(time.Now()) {
		return "", fmt.Errorf("upload session expired")
	}
	// 验证所有分片是否已上传
	uploadedChunks := m.parseUploadedChunks(session.UploadedChunks)
	if len(uploadedChunks) != session.TotalChunks {
		missingChunks := m.getMissingChunks(uploadedChunks, session.TotalChunks)
		return "", fmt.Errorf("not all chunks uploaded, missing: %v", missingChunks)
	}

	finish, err := m.admitSessionUse(session, "complete")
	if err != nil {
		return "", err
	}
	confirmed := false
	defer func() {
		if err := finish(confirmed); resultErr == nil {
			resultErr = err
		}
	}()

	// 合并分片
	logrus.Infof("Merging %d chunks for session %s", session.TotalChunks, sessionID)
	err = m.mergeSessionChunks(sessionID, session.StoragePath, session.TotalChunks)
	if err != nil {
		session.Status = "failed"
		db.GetManager().UploadSessionDao().UpdateModel(session)
		return "", fmt.Errorf("failed to merge chunks: %v", err)
	}

	// Publish durable completion before removing the only resumable chunk set.
	// Do not mutate the cached record if the database cannot acknowledge it.
	completed := *session
	completed.Status = "completed"
	completed.UpdatedAt = time.Now()
	if err := db.GetManager().UploadSessionDao().UpdateModel(&completed); err != nil {
		return "", fmt.Errorf("failed to persist upload completion: %w", err)
	}
	m.cacheSession(&completed)
	if err := m.cleanupSessionChunks(sessionID); err != nil {
		// The package is complete; retain its durable record for later chunk cleanup.
		logrus.Warnf("Failed to cleanup chunks for session %s: %v", sessionID, err)
	} else {
		confirmed = true
	}

	m.removeFromCache(sessionID)
	logrus.Infof("Upload completed for session %s, file: %s", sessionID, session.StoragePath)
	return session.StoragePath, nil
}

// GetUploadStatus 获取上传状态
func (m *ChunkUploadManager) GetUploadStatus(sessionID string) (*UploadStatusResponse, error) {
	session, err := m.getSession(sessionID)
	if err != nil {
		return nil, err
	}

	uploadedChunks := m.parseUploadedChunks(session.UploadedChunks)
	missingChunks := m.getMissingChunks(uploadedChunks, session.TotalChunks)
	progress := float64(len(uploadedChunks)) / float64(session.TotalChunks) * 100

	return &UploadStatusResponse{
		SessionID:      session.ID,
		FileName:       session.FileName,
		FileSize:       session.FileSize,
		ChunkSize:      session.ChunkSize,
		TotalChunks:    session.TotalChunks,
		UploadedChunks: uploadedChunks,
		MissingChunks:  missingChunks,
		Progress:       progress,
		Status:         session.Status,
		CreatedAt:      session.CreatedAt,
		UpdatedAt:      session.UpdatedAt,
	}, nil
}

// CancelUpload 取消上传
func (m *ChunkUploadManager) CancelUpload(sessionID string) (resultErr error) {
	session, err := m.getSession(sessionID)
	if err != nil {
		return err
	}

	finish, err := m.admitSessionUse(session, "cancel")
	if err != nil {
		return err
	}
	confirmed := false
	defer func() {
		if err := finish(confirmed); resultErr == nil {
			resultErr = err
		}
	}()

	// 清理分片文件
	if err := m.cleanupSessionChunks(sessionID); err != nil {
		return fmt.Errorf("failed to cleanup upload chunks: %w", err)
	}

	// 删除会话记录
	if err := db.GetManager().UploadSessionDao().DeleteByID(sessionID); err != nil {
		return fmt.Errorf("failed to delete session: %v", err)
	}

	confirmed = true
	m.removeFromCache(sessionID)
	logrus.Infof("Cancelled upload session: %s", sessionID)
	return nil
}

func (m *ChunkUploadManager) mergeSessionChunks(sessionID, path string, count int) error {
	if m.mergeChunks != nil {
		return m.mergeChunks(sessionID, path, count)
	}
	configured := storage.Default()
	if configured == nil || configured.StorageCli == nil {
		return fmt.Errorf("upload storage unavailable")
	}
	return configured.StorageCli.MergeChunks(sessionID, path, count)
}

func (m *ChunkUploadManager) cleanupSessionChunks(sessionID string) error {
	if m.cleanupChunks != nil {
		return m.cleanupChunks(sessionID)
	}
	configured := storage.Default()
	if configured == nil || configured.StorageCli == nil {
		return fmt.Errorf("upload storage unavailable")
	}
	return configured.StorageCli.CleanupChunks(sessionID)
}

// 内部方法

// getSession 获取会话（优先从缓存）
func (m *ChunkUploadManager) getSession(sessionID string) (*model.UploadSession, error) {
	// 先从缓存获取
	m.cacheMutex.RLock()
	if session, ok := m.sessionCache[sessionID]; ok {
		m.cacheMutex.RUnlock()
		return session, nil
	}
	m.cacheMutex.RUnlock()

	// 从数据库获取
	session, err := db.GetManager().UploadSessionDao().GetByID(sessionID)
	if err != nil {
		return nil, fmt.Errorf("session not found: %v", err)
	}

	m.cacheSession(session)
	return session, nil
}

// cacheSession 缓存会话
func (m *ChunkUploadManager) cacheSession(session *model.UploadSession) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	m.sessionCache[session.ID] = session
}

// removeFromCache 从缓存移除
func (m *ChunkUploadManager) removeFromCache(sessionID string) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	delete(m.sessionCache, sessionID)
}

// updateUploadedChunks 更新已上传的分片列表
func (m *ChunkUploadManager) updateUploadedChunks(session *model.UploadSession, chunkIndex int) error {
	uploadedChunks := m.parseUploadedChunks(session.UploadedChunks)

	// 检查是否已记录
	for _, idx := range uploadedChunks {
		if idx == chunkIndex {
			return nil // 已存在
		}
	}

	// 添加新的分片索引
	uploadedChunks = append(uploadedChunks, chunkIndex)
	sort.Ints(uploadedChunks)

	// 转换为字符串
	strChunks := make([]string, len(uploadedChunks))
	for i, idx := range uploadedChunks {
		strChunks[i] = strconv.Itoa(idx)
	}
	session.UploadedChunks = strings.Join(strChunks, ",")
	session.UpdatedAt = time.Now()

	// 更新数据库
	if err := db.GetManager().UploadSessionDao().UpdateModel(session); err != nil {
		return fmt.Errorf("failed to update session: %v", err)
	}

	return nil
}

// parseUploadedChunks 解析已上传的分片列表
func (m *ChunkUploadManager) parseUploadedChunks(uploadedChunksStr string) []int {
	if uploadedChunksStr == "" {
		return []int{}
	}

	parts := strings.Split(uploadedChunksStr, ",")
	chunks := make([]int, 0, len(parts))
	for _, part := range parts {
		if idx, err := strconv.Atoi(strings.TrimSpace(part)); err == nil {
			chunks = append(chunks, idx)
		}
	}
	return chunks
}

// getMissingChunks 获取缺失的分片列表
func (m *ChunkUploadManager) getMissingChunks(uploadedChunks []int, totalChunks int) []int {
	uploaded := make(map[int]bool)
	for _, idx := range uploadedChunks {
		uploaded[idx] = true
	}

	missing := make([]int, 0)
	for i := 0; i < totalChunks; i++ {
		if !uploaded[i] {
			missing = append(missing, i)
		}
	}
	return missing
}

// startCleanupWorker 启动定时清理任务
func (m *ChunkUploadManager) startCleanupWorker() {
	for range m.cleanupTicker.C {
		if err := m.cleanExpiredSessions(); err != nil {
			logrus.Errorf("Failed to clean expired sessions: %v", err)
		}
	}
}

// cleanExpiredSessions 清理过期的会话
func (m *ChunkUploadManager) cleanExpiredSessions() error {
	// Preserve uncached records after a restart: expiry alone does not prove
	// their chunks were removed. Keep the existing cached-session cleanup scope.
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	now := time.Now()
	var first error
	for id, session := range m.sessionCache {
		if !session.ExpiresAt.Before(now) {
			continue
		}
		if err := m.expireSessionChunks(session); err != nil {
			if first == nil {
				first = fmt.Errorf("expired upload chunk cleanup failed: %w", err)
			}
			continue
		}
		// Expiration is not an explicit cancellation. Retain the durable row,
		// including a row another API instance may have completed meanwhile.
		delete(m.sessionCache, id)
	}
	return first
}

// UploadStatusResponse 上传状态响应
type UploadStatusResponse struct {
	SessionID      string    `json:"session_id"`
	FileName       string    `json:"file_name"`
	FileSize       int64     `json:"file_size"`
	ChunkSize      int       `json:"chunk_size"`
	TotalChunks    int       `json:"total_chunks"`
	UploadedChunks []int     `json:"uploaded_chunks"`
	MissingChunks  []int     `json:"missing_chunks"`
	Progress       float64   `json:"progress"`
	Status         string    `json:"status"`
	CreatedAt      time.Time `json:"created_at"`
	UpdatedAt      time.Time `json:"updated_at"`
}