feat(upload): Added the function of cleaning up orphan shard upload tasks and optimized temporary file management

Suyunmeng · Suyunmeng · commit a58e939d4c5f · 2025-09-05T23:54:06.000+08:00
diff --git a/internal/db/db.go b/internal/db/db.go
@@ -17,6 +17,11 @@ func Init(d *gorm.DB) {
 	if err != nil {
 		log.Fatalf("failed migrate database: %s", err.Error())
 	}
+	
+	// 清理启动前遗留的孤儿分片上传任务
+	if err := CleanupOrphanedSliceUploads(); err != nil {
+		log.Errorf("Failed to cleanup orphaned slice uploads: %v", err)
+	}
 }
 
 func AutoMigrate(dst ...interface{}) error {
diff --git a/internal/db/slice_upload.go b/internal/db/slice_upload.go
@@ -1,9 +1,17 @@
 package db
 
 import (
+	"io/fs"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+	
 	"github.com/OpenListTeam/OpenList/v4/internal/model/tables"
+	"github.com/OpenListTeam/OpenList/v4/pkg/tempdir"
 	"github.com/pkg/errors"
 	"gorm.io/gorm"
+	log "github.com/sirupsen/logrus"
 )
 
 func CreateSliceUpload(su *tables.SliceUpload) error {
@@ -58,3 +66,130 @@ func UpdateSliceStatusAtomic(taskID string, sliceNum int, status []byte) error {
 		return tx.Save(&su).Error
 	}))
 }
+
+// CleanupOrphanedSliceUploads 清理孤儿分片上传记录（启动时调用）
+func CleanupOrphanedSliceUploads() error {
+	// 清理超过24小时的未完成任务
+	cutoff := time.Now().Add(-24 * time.Hour)
+	
+	var orphanedTasks []tables.SliceUpload
+	if err := db.Where("status IN (?, ?) AND updated_at < ?", 
+		tables.SliceUploadStatusWaiting, 
+		tables.SliceUploadStatusUploading, 
+		cutoff).Find(&orphanedTasks).Error; err != nil {
+		return errors.WithStack(err)
+	}
+	
+	cleanedCount := 0
+	for _, task := range orphanedTasks {
+		// 清理临时文件
+		if task.TmpFile != "" {
+			if err := os.Remove(task.TmpFile); err != nil && !os.IsNotExist(err) {
+				log.Warnf("Failed to remove orphaned tmp file %s: %v", task.TmpFile, err)
+			} else if err == nil {
+				log.Debugf("Removed orphaned tmp file: %s", task.TmpFile)
+			}
+		}
+		
+		// 删除数据库记录
+		if err := db.Delete(&task).Error; err != nil {
+			log.Errorf("Failed to delete orphaned slice upload task %s: %v", task.TaskID, err)
+		} else {
+			cleanedCount++
+		}
+	}
+	
+	if cleanedCount > 0 {
+		log.Infof("Cleaned up %d orphaned slice upload tasks", cleanedCount)
+	}
+	
+	// 额外清理：扫描临时目录中的孤儿文件
+	return cleanupOrphanedTempFiles()
+}
+
+// cleanupOrphanedTempFiles 清理临时目录中的孤儿文件
+func cleanupOrphanedTempFiles() error {
+	// 获取临时目录路径，使用共享的tempdir包
+	tempDir := tempdir.GetPersistentTempDir()
+	
+	// 检查临时目录是否存在
+	if _, err := os.Stat(tempDir); os.IsNotExist(err) {
+		log.Debugf("Temp directory does not exist: %s", tempDir)
+		return nil
+	}
+	
+	// 获取所有活跃的分片上传任务的临时文件列表
+	var activeTasks []tables.SliceUpload
+	if err := db.Where("tmp_file IS NOT NULL AND tmp_file != '' AND status IN (?, ?)", 
+		tables.SliceUploadStatusWaiting, 
+		tables.SliceUploadStatusUploading).Find(&activeTasks).Error; err != nil {
+		return errors.WithStack(err)
+	}
+	
+	// 构建活跃文件的映射表
+	activeFiles := make(map[string]bool)
+	for _, task := range activeTasks {
+		if task.TmpFile != "" {
+			activeFiles[task.TmpFile] = true
+		}
+	}
+	
+	cleanedCount := 0
+	cutoff := time.Now().Add(-24 * time.Hour) // 只清理超过24小时的文件
+	
+	// 遍历临时目录
+	err := filepath.WalkDir(tempDir, func(path string, d fs.DirEntry, err error) error {
+		if err != nil {
+			log.Warnf("Failed to access path %s: %v", path, err)
+			return nil // 继续处理其他文件
+		}
+		
+		// 跳过目录
+		if d.IsDir() {
+			return nil
+		}
+		
+		// 只处理分片上传临时文件（以slice_upload_开头）
+		if !strings.HasPrefix(d.Name(), "slice_upload_") {
+			return nil
+		}
+		
+		// 检查文件是否在活跃任务列表中
+		if activeFiles[path] {
+			return nil // 文件仍在使用中，跳过
+		}
+		
+		// 检查文件修改时间
+		info, err := d.Info()
+		if err != nil {
+			log.Warnf("Failed to get file info for %s: %v", path, err)
+			return nil
+		}
+		
+		// 只清理超过24小时的文件
+		if info.ModTime().After(cutoff) {
+			return nil
+		}
+		
+		// 删除孤儿文件
+		if err := os.Remove(path); err != nil {
+			log.Warnf("Failed to remove orphaned temp file %s: %v", path, err)
+		} else {
+			log.Debugf("Removed orphaned temp file: %s", path)
+			cleanedCount++
+		}
+		
+		return nil
+	})
+	
+	if err != nil {
+		log.Errorf("Failed to walk temp directory %s: %v", tempDir, err)
+		return errors.WithStack(err)
+	}
+	
+	if cleanedCount > 0 {
+		log.Infof("Cleaned up %d orphaned temp files from %s", cleanedCount, tempDir)
+	}
+	
+	return nil
+}
diff --git a/internal/fs/sliceup.go b/internal/fs/sliceup.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"mime/multipart"
 	"os"
