package task import ( "context" "encoding/csv" "errors" "fmt" "os" "path/filepath" "time" "github.com/bytedance/sonic" "github.com/hibiken/asynq" "github.com/redis/go-redis/v9" "github.com/xuri/excelize/v2" "github.com/break/junhong_cmp_fiber/pkg/constants" "github.com/break/junhong_cmp_fiber/pkg/storage" ) const ( exportLockTTL = 5 * time.Minute exportShardLockTTL = 10 * time.Minute exportFinalizeLockTTL = 2 * time.Minute ) // ExportDispatchPayload 导出 dispatch 任务载荷。 type ExportDispatchPayload struct { TaskID uint `json:"task_id"` } // ExportShardPayload 导出 shard 任务载荷。 type ExportShardPayload struct { TaskID uint `json:"task_id"` ShardID uint `json:"shard_id"` } // ExportFinalizePayload 导出 finalize 任务载荷。 type ExportFinalizePayload struct { TaskID uint `json:"task_id"` } func enqueueTask(ctx context.Context, client *asynq.Client, taskType string, payload any, opts ...asynq.Option) error { if client == nil { return errors.New("asynq client 未初始化") } bytes, err := sonic.Marshal(payload) if err != nil { return err } task := asynq.NewTask(taskType, bytes, opts...) _, err = client.EnqueueContext(ctx, task) return err } func tryAcquireLock(ctx context.Context, rdb *redis.Client, key string, ttl time.Duration) (bool, error) { if rdb == nil { return true, nil } return rdb.SetNX(ctx, key, "1", ttl).Result() } func releaseLock(ctx context.Context, rdb *redis.Client, key string) { if rdb == nil { return } _, _ = rdb.Del(ctx, key).Result() } func isFinalRetry(ctx context.Context) bool { retryCount, ok := asynq.GetRetryCount(ctx) if !ok { return false } maxRetry, ok := asynq.GetMaxRetry(ctx) if !ok { return false } return retryCount >= maxRetry } func writeExportFile(format string, headers []string, rows [][]string) (string, int64, error) { switch format { case constants.ExportTaskFormatCSV: return writeCSVFile(headers, rows) case constants.ExportTaskFormatXLSX: return writeXLSXFile(headers, rows) default: return "", 0, fmt.Errorf("不支持的导出格式: %s", format) } } func writeCSVFile(headers []string, rows [][]string) (string, int64, error) { file, err := os.CreateTemp("", "export-*.csv") if err != nil { return "", 0, err } defer file.Close() writer := csv.NewWriter(file) if err := writer.Write(headers); err != nil { return "", 0, err } for _, row := range rows { if err := writer.Write(row); err != nil { return "", 0, err } } writer.Flush() if err := writer.Error(); err != nil { return "", 0, err } if err := file.Sync(); err != nil { return "", 0, err } stat, err := file.Stat() if err != nil { return "", 0, err } return file.Name(), stat.Size(), nil } func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) { file := excelize.NewFile() sheet := file.GetSheetName(file.GetActiveSheetIndex()) streamWriter, err := file.NewStreamWriter(sheet) if err != nil { return "", 0, err } headerValues := make([]interface{}, 0, len(headers)) for _, item := range headers { headerValues = append(headerValues, item) } if err := streamWriter.SetRow("A1", headerValues); err != nil { return "", 0, err } for i, row := range rows { axis, err := excelize.CoordinatesToCellName(1, i+2) if err != nil { return "", 0, err } values := make([]interface{}, 0, len(row)) for _, item := range row { values = append(values, item) } if err := streamWriter.SetRow(axis, values); err != nil { return "", 0, err } } if err := streamWriter.Flush(); err != nil { return "", 0, err } path := filepath.Join(os.TempDir(), fmt.Sprintf("export-%d.xlsx", time.Now().UnixNano())) if err := file.SaveAs(path); err != nil { return "", 0, err } if err := file.Close(); err != nil { return "", 0, err } stat, err := os.Stat(path) if err != nil { return "", 0, err } return path, stat.Size(), nil } func uploadExportFile(ctx context.Context, storageSvc *storage.Service, localPath, fileName string) (string, error) { if storageSvc == nil || storageSvc.Provider() == nil { return "", errors.New("对象存储服务未配置") } key, err := storageSvc.GenerateFileKey("export", fileName) if err != nil { return "", err } file, err := os.Open(localPath) if err != nil { return "", err } defer file.Close() contentType := "application/octet-stream" if filepath.Ext(fileName) == ".csv" { contentType = "text/csv" } if filepath.Ext(fileName) == ".xlsx" { contentType = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" } if err := storageSvc.Provider().Upload(ctx, key, file, contentType); err != nil { return "", err } return key, nil } func parseExportPayload[T any](payload []byte, target *T) error { return sonic.Unmarshal(payload, target) }