FranklinWH-Toolbox/service/task/queue/dataDownloadQueue/dataDownloadQueue.go

package dataDownloadQueue

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/jobs"
	minioRepo "fwh-toolbox-server/repositories/minio"
	mongoRepo "fwh-toolbox-server/repositories/mongo"
	"fwh-toolbox-server/services"
	task_core "fwh-toolbox-server/task/core"

	"github.com/hibiken/asynq"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.uber.org/zap"
)

const (
	TaskTypePlan    = "data-download:plan"
	TaskTypeItem    = "data-download:item"
	TaskTypePackage = "data-download:package"
	QueuePlan       = "data-download:plan"
	QueueItem       = "data-download:item"
)

var client *asynq.Client

func InitClient(redisOpt asynq.RedisClientOpt) {
	if client == nil {
		client = asynq.NewClient(redisOpt)
	}
}

type planPayload struct {
	TaskID string `json:"taskId"`
}
type itemPayload struct {
	ItemID         string `json:"itemId"`
	AwaitRequestID string `json:"awaitRequestId,omitempty"`
	DeferredCount  int    `json:"deferredCount,omitempty"`
}
type packagePayload struct {
	JobID string `json:"jobId"`
}

func Init(tc *task_core.TaskCore) {
	mux := asynq.NewServeMux()
	mux.HandleFunc(TaskTypePlan, handlePlan)
	mux.HandleFunc(TaskTypeItem, handleItem)
	mux.HandleFunc(TaskTypePackage, handlePackage)
	client = tc.CreateTask(task_core.TaskConfig{
		Name: "DataDownload",
		SvrCfg: asynq.Config{Concurrency: 8, Queues: map[string]int{QueueItem: 6, QueuePlan: 2}, ErrorHandler: asynq.ErrorHandlerFunc(func(_ context.Context, task *asynq.Task, err error) {
			zap.L().Error("data-download task failed", zap.String("type", task.Type()), zap.String("err", err.Error()))
		})},
		Mux: mux,
	})
}

func EnqueuePackage(jobID primitive.ObjectID) error {
	payload, _ := json.Marshal(packagePayload{JobID: jobID.Hex()})
	_, err := client.Enqueue(asynq.NewTask(TaskTypePackage, payload, asynq.Queue(QueueItem), asynq.MaxRetry(1), asynq.Timeout(2*time.Hour), asynq.TaskID("dd-package-"+jobID.Hex()), asynq.Retention(30*24*time.Hour)))
	return err
}

func handlePackage(ctx context.Context, task *asynq.Task) error {
	var payload packagePayload
	if err := json.Unmarshal(task.Payload(), &payload); err != nil {
		return err
	}
	id, err := primitive.ObjectIDFromHex(payload.JobID)
	if err != nil {
		return err
	}
	job, err := mongoRepo.GetDataDownloadPackageJob(ctx, id)
	if err != nil {
		return err
	}
	if job.Status == domain.DataDownloadCancelled || job.Status == domain.DataDownloadCancelling {
		if job.Status == domain.DataDownloadCancelling {
			_ = mongoRepo.UpdateDataDownloadPackageJob(ctx, id, bson.M{"status": domain.DataDownloadCancelled})
		}
		return nil
	}
	now := time.Now().UTC()
	_ = mongoRepo.UpdateDataDownloadPackageJob(ctx, id, bson.M{"status": domain.DataDownloadRunning, "startedAt": now})
	zap.L().Info("[DataDownload] 打包开始",
		zap.String("jobId", id.Hex()),
		zap.Int("taskCount", len(job.TaskIDs)),
		zap.String("filename", job.Filename),
	)
	data, count, err := services.BuildDataDownloadPackage(ctx, job, func(done, total int64) error {
		current, getErr := mongoRepo.GetDataDownloadPackageJob(ctx, id)
		if getErr != nil {
			return getErr
		}
		if current.Status == domain.DataDownloadCancelling || current.Status == domain.DataDownloadCancelled {
			return context.Canceled
		}
		return mongoRepo.UpdateDataDownloadPackageJob(ctx, id, bson.M{"doneCount": done, "totalCount": total})
	})
	if err != nil {
		status, reason := domain.DataDownloadFailed, err.Error()
		if err == context.Canceled {
			status, reason = domain.DataDownloadCancelled, ""
		}
		zap.L().Error("[DataDownload] 打包失败", zap.String("jobId", id.Hex()), zap.String("status", string(status)), zap.String("err", err.Error()))
		_ = mongoRepo.UpdateDataDownloadPackageJob(ctx, id, bson.M{"status": status, "failureReason": reason})
		return nil
	}
	objectKey, err := minioRepo.PutPackage(ctx, job.Filename, data)
	if err != nil {
		zap.L().Error("[DataDownload] 打包上传 MinIO 失败", zap.String("jobId", id.Hex()), zap.String("err", err.Error()))
		_ = mongoRepo.UpdateDataDownloadPackageJob(ctx, id, bson.M{"status": domain.DataDownloadFailed, "failureReason": err.Error()})
		return nil
	}
	completed := time.Now().UTC()
	expires := completed.Add(7 * 24 * time.Hour)
	purge := expires.Add(30 * 24 * time.Hour)
	zap.L().Info("[DataDownload] 打包完成",
		zap.String("jobId", id.Hex()),
		zap.Int64("fileCount", count),
		zap.Int("bytes", len(data)),
		zap.String("objectKey", objectKey),
	)
	return mongoRepo.UpdateDataDownloadPackageJob(ctx, id, bson.M{"status": domain.DataDownloadSucceeded, "objectKey": objectKey, "size": len(data), "doneCount": count, "totalCount": count, "completedAt": completed, "expiresAt": expires, "purgeAt": purge, "failureReason": ""})
}