+	"path/filepath"
 	"strings"
 	"sync"
 	"time"
@@ -18,6 +19,7 @@ import (
 	"github.com/OpenListTeam/OpenList/v4/internal/op"
 	"github.com/OpenListTeam/OpenList/v4/internal/stream"
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
+	"github.com/OpenListTeam/OpenList/v4/pkg/tempdir"
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"gorm.io/gorm"
@@ -39,12 +41,9 @@ type SliceUploadSession struct {
 
 // NewSliceUploadManager 创建分片上传管理器
 func NewSliceUploadManager() *SliceUploadManager {
-	tempDir := os.TempDir() // 默认使用系统临时目录
-	if conf.Conf != nil && conf.Conf.TempDir != "" {
-		tempDir = conf.Conf.TempDir
-	}
+	tempDirPath := tempdir.GetPersistentTempDir()
 	return &SliceUploadManager{
-		tempDir: tempDir,
+		tempDir: tempDirPath,
 	}
 }
 
@@ -74,14 +73,40 @@ func (m *SliceUploadManager) CreateSession(ctx context.Context, storage driver.D
 	}
 
 	if su.ID != 0 { // 找到未完成的上传任务，支持断点续传
-		session := &SliceUploadSession{SliceUpload: su}
-		m.cache.Store(su.TaskID, session)
-		return &reqres.PreupResp{
-			TaskID:            su.TaskID,
-			SliceSize:         su.SliceSize,
-			SliceCnt:          su.SliceCnt,
-			SliceUploadStatus: su.SliceUploadStatus,
-		}, nil
+		// 验证临时文件是否仍然存在（重启后可能被清理）
+		if su.TmpFile != "" {
+			if _, err := os.Stat(su.TmpFile); os.IsNotExist(err) {
+				// 临时文件丢失，清理数据库记录，重新开始
+				log.Warnf("Temporary file lost after restart, cleaning up task: %s", su.TaskID)
+				if deleteErr := db.DeleteSliceUploadByTaskID(su.TaskID); deleteErr != nil {
+					log.Errorf("Failed to delete lost slice upload task: %v", deleteErr)
+				}
+				// 继续创建新任务
+			} else {
+				// 临时文件存在，可以继续断点续传
+				session := &SliceUploadSession{SliceUpload: su}
+				m.cache.Store(su.TaskID, session)
+				log.Infof("Resuming slice upload after restart: %s, completed slices: %d/%d", 
+					su.TaskID, tables.CountUploadedSlices(su.SliceUploadStatus), su.SliceCnt)
+				return &reqres.PreupResp{
+					TaskID:            su.TaskID,
+					SliceSize:         su.SliceSize,
+					SliceCnt:          su.SliceCnt,
+					SliceUploadStatus: su.SliceUploadStatus,
+				}, nil
+			}
+		} else {
+			// 原生分片上传（如123open/baidu），无需临时文件
+			session := &SliceUploadSession{SliceUpload: su}
+			m.cache.Store(su.TaskID, session)
+			log.Infof("Resuming native slice upload after restart: %s", su.TaskID)
+			return &reqres.PreupResp{
+				TaskID:            su.TaskID,
+				SliceSize:         su.SliceSize,
+				SliceCnt:          su.SliceCnt,
+				SliceUploadStatus: su.SliceUploadStatus,
+			}, nil
+		}
 	}
 	
 	srcobj, err := op.Get(ctx, storage, actualPath)
@@ -419,39 +444,48 @@ func (m *SliceUploadManager) CompleteUpload(ctx context.Context, storage driver.
 	}
 }
 
-// ensureTmpFile 确保临时文件存在且正确初始化，线程安全 - 保持原始实现
+// ensureTmpFile 确保临时文件存在且正确初始化，线程安全 - 使用持久化目录
 func (s *SliceUploadSession) ensureTmpFile() error {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 	
 	if s.TmpFile == "" {
-		tempDir := os.TempDir() // 默认使用系统临时目录
-		if conf.Conf != nil && conf.Conf.TempDir != "" {
-			tempDir = conf.Conf.TempDir
-		}
-		tf, err := os.CreateTemp(tempDir, "file-*")
+		tempDirPath := tempdir.GetPersistentTempDir()
+		
+		// 使用TaskID作为文件名的一部分，确保唯一性和可识别性
+		filename := fmt.Sprintf("slice_upload_%s_%s", s.TaskID, s.Name)
+		// 清理文件名中的特殊字符
+		filename = strings.ReplaceAll(filename, "/", "_")
+		filename = strings.ReplaceAll(filename, "\\", "_")
+		filename = strings.ReplaceAll(filename, ":", "_")
+		
+		tmpPath := filepath.Join(tempDirPath, filename)
+		
+		tf, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_RDWR, 0644)
 		if err != nil {
-			return fmt.Errorf("CreateTemp error: %w", err)
+			return fmt.Errorf("create persistent temp file error: %w", err)
 		}
 		
-		abspath := tf.Name() //这里返回的是绝对路径
-		if err = os.Truncate(abspath, int64(s.Size)); err != nil {
+		if err = os.Truncate(tmpPath, int64(s.Size)); err != nil {
 			tf.Close() // 确保文件被关闭
-			os.Remove(abspath) // 清理文件
-			return fmt.Errorf("Truncate error: %w", err)
+			os.Remove(tmpPath) // 清理文件
+			return fmt.Errorf("truncate persistent temp file error: %w", err)
 		}
 		
-		s.TmpFile = abspath
+		s.TmpFile = tmpPath
 		s.tmpFile = tf
+		
+		log.Debugf("Created persistent temp file: %s", tmpPath)
 		return nil
 	}
 	
 	if s.tmpFile == nil {
 		var err error
 		s.tmpFile, err = os.OpenFile(s.TmpFile, os.O_RDWR, 0644)
 		if err != nil {
-			return fmt.Errorf("OpenFile error: %w", err)
+			return fmt.Errorf("reopen persistent temp file error: %w", err)
 		}
+		log.Debugf("Reopened persistent temp file: %s", s.TmpFile)
 	}
 	return nil
 }
diff --git a/internal/model/tables/slice_upload.go b/internal/model/tables/slice_upload.go
@@ -60,3 +60,14 @@ func IsAllSliceUploaded(status []byte, sliceCnt uint) bool {
 	}
 	return true
 }
