junhong_cmp_fiber/internal/task/export_dispatch.go

package task

import (
	"context"
	"fmt"
	"time"

	"github.com/hibiken/asynq"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
	"gorm.io/gorm"

	"github.com/break/junhong_cmp_fiber/internal/exporter"
	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/internal/store/postgres"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// ExportDispatchHandler 导出 dispatch 处理器。
type ExportDispatchHandler struct {
	db            *gorm.DB
	redis         *redis.Client
	taskStore     *postgres.ExportTaskStore
	shardStore    *postgres.ExportShardTaskStore
	asynqClient   *asynq.Client
	sceneRegistry *exporter.Registry
	logger        *zap.Logger
}

// NewExportDispatchHandler 创建导出 dispatch 处理器。
func NewExportDispatchHandler(
	db *gorm.DB,
	redis *redis.Client,
	taskStore *postgres.ExportTaskStore,
	shardStore *postgres.ExportShardTaskStore,
	asynqClient *asynq.Client,
	sceneRegistry *exporter.Registry,
	logger *zap.Logger,
) *ExportDispatchHandler {
	return &ExportDispatchHandler{
		db:            db,
		redis:         redis,
		taskStore:     taskStore,
		shardStore:    shardStore,
		asynqClient:   asynqClient,
		sceneRegistry: sceneRegistry,
		logger:        logger,
	}
}

// HandleExportDispatch 处理导出 dispatch 任务。
func (h *ExportDispatchHandler) HandleExportDispatch(ctx context.Context, task *asynq.Task) error {
	var payload ExportDispatchPayload
	if err := parseExportPayload(task.Payload(), &payload); err != nil {
		h.logger.Error("解析导出 dispatch 任务载荷失败", zap.Error(err))
		return asynq.SkipRetry
	}

	lockKey := constants.RedisExportDispatchLockKey(payload.TaskID)
	locked, err := tryAcquireLock(ctx, h.redis, lockKey, exportLockTTL)
	if err != nil {
		h.logger.Error("获取导出 dispatch 锁失败", zap.Uint("task_id", payload.TaskID), zap.Error(err))
		return err
	}
	if !locked {
		h.logger.Info("导出 dispatch 已在执行，跳过", zap.Uint("task_id", payload.TaskID))
		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
	}

	if exportTask.Status == constants.ExportTaskStatusCompleted || exportTask.Status == constants.ExportTaskStatusFailed || exportTask.Status == constants.ExportTaskStatusCancelled {
		return nil
	}

	updater := dispatchUpdater(exportTask)

	if exportTask.CancelRequested {
		if err := h.shardStore.CancelByTaskID(ctx, exportTask.ID, updater, "任务已取消"); err != nil {
			return err
		}
		if err := h.taskStore.MarkCancelledByWorker(ctx, exportTask.ID, updater, "任务已取消"); err != nil {
			return err
		}
		return nil
	}

	started, err := h.taskStore.TryMarkProcessingIfPending(ctx, exportTask.ID, updater)
	if err != nil {
		h.logger.Error("更新导出任务状态为处理中失败", zap.Uint("task_id", exportTask.ID), zap.Error(err))
		return err
	}

	if !started {
		if exportTask.Status != constants.ExportTaskStatusProcessing {
			return nil
		}
		existingShards, err := h.shardStore.ListByTaskID(ctx, exportTask.ID)
		if err != nil {
			return err
		}
		if len(existingShards) == 0 {
			return nil
		}
		if err := h.enqueueShardTasks(ctx, exportTask, existingShards); err != nil {
			return err
		}
		return h.enqueueFinalizeTask(ctx, exportTask.ID)
	}

	strategy, ok := h.sceneRegistry.Get(exportTask.Scene)
	if !ok {
		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "导出场景未注册")
		return asynq.SkipRetry
	}

	params := exporter.ParseExportParams(exportTask)
	if err := exporter.ValidateTaskTimeFilters(exportTask.Scene, params.Filters); err != nil {
		// 冻结值非法（含变更前遗留任务的旧格式值）必须落失败，不得忽略条件后放行全量数据。
		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, constants.ExportTaskInvalidTimeFilterMessage)
		h.logger.Error("导出任务冻结的时间边界非法", zap.Uint("task_id", exportTask.ID), zap.Error(err))
		return asynq.SkipRetry
	}
	headers, err := strategy.Headers(ctx, params)
	if err != nil {
		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "解析导出表头失败")
		h.logger.Error("解析导出表头失败", zap.Uint("task_id", exportTask.ID), zap.Error(err))
		return err
	}

	shards, totalRows, err := h.buildShards(ctx, strategy, params, exportTask.CancelRequested, exportTask.ID, updater)
	if err != nil {
		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "构建导出分片失败")
		h.logger.Error("构建导出分片失败", zap.Uint("task_id", exportTask.ID), zap.Error(err))
		return err
	}

	if err := h.persistShardsAndPlan(ctx, exportTask.ID, updater, shards, totalRows, headers); err != nil {
		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "保存导出分片失败")
		return err
	}

	if err := h.enqueueShardTasks(ctx, exportTask, shards); err != nil {
		return err
	}

	if err := h.enqueueFinalizeTask(ctx, exportTask.ID); err != nil {
		return err
	}

	h.logger.Info("导出 dispatch 完成",
		zap.Uint("task_id", exportTask.ID),
		zap.Int("total_rows", totalRows),
		zap.Int("total_shards", len(shards)),
	)

	return nil
}