func EnqueuePlan(taskID primitive.ObjectID) error {
	if client == nil {
		return fmt.Errorf("data-download queue is not initialized")
	}
	payload, _ := json.Marshal(planPayload{TaskID: taskID.Hex()})
	_, err := client.Enqueue(asynq.NewTask(TaskTypePlan, payload, asynq.Queue(QueuePlan), asynq.MaxRetry(3), asynq.Timeout(15*time.Minute), asynq.TaskID("dd-plan-"+taskID.Hex()), asynq.Retention(30*24*time.Hour)))
	if err != nil {
		zap.L().Error("[DataDownload] plan 入队失败", zap.String("taskId", taskID.Hex()), zap.String("err", err.Error()))
	} else {
		zap.L().Info("[DataDownload] plan 已入队", zap.String("taskId", taskID.Hex()))
	}
	return err
}

func enqueueItem(itemID primitive.ObjectID) error {
	payload, _ := json.Marshal(itemPayload{ItemID: itemID.Hex()})
	_, err := client.Enqueue(asynq.NewTask(TaskTypeItem, payload, asynq.Queue(QueueItem), asynq.MaxRetry(3), asynq.Timeout(40*time.Minute), asynq.TaskID("dd-item-"+itemID.Hex()), asynq.Retention(30*24*time.Hour)))
	if errors.Is(err, asynq.ErrTaskIDConflict) {
		return nil
	}
	return err
}

func enqueueDeferredItem(itemID primitive.ObjectID, requestID string, deferredCount int, after time.Duration) error {
	payload, _ := json.Marshal(itemPayload{ItemID: itemID.Hex(), AwaitRequestID: requestID, DeferredCount: deferredCount})
	_, err := client.Enqueue(asynq.NewTask(TaskTypeItem, payload,
		asynq.Queue(QueueItem),
		asynq.MaxRetry(3),
		asynq.Timeout(40*time.Minute),
		asynq.ProcessIn(after),
		asynq.TaskID(fmt.Sprintf("dd-item-wait-%s-%d", itemID.Hex(), time.Now().UnixNano())),
		asynq.Retention(30*24*time.Hour)))
	return err
}

