junhong_cmp_fiber/internal/task/export_shard.go

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),
	)
}