FranklinWH-Toolbox/service/services/dataDownloadFileRequestService.go

package services

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"fwh-toolbox-server/cache"
	"fwh-toolbox-server/domain"
	mongoRepo "fwh-toolbox-server/repositories/mongo"

	"github.com/redis/go-redis/v9"
	"go.mongodb.org/mongo-driver/mongo"
	"go.uber.org/zap"
)

const (
	dataDownloadRequestLeaseTTL = 2 * time.Minute
	dataDownloadCompletedTTL    = time.Hour
)

// DataDownloadRequestDeferred tells the queue to retry an item later without
// occupying a worker while another request for the same document is running.
type DataDownloadRequestDeferred struct {
	RequestID string
	After     time.Duration
}

func (e *DataDownloadRequestDeferred) Error() string {
	return fmt.Sprintf("data download request %s is still running", e.RequestID)
}

type dataDownloadRequestLease struct {
	documentID string
	requestKey string
	requestID  string
	lock       *cache.RedisLock
	cancel     context.CancelFunc
	done       chan struct{}
	lost       chan struct{}
	stopOnce   sync.Once
}

func dataDownloadRequestKey(documentID string) string {
	return "data-download:request:" + documentID
}

func dataDownloadCompletedKey(documentID string) string {
	return "data-download:completed:" + documentID
}

func prepareDataDownloadRequest(ctx context.Context, item *domain.DataDownloadItem, awaitRequestID string) (*domain.DataDownloadDocument, *dataDownloadRequestLease, error) {
	requestKey, err := buildDataDownloadRequestKey(item)
	if err != nil {
		return nil, nil, err
	}
	documentID := stableDataDownloadDocumentID(requestKey)
	if existing, err := mongoRepo.GetDataDownloadDocument(ctx, documentID); err == nil && existing.State == domain.DataDownloadFrozen {
		return existing, nil, nil
	} else if err != nil && !errors.Is(err, mongo.ErrNoDocuments) {
		return nil, nil, err
	}

	if awaitRequestID != "" && item.Kind != domain.DataDownloadUpdateLog {
		completed, err := cache.Get(ctx, dataDownloadCompletedKey(documentID))
		if err == nil && completed == awaitRequestID {
			document, getErr := mongoRepo.GetDataDownloadDocument(ctx, documentID)
			if getErr == nil {
				return document, nil, nil
			}
			if !errors.Is(getErr, mongo.ErrNoDocuments) {
				return nil, nil, getErr
			}
		}
		if err != nil && !errors.Is(err, redis.Nil) {
			return nil, nil, err
		}
	}

	lock, ok, err := cache.TryLock(ctx, dataDownloadRequestKey(documentID), dataDownloadRequestLeaseTTL)
	if err != nil {
		return nil, nil, err
	}
	if !ok {
		requestID, getErr := cache.Get(ctx, dataDownloadRequestKey(documentID))
		if getErr != nil && !errors.Is(getErr, redis.Nil) {
			return nil, nil, getErr
		}
		if requestID == "" {
			requestID = awaitRequestID
		}
		zap.L().Debug("[DataDownload] 请求租约被占用，本 item 延迟等待",
			zap.String("itemId", item.ID.Hex()),
			zap.String("documentId", documentID),
			zap.String("kind", string(item.Kind)),
			zap.String("sn", item.SiteSN),
			zap.String("holderRequestId", requestID),
		)
		return nil, nil, &DataDownloadRequestDeferred{RequestID: requestID, After: 15 * time.Second}
	}

	// The document may have become frozen between the first lookup and lease acquisition.
	if existing, getErr := mongoRepo.GetDataDownloadDocument(ctx, documentID); getErr == nil && existing.State == domain.DataDownloadFrozen {
		lock.SafeRelease(context.Background())
		return existing, nil, nil
	} else if getErr != nil && !errors.Is(getErr, mongo.ErrNoDocuments) {
		lock.SafeRelease(context.Background())
		return nil, nil, getErr
	}

	renewCtx, cancel := context.WithCancel(context.Background())
	lease := &dataDownloadRequestLease{
		documentID: documentID,
		requestKey: requestKey,
		requestID:  lock.GetToken(),
		lock:       lock,
		cancel:     cancel,
		done:       make(chan struct{}),
		lost:       make(chan struct{}),
	}
	go lease.renew(renewCtx)
	zap.L().Debug("[DataDownload] 已获取请求租约，进入执行",
		zap.String("itemId", item.ID.Hex()),
		zap.String("documentId", documentID),
		zap.String("kind", string(item.Kind)),
		zap.String("sn", item.SiteSN),
		zap.String("subtype", item.Subtype),
		zap.String("date", item.Date),
	)
	return nil, lease, nil
}

func (l *dataDownloadRequestLease) renew(ctx context.Context) {
	defer close(l.done)
	ticker := time.NewTicker(dataDownloadRequestLeaseTTL / 3)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			renewCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			err := l.lock.Extend(renewCtx, dataDownloadRequestLeaseTTL)
			cancel()
			if err != nil {
				zap.L().Warn("[DataDownload] 请求租约续期失败，租约丢失",
					zap.String("documentId", l.documentID),
					zap.String("requestId", l.requestID),
					zap.String("err", err.Error()),
				)
				close(l.lost)
				return
			}
		}
	}
}

func (l *dataDownloadRequestLease) validate(ctx context.Context) error {
	select {
	case <-l.lost:
		return fmt.Errorf("data download request lease was lost")
	default:
	}
	ok, err := l.lock.IsOwner(ctx)
	if err != nil {
		return err
	}
	if !ok {
		return fmt.Errorf("data download request lease was lost")
	}
	return nil
}

func (l *dataDownloadRequestLease) complete(ctx context.Context) error {
	if err := l.lock.SetIfOwner(ctx, dataDownloadCompletedKey(l.documentID), l.requestID, dataDownloadCompletedTTL); err != nil {
		l.stop()
		return err
	}
	l.stop()
	return nil
}

func (l *dataDownloadRequestLease) abort() {
	l.stop()
}

func (l *dataDownloadRequestLease) stop() {
	l.stopOnce.Do(func() {
		l.cancel()
		<-l.done
		l.lock.SafeRelease(context.Background())
	})
}