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

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

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