func handlePlan(ctx context.Context, job *asynq.Task) error {
	var payload planPayload
	if err := json.Unmarshal(job.Payload(), &payload); err != nil {
		return err
	}
	id, err := primitive.ObjectIDFromHex(payload.TaskID)
	if err != nil {
		return err
	}
	task, err := mongoRepo.GetDataDownloadTask(ctx, id)
	if err != nil {
		return err
	}
	zap.L().Info("[DataDownload] plan 开始",
		zap.String("taskId", id.Hex()),
		zap.Int("snCount", len(task.SNs)),
		zap.String("status", string(task.Status)),
	)
	if task.Status == domain.DataDownloadCancelled || task.Status == domain.DataDownloadCancelling {
		if task.Status == domain.DataDownloadCancelling {
			_ = mongoRepo.UpdateDataDownloadTaskStatus(ctx, id, domain.DataDownloadCancelled)
		}
		zap.L().Info("[DataDownload] plan 因任务取消而终止", zap.String("taskId", id.Hex()))
		return nil
	}
	if err := mongoRepo.UpdateDataDownloadTaskStatus(ctx, id, domain.DataDownloadRunning); err != nil {
		return err
	}
	totalItems := 0
	// Plan and enqueue one site at a time so an unbounded SN list never expands
	// into one large in-memory slice or MongoDB document.
	for _, sn := range task.SNs {
		latest, getErr := mongoRepo.GetDataDownloadTask(ctx, id)
		if getErr != nil {
			return getErr
		}
		if latest.Status == domain.DataDownloadCancelling || latest.Status == domain.DataDownloadCancelled {
			_ = mongoRepo.UpdateDataDownloadTaskStatus(ctx, id, domain.DataDownloadCancelled)
			break
		}
		siteTask := *task
		siteTask.SNs = []string{sn}
		items, planErr := services.PlanDataDownloadTask(ctx, &siteTask)
		if planErr != nil {
			zap.L().Error("[DataDownload] plan 站点规划失败",
				zap.String("taskId", id.Hex()),
				zap.String("sn", sn),
				zap.String("err", planErr.Error()),
			)
			_ = mongoRepo.UpdateDataDownloadTaskStatus(ctx, id, domain.DataDownloadError)
			return planErr
		}
		if err := mongoRepo.InsertDataDownloadItems(ctx, items); err != nil {
			return err
		}
		totalItems += len(items)
		zap.L().Debug("[DataDownload] plan 站点拆分完成",
			zap.String("taskId", id.Hex()),
			zap.String("sn", sn),
			zap.Int("itemCount", len(items)),
		)
		for i := range items {
			if items[i].Status == domain.DataDownloadFailed || items[i].Status == domain.DataDownloadSkipped {
				continue
			}
			if err := enqueueItem(items[i].ID); err != nil {
				zap.L().Error("[DataDownload] item 入队失败",
					zap.String("taskId", id.Hex()),
					zap.String("itemId", items[i].ID.Hex()),
					zap.String("kind", string(items[i].Kind)),
					zap.String("err", err.Error()),
				)
				_ = mongoRepo.UpdateDataDownloadItem(ctx, items[i].ID, domain.DataDownloadFailed, "enqueue", err.Error())
			}
		}
		_ = mongoRepo.RecountDataDownloadTask(ctx, id)
	}
	if totalItems == 0 {
		latest, getErr := mongoRepo.GetDataDownloadTask(ctx, id)
		if getErr == nil && latest.Status == domain.DataDownloadCancelled {
			return nil
		}
		zap.L().Info("[DataDownload] plan 完成，无有效 item，任务直接置为成功", zap.String("taskId", id.Hex()))
		_ = mongoRepo.UpdateDataDownloadTaskStatus(ctx, id, domain.DataDownloadSucceeded)
		return nil
	}
	zap.L().Info("[DataDownload] plan 完成，item 已全部入队",
		zap.String("taskId", id.Hex()),
		zap.Int("totalItems", totalItems),
	)
	return aggregate(ctx, id)
}

func EnqueueItem(itemID primitive.ObjectID) error {
	payload, _ := json.Marshal(itemPayload{ItemID: itemID.Hex()})
	_, err := client.Enqueue(asynq.NewTask(TaskTypeItem, payload, asynq.Queue(QueueItem), asynq.MaxRetry(3), asynq.Timeout(40*time.Minute), asynq.TaskID(fmt.Sprintf("dd-item-%s-%d", itemID.Hex(), time.Now().UnixNano())), asynq.Retention(30*24*time.Hour)))
	return err
}

