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