+
+// CountUploadedSlices 统计已上传的分片数量
+func CountUploadedSlices(status []byte) uint {
+	count := uint(0)
+	for i := 0; i < len(status)*8; i++ {
+		if status[i/8]&(1<<(i%8)) != 0 {
+			count++
+		}
+	}
+	return count
+}
diff --git a/pkg/tempdir/tempdir.go b/pkg/tempdir/tempdir.go
@@ -0,0 +1,35 @@
+package tempdir
+
+import (
+	"os"
+	"path/filepath"
+	"github.com/OpenListTeam/OpenList/v4/internal/conf"
+)
+
+// GetPersistentTempDir 获取持久化临时目录
+// 这个函数被多个包共享使用，避免代码重复
+func GetPersistentTempDir() string {
+	var tempDir string
+	
+	// 优先使用配置的临时目录
+	if conf.Conf != nil && conf.Conf.TempDir != "" {
+		tempDir = conf.Conf.TempDir
+	} else {
+		// 使用数据目录下的slice_temp子目录
+		if conf.Conf != nil && conf.Conf.Database.DBFile != "" {
+			// 从数据库文件路径推断数据目录
+			dataDir := filepath.Dir(conf.Conf.Database.DBFile)
+			tempDir = filepath.Join(dataDir, "slice_temp")
+		} else {
+			// fallback到当前工作目录下的slice_temp
+			if wd, err := os.Getwd(); err == nil {
+				tempDir = filepath.Join(wd, "slice_temp")
+			} else {
+				// 最后的fallback
+				tempDir = filepath.Join(os.TempDir(), "openlist_slice_temp")
+			}
+		}
+	}
+	
+	return tempDir
+}

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,11 @@ func Init(d *gorm.DB) {`
`17`	`17`	`if err != nil {`
`18`	`18`	`log.Fatalf("failed migrate database: %s", err.Error())`
`19`	`19`	`}`
	`20`	`+`
	`21`	`+ // 清理启动前遗留的孤儿分片上传任务`
	`22`	`+ if err := CleanupOrphanedSliceUploads(); err != nil {`
	`23`	`+ log.Errorf("Failed to cleanup orphaned slice uploads: %v", err)`
	`24`	`+ }`
`20`	`25`	`}`
`21`	`26`
`22`	`27`	`func AutoMigrate(dst ...interface{}) error {`