fix(download): separate torrent progress from transfer rates (#1654)

Co-authored-by: xunguangzlj-cloud <[email protected]>

internal/fetcher/fetcher.go

@@ -46,6 +46,10 @@ type Uploader interface {
 	WaitUpload() error
 }
 
+type DownloadCounter interface {
+	ReceivedBytes() int64
+}
+
 // FetcherMeta defines the meta information of a fetcher.
 type FetcherMeta struct {
 	Req  *base.Request  `json:"req"`

internal/protocol/bt/download_test.go

@@ -290,6 +290,82 @@ func TestBTDownloadSelectedFiles(t *testing.T) {
 	if f.Meta().Res.Size != 48*1024 {
 		t.Fatal("selected resource size is incorrect")
 	}
+	if downloaded := f.Progress().TotalDownloaded(); downloaded != 48*1024 || f.ReceivedBytes() <= downloaded {
+		t.Fatalf("selected progress = %d, received traffic = %d; want 48 KiB of progress and traffic across the shared piece", downloaded, f.ReceivedBytes())
+	}
+}
+
+func TestBTProgressOnlyCountsVerifiedFileBytes(t *testing.T) {
+	const fileSize = 56 * 1024
+	fixture := makeBTFixtureWithPieceSize(t, map[string][]byte{
+		"a.bin": btPayload(fileSize, "verified-a"),
+		"b.bin": btPayload(fileSize, "verified-b"),
+	}, 16*1024, false, true, 64*1024)
+	f := btDownloadFetcher(t, nil)
+	btResolve(t, f, fixture.Magnet)
+	dir := t.TempDir()
+	f.Meta().Opts = &base.Options{Path: dir}
+	btStart(t, f)
+	deadline := time.Now().Add(8 * time.Second)
+	for f.torrent.Files()[0].BytesCompleted() == 0 && time.Now().Before(deadline) {
+		time.Sleep(10 * time.Millisecond)
+	}
+	btPause(t, f)
+	if f.torrent.Files()[0].BytesCompleted() == 0 || f.torrent.Piece(0).State().Complete {
+		t.Fatal("expected some received bytes in a piece that is still incomplete")
+	}
+	if f.ReceivedBytes() == 0 {
+		t.Fatal("received traffic must be visible before a piece is verified")
+	}
+	if progress := f.Progress(); len(progress) != 2 || progress[0] != 0 {
+		t.Errorf("progress before the first file's piece is verified = %v, want zero for the first file", progress)
+	}
+	fixture.Limiter.SetLimit(rate.Inf)
+	btStart(t, f)
+	btAwaitDone(t, f)
+	fixture.AssertFiles(t, dir)
+	progress := f.Progress()
+	if len(progress) != 2 || progress[0] != fileSize || progress[1] != fileSize {
+		t.Fatalf("verified progress at completion = %v, want %d for each file", progress, fileSize)
+	}
+}
+
+func TestBTUploadedBytesSamplesLiveTraffic(t *testing.T) {
+	payload := btPayload(48*1024, "upload-counter")
+	fixture := newBTFixture(t, map[string][]byte{"": payload}, 0)
+	const saved = 1024
+	f := &Fetcher{
+		torrent:       fixture.Torrent,
+		data:          &fetcherData{SeedBytes: saved},
+		seedBaseBytes: saved,
+	}
+	f.torrentReady.Store(true)
+	f.torrentUpload.Store(true)
+
+	cfg := torrent.TestingConfig(t)
+	cfg.DataDir = t.TempDir()
+	cfg.DisableTrackers = true
+	cfg.DisableUTP = true
+	cfg.DisableIPv6 = true
+	peer, err := torrent.NewClient(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { peer.Close() })
+	download, _, err := peer.AddTorrentSpec(torrent.TorrentSpecFromMetaInfo(fixture.MetaInfo))
+	if err != nil {
+		t.Fatal(err)
+	}
+	download.AddPeers([]torrent.PeerInfo{{Addr: fixture.Seeder.ListenAddrs()[0]}})
+	download.DownloadAll()
+	select {
+	case <-download.Complete().On():
+	case <-time.After(10 * time.Second):
+		t.Fatal("loopback peer did not complete")
+	}
+	if got, want := f.UploadedBytes(), int64(saved+len(payload)); got != want {
+		t.Fatalf("uploaded bytes = %d, want %d including the saved bytes", got, want)
+	}
 }
 
 func TestBTPauseAndResume(t *testing.T) {

internal/protocol/bt/fetcher.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -42,9 +43,10 @@ type Fetcher struct {
 	ctl    *controller.Controller
 	config *config
 
-	torrent *torrent.Torrent
-	meta    *fetcher.FetcherMeta
-	data    *fetcherData
+	torrent       *torrent.Torrent
+	meta          *fetcher.FetcherMeta
+	data          *fetcherData
+	seedBaseBytes int64
 
 	torrentReady    atomic.Bool
 	torrentUpload   atomic.Bool
@@ -463,14 +465,46 @@ func (f *Fetcher) Progress() fetcher.Progress {
 	if !f.torrentReady.Load() {
 		return append(fetcher.Progress(nil), f.data.Progress...)
 	}
+	pieceSize := f.torrent.Info().PieceLength
+	var offset int64
+	var verified []verifiedByteRange
+	for _, run := range f.torrent.PieceStateRuns() {
+		end := offset + int64(run.Length)*pieceSize
+		if pieceVerified(run.PieceState) {
+			verified = append(verified, verifiedByteRange{start: offset, end: end})
+		}
+		offset = end
+	}
+	files := f.torrent.Files()
 	for i := range f.data.Progress {
 		selectIndex := f.meta.Opts.SelectFiles[i]
-		file := f.torrent.Files()[selectIndex]
-		f.data.Progress[i] = file.BytesCompleted()
+		file := files[selectIndex]
+		f.data.Progress[i] = verifiedFileBytes(verified, file.Offset(), file.Length())
 	}
 	return append(fetcher.Progress(nil), f.data.Progress...)
 }
 
+type verifiedByteRange struct {
+	start, end int64
+}
+
+func verifiedFileBytes(verified []verifiedByteRange, offset, length int64) int64 {
+	first := sort.Search(len(verified), func(i int) bool { return verified[i].end > offset })
+	fileEnd := offset + length
+	var completed int64
+	for _, run := range verified[first:] {
+		if run.start >= fileEnd {
+			break
+		}
+		completed += min(fileEnd, run.end) - max(offset, run.start)
+	}
+	return completed
+}
+
+func pieceVerified(state torrent.PieceState) bool {
+	return state.Ok && state.Complete && !state.Hashing && !state.QueuedForHash
+}
+
 func (f *Fetcher) Wait() error {
 	f.mu.Lock()
 	var closed <-chan struct{}
@@ -587,7 +621,7 @@ func (f *Fetcher) isDoneLocked() bool {
 		file := f.torrent.Files()[index]
 		for i := file.BeginPieceIndex(); i < file.EndPieceIndex(); i++ {
 			state := f.torrent.Piece(i).State()
-			if !state.Ok || !state.Complete || state.Hashing || state.QueuedForHash {
+			if !pieceVerified(state) {
 				return false
 			}
 		}
@@ -685,11 +719,13 @@ func (f *Fetcher) Upload() (err error) {
 }
 
 func (f *Fetcher) doUpload(fromUpload bool) {
+	f.mu.Lock()
 	if !f.torrentUpload.CompareAndSwap(false, true) {
+		f.mu.Unlock()
 		return
 	}
-	f.mu.Lock()
-	initialBytes, initialTime := f.data.SeedBytes, f.data.SeedTime
+	f.seedBaseBytes = f.data.SeedBytes
+	initialTime := f.data.SeedTime
 	torrentClosed := f.torrent.Closed()
 	f.mu.Unlock()
 	var doneTime int64
@@ -707,8 +743,7 @@ func (f *Fetcher) doUpload(fromUpload bool) {
 				f.mu.Unlock()
 				continue
 			}
-			stats := f.torrentStats()
-			f.data.SeedBytes = initialBytes + stats.BytesWrittenData.Int64()
+			f.refreshSeedBytesLocked()
 			if !fromUpload && !f.isDoneLocked() {
 				f.mu.Unlock()
 				continue
@@ -741,9 +776,27 @@ func (f *Fetcher) torrentStats() torrent.TorrentStats {
 func (f *Fetcher) UploadedBytes() int64 {
 	f.mu.Lock()
 	defer f.mu.Unlock()
+	if f.torrentReady.Load() && f.torrentUpload.Load() {
+		f.refreshSeedBytesLocked()
+	}
 	return f.data.SeedBytes
 }
 
+func (f *Fetcher) refreshSeedBytesLocked() {
+	stats := f.torrentStats()
+	f.data.SeedBytes = max(f.data.SeedBytes, f.seedBaseBytes+stats.BytesWrittenData.Int64())
+}
+
+func (f *Fetcher) ReceivedBytes() int64 {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if !f.torrentReady.Load() {
+		return 0
+	}
+	stats := f.torrentStats()
+	return stats.BytesReadUsefulData.Int64()
+}
+
 func (f *Fetcher) WaitUpload() (err error) {
 	<-f.uploadDoneCh
 	return nil

internal/protocol/bt/fixture_test.go

@@ -46,6 +46,11 @@ func newBTHybridFixture(t testing.TB, payload []byte) *btFixture {
 }
 
 func makeBTFixture(t testing.TB, files map[string][]byte, bytesPerSecond int, v2, private bool) *btFixture {
+	t.Helper()
+	return makeBTFixtureWithPieceSize(t, files, bytesPerSecond, v2, private, 16*1024)
+}
+
+func makeBTFixtureWithPieceSize(t testing.TB, files map[string][]byte, bytesPerSecond int, v2, private bool, pieceSize int64) *btFixture {
 	t.Helper()
 	source := t.TempDir()
 	name := "payload.bin"
@@ -64,7 +69,7 @@ func makeBTFixture(t testing.TB, files map[string][]byte, bytesPerSecond int, v2
 		}
 		expected[filepath.Join(name, filepath.FromSlash(path))] = append([]byte(nil), data...)
 	}
-	info := metainfo.Info{PieceLength: 16 * 1024}
+	info := metainfo.Info{PieceLength: pieceSize}
 	if err := info.BuildFromFilePath(root); err != nil {
 		t.Fatal(err)
 	}

pkg/download/downloader.go

@@ -345,7 +345,12 @@ func (d *Downloader) Setup() error {
 						}
 
 						current := task.fetcher.Progress().TotalDownloaded()
-						delta := current - task.Progress.Downloaded
+						delta := max(0, current-task.Progress.Downloaded)
+						if counter, ok := task.fetcher.(fetcher.DownloadCounter); ok && task.Status == base.DownloadStatusRunning {
+							received := counter.ReceivedBytes()
+							delta = received - task.receivedBytes
+							task.receivedBytes = received
+						}
 						if task.Meta != nil && task.Meta.Req != nil {
 							if producer := d.blob.Production(task.Meta.Req.URL); producer != nil {
 								downloaded, received := producer.Progress()
@@ -353,19 +358,20 @@ func (d *Downloader) Setup() error {
 									task.producer = producer
 									task.producerReceived = 0
 									task.speedArr = nil
+									task.downloadSampledAt = time.Time{}
 								}
 								current = downloaded
 								delta = max(0, received-task.producerReceived)
 								task.producerReceived = received
 							}
 						}
-						tick := float64(d.cfg.RefreshInterval) / 1000
 						downloadDataChanged := false
 						uploading := task.Uploading
 						if task.Status == base.DownloadStatusRunning {
 							downloadDataChanged = current != task.Progress.Downloaded
 							task.Progress.Used = task.timer.Used()
-							task.Progress.Speed = task.updateSpeed(max(0, delta), tick)
+							elapsed := sampleSeconds(&task.downloadSampledAt, time.Now())
+							task.Progress.Speed = task.updateSpeed(delta, elapsed)
 							task.Progress.Downloaded = current
 						}
 
@@ -374,7 +380,8 @@ func (d *Downloader) Setup() error {
 							uploader := task.fetcher.(fetcher.Uploader)
 							currentUploaded := uploader.UploadedBytes()
 							uploadDataChanged = currentUploaded != task.Progress.Uploaded
-							task.Progress.UploadSpeed = task.updateUploadSpeed(currentUploaded-task.Progress.Uploaded, tick)
+							elapsed := sampleSeconds(&task.uploadSampledAt, time.Now())
+							task.Progress.UploadSpeed = task.updateUploadSpeed(currentUploaded-task.Progress.Uploaded, elapsed)
 							task.Progress.Uploaded = currentUploaded
 						}
 						task.statusLock.Unlock()
@@ -2453,6 +2460,8 @@ func (d *Downloader) doStart(task *Task) (err error) {
 			return nil
 		}
 		task.Progress.Speed = 0
+		task.speedArr = nil
+		task.downloadSampledAt = time.Time{}
 		task.timer.Start()
 		task.statusLock.Unlock()
 		if err := task.fetcher.Start(); err != nil {
@@ -3061,8 +3070,10 @@ func initTask(task *Task) {
 	task.lock = &sync.Mutex{}
 	task.statsLock = &sync.Mutex{}
 	task.blobRefLock = &sync.Mutex{}
-	task.speedArr = make([]int64, 0)
-	task.uploadSpeedArr = make([]int64, 0)
+	task.speedArr = nil
+	task.uploadSpeedArr = nil
+	task.downloadSampledAt = time.Time{}
+	task.uploadSampledAt = time.Time{}
 }
 
 func (d *Downloader) releaseBlobTask(task *Task) {

pkg/download/model.go

@@ -75,20 +75,23 @@ type Task struct {
 	producer         production.Source
 	producerReceived int64
 
-	fetcherManager fetcher.FetcherManager
-	fetcher        fetcher.Fetcher
-	timer          *util.Timer
-	statusLock     *sync.Mutex
-	lock           *sync.Mutex
-	statsLock      *sync.Mutex
-	statsLoaded    bool
-	stats          *fetcher.Stats
-	statsSnapshot  json.RawMessage
-	blobRefLock    *sync.Mutex
-	blobURL        string
-	runGeneration  uint64
-	speedArr       []int64
-	uploadSpeedArr []int64
+	fetcherManager    fetcher.FetcherManager
+	fetcher           fetcher.Fetcher
+	timer             *util.Timer
+	statusLock        *sync.Mutex
+	lock              *sync.Mutex
+	statsLock         *sync.Mutex
+	statsLoaded       bool
+	stats             *fetcher.Stats
+	statsSnapshot     json.RawMessage
+	blobRefLock       *sync.Mutex
+	blobURL           string
+	runGeneration     uint64
+	speedArr          []speedSample
+	uploadSpeedArr    []speedSample
+	downloadSampledAt time.Time
+	uploadSampledAt   time.Time
+	receivedBytes     int64
 }
 
 // TaskRuntimeStatus is the lightweight, frequently refreshed view of a task.
@@ -195,7 +198,21 @@ func (t *Task) updateUploadSpeed(downloaded int64, usedTime float64) int64 {
 	return calcSpeed(&t.uploadSpeedArr, downloaded, usedTime)
 }
 
-func calcSpeed(speedArr *[]int64, downloaded int64, usedTime float64) int64 {
+type speedSample struct {
+	bytes   int64
+	seconds float64
+}
+
+func sampleSeconds(lastSample *time.Time, now time.Time) float64 {
+	previous := *lastSample
+	*lastSample = now
+	if previous.IsZero() {
+		return 0
+	}
+	return now.Sub(previous).Seconds()
+}
+
+func calcSpeed(speedArr *[]speedSample, downloaded int64, usedTime float64) int64 {
 	if usedTime <= 0 {
 		return 0
 	}
@@ -204,18 +221,18 @@ func calcSpeed(speedArr *[]int64, downloaded int64, usedTime float64) int64 {
 		return 0
 	}
 
-	*speedArr = append(*speedArr, downloaded)
-	// Record last 5 seconds of download speed to calculate the average speed
-	if len(*speedArr) > int(5.0/usedTime) {
-		*speedArr = (*speedArr)[1:]
-	}
-
+	*speedArr = append(*speedArr, speedSample{bytes: downloaded, seconds: usedTime})
 	var total int64
-	for _, v := range *speedArr {
-		total += v
+	var elapsed float64
+	first := len(*speedArr)
+	for first > 0 && elapsed < 5 {
+		first--
+		sample := (*speedArr)[first]
+		total += sample.bytes
+		elapsed += sample.seconds
 	}
-
-	return int64(float64(total) / float64(len(*speedArr)) / usedTime)
+	*speedArr = (*speedArr)[first:]
+	return int64(float64(total) / elapsed)
 }
 
 type TaskFilter struct {

pkg/download/model_test.go

@@ -1,9 +1,32 @@
 package download
 
-import "testing"
+import (
+	"testing"
+	"time"
+)
+
+func TestTransferSpeedWithUnevenSamples(t *testing.T) {
+	var task Task
+	for _, sample := range []struct {
+		bytes   int64
+		seconds float64
+	}{
+		{1000, 1},
+		{2000, 2},
+		{4000, 4},
+		{6000, 6},
+	} {
+		if got := task.updateSpeed(sample.bytes, sample.seconds); got != 1000 {
+			t.Errorf("download speed after %v seconds = %d, want 1000", sample.seconds, got)
+		}
+		if got := task.updateUploadSpeed(sample.bytes, sample.seconds); got != 1000 {
+			t.Errorf("upload speed after %v seconds = %d, want 1000", sample.seconds, got)
+		}
+	}
+}
 
 func TestCalcSpeedResetOnRollback(t *testing.T) {
-	speedArr := []int64{1024, 2048, 4096}
+	speedArr := []speedSample{{1024, 1}, {2048, 1}, {4096, 1}}
 
 	if got := calcSpeed(&speedArr, -512, 1); got != 0 {
 		t.Fatalf("calcSpeed() = %d, want 0 after rollback", got)
@@ -16,3 +39,17 @@ func TestCalcSpeedResetOnRollback(t *testing.T) {
 		t.Fatalf("calcSpeed() = %d, want 1024 after reset", got)
 	}
 }
+
+func TestSampleSecondsUsesThePreviousSuccessfulSample(t *testing.T) {
+	start := time.Now()
+	var last time.Time
+	if got := sampleSeconds(&last, start); got != 0 {
+		t.Fatalf("first sample elapsed = %v, want 0", got)
+	}
+	if got := sampleSeconds(&last, start.Add(1500*time.Millisecond)); got != 1.5 {
+		t.Fatalf("elapsed after a delayed sample = %v, want 1.5", got)
+	}
+	if got := sampleSeconds(&last, start.Add(2*time.Second)); got != 0.5 {
+		t.Fatalf("elapsed after the next sample = %v, want 0.5", got)
+	}
+}