func (h *ExportDispatchHandler) buildShards(ctx context.Context, source exporter.DataSource, params exporter.ExportParams, cancelRequested bool, taskID uint, updater uint) ([]*model.ExportShardTask, int, error) {
	shards := make([]*model.ExportShardTask, 0)
	if cancelRequested {
		return shards, 0, nil
	}

	totalRows, err := source.Count(ctx, params)
	if err != nil {
		return nil, 0, err
	}
	if totalRows <= 0 {
		return shards, 0, nil
	}

	shardNo := 1
	for offset := 0; offset < totalRows; offset += constants.ExportDefaultShardSize {
		limit := constants.ExportDefaultShardSize
		if remaining := totalRows - offset; remaining < limit {
			limit = remaining
		}

		shard := &model.ExportShardTask{
			TaskID:      taskID,
			ShardNo:     shardNo,
			Status:      constants.ExportShardStatusPending,
			ShardOffset: offset,
			ShardLimit:  limit,
			RowCount:    limit,
		}
		shard.Creator = updater
		shard.Updater = updater
		shards = append(shards, shard)

		shardNo++
	}

	return shards, totalRows, nil
}

func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID uint, updater uint, shards []*model.ExportShardTask, totalRows int, headers []string) error {
	return h.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		txShardStore := postgres.NewExportShardTaskStore(tx, h.redis)
		txTaskStore := postgres.NewExportTaskStore(tx, h.redis)

		if err := txShardStore.CreateBatch(ctx, shards); err != nil {
			return err
		}
		if err := txTaskStore.UpdateResolvedHeaders(ctx, taskID, headers, updater); err != nil {
			return err
		}
		if err := txTaskStore.UpdateDispatchPlan(ctx, taskID, totalRows, len(shards), updater); err != nil {
			return err
		}
		return nil
	})
}

func (h *ExportDispatchHandler) enqueueShardTasks(ctx context.Context, task *model.ExportTask, shards []*model.ExportShardTask) error {
	for _, shard := range shards {
		payload := ExportShardPayload{TaskID: task.ID, ShardID: shard.ID}
		err := enqueueTask(
			ctx,
			h.asynqClient,
			constants.TaskTypeExportShard,
			payload,
			asynq.MaxRetry(constants.ExportShardRetryMax),
			asynq.Timeout(constants.ExportShardTaskTimeout),
			asynq.Queue(constants.QueueForTaskType(constants.TaskTypeExportShard)),
		)
		if err != nil {
			h.logger.Error("导出分片任务入队失败",
				zap.Uint("task_id", task.ID),
				zap.Uint("shard_id", shard.ID),
				zap.Error(err),
			)
			return fmt.Errorf("enqueue shard task failed: %w", err)
		}
	}
	return nil
}

func (h *ExportDispatchHandler) enqueueFinalizeTask(ctx context.Context, taskID uint) error {
	return enqueueTask(
		ctx,
		h.asynqClient,
		constants.TaskTypeExportFinalize,
		ExportFinalizePayload{TaskID: taskID},
		asynq.MaxRetry(constants.ExportFinalizeRetryMax),
		asynq.Timeout(constants.ExportFinalizeTaskTimeout),
		asynq.ProcessIn(5*time.Second),
		asynq.Queue(constants.QueueForTaskType(constants.TaskTypeExportFinalize)),
	)
}

func dispatchUpdater(task *model.ExportTask) uint {
	if task.CreatorUserID > 0 {
		return task.CreatorUserID
	}
	if task.Creator > 0 {
		return task.Creator
	}
	return 0
}