xinmi/internal/resumable/worker.go

package resumable

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"os"
	"time"

	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// Start runs one bounded ingest worker per process. Database leases support
// recovery after restarts and prevent the same upload running on two workers.
func (s *Service) Start(ctx context.Context) {
	s.once.Do(func() {
		go func() {
			timer := time.NewTicker(2 * time.Second)
			defer timer.Stop()
			lastCleanup := time.Time{}
			for {
				if ctx.Err() != nil {
					return
				}
				if time.Since(lastCleanup) > time.Minute {
					if e := s.Cleanup(ctx); e != nil {
						logger.Errorf(ctx, "Upload cleanup failed: %v", e)
					}
					lastCleanup = time.Now()
				}
				worked, e := s.ProcessNext(ctx)
				if e != nil {
					logger.Errorf(ctx, "Upload worker failed: %v", e)
				}
				if worked && e == nil {
					continue
				}
				select {
				case <-ctx.Done():
					return
				case <-timer.C:
				}
			}
		}()
	})
}

func (s *Service) ProcessNext(ctx context.Context) (bool, error) {
	if s.ingest == nil || s.workerContext == nil {
		return false, nil
	}
	var row Upload
	e := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var rows []Upload
		e := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).Where("state IN ('queued','verifying','processing') AND expires_at > ? AND (lease_until IS NULL OR lease_until < ?)", time.Now(), time.Now()).Order("updated_at, created_at").Limit(1).Find(&rows).Error
		if e != nil {
			return e
		}
		if len(rows) == 0 {
			return gorm.ErrRecordNotFound
		}
		row = rows[0]
		lease := time.Now().Add(2 * time.Minute)
		row.LeaseUntil = &lease
		row.LeaseToken = uuid.NewString()
		if row.State == "queued" {
			row.State = "verifying"
		}
		return tx.Model(&row).Updates(map[string]any{"state": row.State, "lease_until": lease, "lease_token": row.LeaseToken, "updated_at": time.Now()}).Error
	})
	if errors.Is(e, gorm.ErrRecordNotFound) {
		return false, nil
	}
	if e != nil {
		return false, e
	}
	s.process(ctx, &row)
	return true, nil
}

func (s *Service) process(parent context.Context, row *Upload) {
	ctx, cancel := context.WithTimeout(parent, 20*time.Minute)
	defer cancel()
	heartbeatDone := make(chan struct{})
	go func() {
		defer close(heartbeatDone)
		ticker := time.NewTicker(20 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				r := s.db.WithContext(ctx).Model(&Upload{}).Where("id=? AND lease_token=? AND state IN ('verifying','processing')", row.ID, row.LeaseToken).Updates(map[string]any{"lease_until": time.Now().Add(2 * time.Minute), "expires_at": time.Now().Add(7 * 24 * time.Hour), "updated_at": time.Now()})
				if r.Error != nil || r.RowsAffected != 1 {
					cancel()
					return
				}
			}
		}
	}()
	defer func() { cancel(); <-heartbeatDone }()
	finish := func(failure error) {
		finishCtx, done := context.WithTimeout(context.WithoutCancel(parent), 10*time.Second)
		defer done()
		logger.Errorf(finishCtx, "Upload %s processing failed: %v", row.ID, failure)
		detail := publicError(failure)
		if len(detail) > 2000 {
			detail = detail[:2000]
		}
		if e := s.db.WithContext(finishCtx).Model(&Upload{}).Where("id=? AND lease_token=?", row.ID, row.LeaseToken).Updates(map[string]any{"state": "failed", "error": detail, "lease_until": nil, "lease_token": "", "updated_at": time.Now()}).Error; e != nil {
			logger.Errorf(finishCtx, "Upload finalization failed: %v", e)
		}
	}
	defer func() {
		if recovered := recover(); recovered != nil {
			logger.Errorf(ctx, "Upload worker panicked: %v", recovered)
			finish(UserError("上传解析意外中断，请重新提交"))
		}
	}()
	workerCtx, e := s.workerContext(ctx, row)
	if e != nil {
		finish(e)
		return
	}
	if e = s.authorize(workerCtx, row.Kind, row.TargetID); e != nil {
		finish(e)
		return
	}
	if row.State == "verifying" {
		hash, e := s.verifyFile(ctx, row)
		if e != nil {
			finish(e)
			return
		}
		if row.SHA256 != "" && row.SHA256 != hash {
			finish(UserError("完整文件校验失败，请重新上传"))
			return
		}
		row.SHA256 = hash
		row.State = "processing"
		result := s.db.WithContext(ctx).Model(&Upload{}).Where("id=? AND lease_token=?", row.ID, row.LeaseToken).Updates(map[string]any{"state": "processing", "sha256": hash, "updated_at": time.Now()})
		if result.Error != nil {
			finish(result.Error)
			return
		}
		if result.RowsAffected != 1 {
			return
		}
	}
	offset, result, done, e := s.ingest(workerCtx, row, s.path(row.ID))
	if e != nil {
		finish(e)
		return
	}
	if offset < row.IngestOffset || offset > row.Size || (!done && offset == row.IngestOffset) {
		finish(UserError("解析进度无效"))
		return
	}
	state := "processing"
	if done {
		state = "completed"
	}
	update := s.db.WithContext(ctx).Model(&Upload{}).Where("id=? AND lease_token=?", row.ID, row.LeaseToken).Updates(map[string]any{"state": state, "ingest_offset": offset, "result": result, "part": row.Part + 1, "lease_until": nil, "lease_token": "", "expires_at": time.Now().Add(7 * 24 * time.Hour), "updated_at": time.Now()})
	if update.Error != nil {
		finish(update.Error)
	}
}

// Revalidate the persisted chunks while computing the full digest. This also
// detects storage corruption when a browser cannot provide a whole-file digest
// without first reading 20 GiB. Chunk metadata is fetched in bounded pages.
func (s *Service) verifyFile(ctx context.Context, row *Upload) (string, error) {
	f, err := os.Open(s.path(row.ID))
	if err != nil {
		return "", err
	}
	defer f.Close()
	stat, err := f.Stat()
	if err != nil {
		return "", err
	}
	if stat.Size() != row.Size {
		return "", UserError("完整文件长度校验失败，请重新上传")
	}
	whole := sha256.New()
	offset := int64(0)
	for offset < row.Size {
		var chunks []Chunk
		if err := s.db.WithContext(ctx).Where("upload_id=? AND \"offset\">=?", row.ID, offset).Order("\"offset\"").Limit(128).Find(&chunks).Error; err != nil {
			return "", err
		}
		if len(chunks) == 0 {
			return "", UserError("上传分块记录缺失，请重新上传")
		}
		for _, chunk := range chunks {
			if chunk.Offset != offset || chunk.Size != min(row.ChunkSize, row.Size-offset) {
				return "", UserError("上传分块长度不一致，请重新上传")
			}
			h := sha256.New()
			n, err := io.CopyN(io.MultiWriter(whole, h), &contextReader{ctx: ctx, r: f}, chunk.Size)
			if err != nil {
				return "", err
			}
			if n != chunk.Size || hex.EncodeToString(h.Sum(nil)) != chunk.SHA256 {
				return "", UserError("上传分块校验失败，请重新上传")
			}
			offset += n
		}
	}
	return hex.EncodeToString(whole.Sum(nil)), nil
}

type contextReader struct {
	ctx context.Context
	r   io.Reader
}

func (r *contextReader) Read(p []byte) (int, error) {
	if e := r.ctx.Err(); e != nil {
		return 0, e
	}
	return r.r.Read(p)
}