统一导出任务

internal/task/export_shard.go

@@ -0,0 +1,183 @@
+package task
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/hibiken/asynq"
+	"github.com/redis/go-redis/v9"
+	"go.uber.org/zap"
+
+	"github.com/break/junhong_cmp_fiber/internal/exporter"
+	"github.com/break/junhong_cmp_fiber/internal/store/postgres"
+	"github.com/break/junhong_cmp_fiber/pkg/constants"
+	"github.com/break/junhong_cmp_fiber/pkg/storage"
+)
+
+// ExportShardHandler 导出 shard 处理器。
+type ExportShardHandler struct {
+	redis         *redis.Client
+	taskStore     *postgres.ExportTaskStore
+	shardStore    *postgres.ExportShardTaskStore
+	asynqClient   *asynq.Client
+	storageSvc    *storage.Service
+	sceneRegistry *exporter.Registry
+	logger        *zap.Logger
+}
+
+// NewExportShardHandler 创建导出 shard 处理器。
+func NewExportShardHandler(
+	redis *redis.Client,
+	taskStore *postgres.ExportTaskStore,
+	shardStore *postgres.ExportShardTaskStore,
+	asynqClient *asynq.Client,
+	storageSvc *storage.Service,
+	sceneRegistry *exporter.Registry,
+	logger *zap.Logger,
+) *ExportShardHandler {
+	return &ExportShardHandler{
+		redis:         redis,
+		taskStore:     taskStore,
+		shardStore:    shardStore,
+		asynqClient:   asynqClient,
+		storageSvc:    storageSvc,
+		sceneRegistry: sceneRegistry,
+		logger:        logger,
+	}
+}
+
+// HandleExportShard 处理导出 shard 任务。
+func (h *ExportShardHandler) HandleExportShard(ctx context.Context, task *asynq.Task) error {
+	var payload ExportShardPayload
+	if err := parseExportPayload(task.Payload(), &payload); err != nil {
+		h.logger.Error("解析导出 shard 任务载荷失败", zap.Error(err))
+		return asynq.SkipRetry
+	}
+
+	lockKey := constants.RedisExportShardLockKey(payload.TaskID, payload.ShardID)
+	locked, err := tryAcquireLock(ctx, h.redis, lockKey, exportShardLockTTL)
+	if err != nil {
+		h.logger.Error("获取导出 shard 锁失败", zap.Uint("task_id", payload.TaskID), zap.Uint("shard_id", payload.ShardID), zap.Error(err))
+		return err
+	}
+	if !locked {
+		return nil
+	}
+	defer releaseLock(ctx, h.redis, lockKey)
+
+	exportTask, err := h.taskStore.GetByIDForWorker(ctx, payload.TaskID)
+	if err != nil {
+		h.logger.Error("查询导出任务失败", zap.Uint("task_id", payload.TaskID), zap.Error(err))
+		return asynq.SkipRetry
+	}
+
+	shardTask, err := h.shardStore.GetByIDForWorker(ctx, payload.ShardID)
+	if err != nil {
+		h.logger.Error("查询导出分片失败", zap.Uint("task_id", payload.TaskID), zap.Uint("shard_id", payload.ShardID), zap.Error(err))
+		return asynq.SkipRetry
+	}
+
+	if shardTask.Status == constants.ExportShardStatusSuccess || shardTask.Status == constants.ExportShardStatusFailed || shardTask.Status == constants.ExportShardStatusCancelled {
+		return nil
+	}
+
+	updater := dispatchUpdater(exportTask)
+	if exportTask.CancelRequested {
+		_, _ = h.shardStore.MarkCancelled(ctx, shardTask.ID, updater, "任务已取消")
+		_ = h.enqueueFinalize(ctx, exportTask.ID)
+		return nil
+	}
+
+	if shardTask.Status == constants.ExportShardStatusPending {
+		started, err := h.shardStore.TryMarkProcessingIfPending(ctx, shardTask.ID, updater)
+		if err != nil {
+			return err
+		}
+		if !started {
+			latest, err := h.shardStore.GetByIDForWorker(ctx, shardTask.ID)
+			if err != nil {
+				return err
+			}
+			if latest.Status != constants.ExportShardStatusProcessing {
+				return nil
+			}
+			shardTask = latest
+		}
+	}
+
+	strategy, ok := h.sceneRegistry.Get(exportTask.Scene)
+	if !ok {
+		_ = h.markShardFinalFailed(ctx, exportTask.ID, shardTask.ID, updater, "导出场景未注册")
+		return asynq.SkipRetry
+	}
+
+	rows, err := strategy.QueryRows(ctx, exportTask, shardTask.CursorStart, shardTask.CursorEnd)
+	if err != nil {
+		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "查询分片数据失败", err)
+	}
+
+	localPath, fileSize, err := writeExportFile(exportTask.Format, strategy.Headers(), rows)
+	if err != nil {
+		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "生成分片文件失败", err)
+	}
+	defer func() { _ = os.Remove(localPath) }()
+
+	fileName := fmt.Sprintf("%s-shard-%04d.%s", exportTask.TaskNo, shardTask.ShardNo, exportTask.Format)
+	fileKey, err := uploadExportFile(ctx, h.storageSvc, localPath, fileName)
+	if err != nil {
+		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "上传分片文件失败", err)
+	}
+
+	updated, err := h.shardStore.MarkSuccess(ctx, shardTask.ID, updater, len(rows), fileKey, fileSize)
+	if err != nil {
+		return err
+	}
+	if updated {
+		if err := h.taskStore.ApplyShardSuccess(ctx, exportTask.ID, updater, len(rows)); err != nil {
+			return err
+		}
+	}
+
+	if err := h.enqueueFinalize(ctx, exportTask.ID); err != nil {
+		h.logger.Warn("触发 finalize 入队失败", zap.Uint("task_id", exportTask.ID), zap.Error(err))
+	}
+
+	return nil
+}
+
+func (h *ExportShardHandler) handleShardError(ctx context.Context, taskID, shardID uint, updater uint, prefix string, err error) error {
+	if isFinalRetry(ctx) {
+		_ = h.markShardFinalFailed(ctx, taskID, shardID, updater, prefix+": "+err.Error())
+		return nil
+	}
+	return err
+}
+
+func (h *ExportShardHandler) markShardFinalFailed(ctx context.Context, taskID, shardID uint, updater uint, message string) error {
+	updated, err := h.shardStore.MarkFailed(ctx, shardID, updater, message)
+	if err != nil {
+		return err
+	}
+	if updated {
+		if err := h.taskStore.ApplyShardFailed(ctx, taskID, updater, message); err != nil {
+			return err
+		}
+	}
+	_ = h.enqueueFinalize(ctx, taskID)
+	return nil
+}
+
+func (h *ExportShardHandler) enqueueFinalize(ctx context.Context, taskID uint) error {
+	return enqueueTask(
+		ctx,
+		h.asynqClient,
+		constants.TaskTypeExportFinalize,
+		ExportFinalizePayload{TaskID: taskID},
+		asynq.Queue(constants.QueueDefault),
+		asynq.MaxRetry(constants.ExportFinalizeRetryMax),
+		asynq.Timeout(constants.ExportFinalizeTaskTimeout),
+		asynq.ProcessIn(2*time.Second),
+	)
+}