func handleItem(ctx context.Context, job *asynq.Task) error {
	var payload itemPayload
	if err := json.Unmarshal(job.Payload(), &payload); err != nil {
		return err
	}
	id, err := primitive.ObjectIDFromHex(payload.ItemID)
	if err != nil {
		return err
	}
	item, err := mongoRepo.GetDataDownloadItem(ctx, id)
	if err != nil {
		return err
	}
	task, err := mongoRepo.GetDataDownloadTask(ctx, item.TaskID)
	if err != nil {
		return err
	}
	retried, _ := asynq.GetRetryCount(ctx)
	zap.L().Debug("[DataDownload] item 开始处理",
		zap.String("itemId", id.Hex()),
		zap.String("taskId", item.TaskID.Hex()),
		zap.String("kind", string(item.Kind)),
		zap.String("sn", item.SiteSN),
		zap.String("subtype", item.Subtype),
		zap.String("date", item.Date),
		zap.Int("attempt", retried),
	)
	if task.Status == domain.DataDownloadCancelled || task.Status == domain.DataDownloadCancelling {
		zap.L().Info("[DataDownload] item 因任务取消而跳过", zap.String("itemId", id.Hex()), zap.String("taskId", item.TaskID.Hex()))
		_ = services.ReleaseFaultCallOwnership(context.Background(), item)
		_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadCancelled, "cancelled", "")
		return aggregate(ctx, item.TaskID)
	}
	if item.Kind == "site_validation" {
		_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadRunning, "validation", "")
		retryTask := *task
		retryTask.SNs = []string{item.SiteSN}
		planned, planErr := services.PlanDataDownloadTask(ctx, &retryTask)
		if planErr != nil || (len(planned) == 1 && planned[0].Kind == "site_validation") {
			reason := "site does not exist"
			if planErr != nil {
				reason = planErr.Error()
			} else if len(planned) == 1 {
				reason = planned[0].FailureReason
			}
			_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadFailed, "validation", reason)
			return aggregate(ctx, item.TaskID)
		}
		if err := mongoRepo.InsertDataDownloadItems(ctx, planned); err != nil {
			_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadFailed, "validation", err.Error())
			return aggregate(ctx, item.TaskID)
		}
		_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadItemSucceeded, "validation", "")
		for i := range planned {
			if err := enqueueItem(planned[i].ID); err != nil {
				_ = mongoRepo.UpdateDataDownloadItem(ctx, planned[i].ID, domain.DataDownloadFailed, "enqueue", err.Error())
			}
		}
		return aggregate(ctx, item.TaskID)
	}
	_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadRunning, "fetch", "")
	err = services.ExecuteStructuredDataDownloadItem(ctx, item, payload.AwaitRequestID)
	if err != nil {
		var faultCallDeferred *services.DataDownloadFaultCallDeferred
		if errors.As(err, &faultCallDeferred) {
			_ = mongoRepo.DeferDataDownloadItem(ctx, id)
			if enqueueErr := enqueueDeferredItem(id, payload.AwaitRequestID, payload.DeferredCount+1, faultCallDeferred.After); enqueueErr != nil {
				_ = services.ReleaseFaultCallOwnership(context.Background(), item)
				_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadFailed, "enqueue", enqueueErr.Error())
			}
			return aggregate(ctx, item.TaskID)
		}
		var snDeferred *services.DataDownloadSNDeferred
		if errors.As(err, &snDeferred) {
			delay := snDeferred.After
			zap.L().Debug("[DataDownload] EMS 召唤名额被占用，延迟处理",
				zap.String("itemId", id.Hex()), zap.String("sn", item.SiteSN), zap.Duration("delay", delay))
			_ = mongoRepo.DeferDataDownloadItem(ctx, id)
			if enqueueErr := enqueueDeferredItem(id, payload.AwaitRequestID, payload.DeferredCount+1, delay); enqueueErr != nil {
				_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadFailed, "enqueue", enqueueErr.Error())
			}
			return aggregate(ctx, item.TaskID)
		}
		var deferred *services.DataDownloadRequestDeferred
		if errors.As(err, &deferred) {
			delay := deferred.After
			switch payload.DeferredCount {
			case 1:
				delay = 30 * time.Second
			case 2:
				delay = time.Minute
			default:
				if payload.DeferredCount > 2 {
					delay = time.Minute
				}
			}
			zap.L().Debug("[DataDownload] item 等待请求租约，延迟重入队",
				zap.String("itemId", id.Hex()),
				zap.String("requestId", deferred.RequestID),
				zap.Int("deferredCount", payload.DeferredCount+1),
				zap.Duration("delay", delay),
			)
			_ = mongoRepo.DeferDataDownloadItem(ctx, id)
			if enqueueErr := enqueueDeferredItem(id, deferred.RequestID, payload.DeferredCount+1, delay); enqueueErr != nil {
				zap.L().Error("[DataDownload] item 延迟重入队失败", zap.String("itemId", id.Hex()), zap.String("err", enqueueErr.Error()))
				_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadFailed, "enqueue", enqueueErr.Error())
				return aggregate(ctx, item.TaskID)
			}
			return aggregate(ctx, item.TaskID)
		}
		maxRetry, _ := asynq.GetMaxRetry(ctx)
		if shouldRetryItem(err, retried, maxRetry) {
			zap.L().Warn("[DataDownload] item 执行失败，将重试",
				zap.String("itemId", id.Hex()),
				zap.String("kind", string(item.Kind)),
				zap.Int("attempt", retried),
				zap.Int("maxRetry", maxRetry),
				zap.String("err", err.Error()),
			)
			return err
		}
		var partial *services.DataDownloadPartialFailure
		if errors.As(err, &partial) && partial.Succeeded > 0 {
			zap.L().Warn("[DataDownload] item 部分成功（重试耗尽）",
				zap.String("itemId", id.Hex()),
				zap.String("kind", string(item.Kind)),
				zap.Int("succeeded", partial.Succeeded),
				zap.String("err", err.Error()),
			)
			_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadPartial, "fetch", err.Error())
		} else {
			zap.L().Error("[DataDownload] item 最终失败（重试耗尽）",
				zap.String("itemId", id.Hex()),
				zap.String("taskId", item.TaskID.Hex()),
				zap.String("kind", string(item.Kind)),
				zap.String("sn", item.SiteSN),
				zap.String("err", err.Error()),
			)
			_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadFailed, "fetch", err.Error())
		}
		_ = services.ReleaseFaultCallOwnership(context.Background(), item)
	} else {
		latest, getErr := mongoRepo.GetDataDownloadTask(ctx, item.TaskID)
		if getErr == nil && (latest.Status == domain.DataDownloadCancelling || latest.Status == domain.DataDownloadCancelled) {
			_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadCancelled, "cancelled", "")
		} else {
			zap.L().Info("[DataDownload] item 执行成功",
				zap.String("itemId", id.Hex()),
				zap.String("taskId", item.TaskID.Hex()),
				zap.String("kind", string(item.Kind)),
				zap.String("sn", item.SiteSN),
			)
			_ = mongoRepo.UpdateDataDownloadItem(ctx, id, domain.DataDownloadItemSucceeded, "complete", "")
		}
	}
	return aggregate(ctx, item.TaskID)
}

