huang
b7369a9c71
All checks were successful
构建并部署到测试环境（无 SSH） / build-and-deploy (push) Successful in 9m22s
256 lines
7.2 KiB
Go
256 lines
7.2 KiB
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
|
|
}
|
|
|
|
shards, totalRows, err := h.buildShards(ctx, exportTask, strategy, 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); 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.Uint("total_rows", totalRows),
|
|
zap.Int("total_shards", len(shards)),
|
|
)
|
|
|
|
return nil
|
|
}
|
|
|
|
func (h *ExportDispatchHandler) buildShards(ctx context.Context, task *model.ExportTask, strategy exporter.SceneStrategy, updater uint) ([]*model.ExportShardTask, uint, error) {
|
|
shards := make([]*model.ExportShardTask, 0)
|
|
totalRows := uint(0)
|
|
afterID := uint64(0)
|
|
shardNo := 1
|
|
|
|
for {
|
|
if task.CancelRequested {
|
|
break
|
|
}
|
|
|
|
boundary, err := strategy.NextShardBoundary(ctx, task, afterID, constants.ExportDefaultShardSize)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
if boundary == nil {
|
|
break
|
|
}
|
|
|
|
rowCount := boundary.RowCount
|
|
if rowCount < 0 {
|
|
rowCount = 0
|
|
}
|
|
|
|
totalRows += uint(rowCount)
|
|
shard := &model.ExportShardTask{
|
|
TaskID: task.ID,
|
|
ShardNo: shardNo,
|
|
Status: constants.ExportShardStatusPending,
|
|
CursorStart: boundary.StartID,
|
|
CursorEnd: boundary.EndID,
|
|
RowCount: rowCount,
|
|
}
|
|
shard.Creator = updater
|
|
shard.Updater = updater
|
|
shards = append(shards, shard)
|
|
|
|
afterID = boundary.EndID
|
|
shardNo++
|
|
}
|
|
|
|
return shards, totalRows, nil
|
|
}
|
|
|
|
func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID uint, updater uint, shards []*model.ExportShardTask, totalRows uint) 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.UpdateDispatchPlan(ctx, taskID, totalRows, uint(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
|
|
}
|