func shouldRetryItem(err error, retried, maxRetry int) bool {
	return retried < maxRetry && !errors.Is(err, jobs.ErrEmsLogCallTimeout) && !errors.Is(err, jobs.ErrEmsLogCallSlotLost) && !errors.Is(err, services.ErrFaultRecordCallTimeout) && !errors.Is(err, services.ErrFaultRecordCallFailed)
}

func aggregate(ctx context.Context, taskID primitive.ObjectID) error {
	task, taskErr := mongoRepo.GetDataDownloadTask(ctx, taskID)
	if taskErr != nil {
		return taskErr
	}
	items, err := mongoRepo.ListDataDownloadItems(ctx, taskID)
	if err != nil {
		return err
	}
	var succeeded, partial, failed, cancelled int
	for _, item := range items {
		switch item.Status {
		case domain.DataDownloadItemSucceeded, domain.DataDownloadSkipped:
			succeeded++
		case domain.DataDownloadPartial:
			partial++
		case domain.DataDownloadFailed:
			failed++
		case domain.DataDownloadCancelled:
			cancelled++
		default:
			_ = mongoRepo.RecountDataDownloadTask(ctx, taskID)
			return nil
		}
	}
	status := domain.DataDownloadSucceeded
	if task.Status == domain.DataDownloadCancelling || task.Status == domain.DataDownloadCancelled {
		status = domain.DataDownloadCancelled
	} else if cancelled == len(items) {
		status = domain.DataDownloadCancelled
	} else if failed == len(items) {
		status = domain.DataDownloadFailed
	} else if partial > 0 || failed > 0 || cancelled > 0 {
		status = domain.DataDownloadPartial
	}
	if err := mongoRepo.RecountDataDownloadTask(ctx, taskID); err != nil {
		return err
	}
	zap.L().Info("[DataDownload] 任务完成",
		zap.String("taskId", taskID.Hex()),
		zap.String("status", string(status)),
		zap.Int("total", len(items)),
		zap.Int("succeeded", succeeded),
		zap.Int("partial", partial),
		zap.Int("failed", failed),
		zap.Int("cancelled", cancelled),
	)
	return mongoRepo.UpdateDataDownloadTaskStatus(ctx, taskID, status)
}