<commit_before>package test_storage

import (
	"bytes"
	"io"
	"io/ioutil"
	"math/rand"
	"sync"
	"testing"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
	"github.com/bradfitz/iter"
	qt "github.com/frankban/quicktest"
)

const (
	ChunkSize        = 1 << 14
	DefaultPieceSize = 2 << 20
	DefaultCapacity  = 0
	DefaultNumPieces = 16
)

func BenchmarkPieceMarkComplete(
	b *testing.B, ci storage.ClientImpl,
	pieceSize int64, numPieces int,
	// This drives any special handling around capacity that may be configured into the storage
	// implementation.
	capacity int64,
) {
	const check = true
	c := qt.New(b)
	info := &metainfo.Info{
		Pieces:      make([]byte, numPieces*metainfo.HashSize),
		PieceLength: pieceSize,
		Length:      pieceSize * int64(numPieces),
		Name:        "TorrentName",
	}
	ti, err := ci.OpenTorrent(info, metainfo.Hash{})
	c.Assert(err, qt.IsNil)
	defer ti.Close()
	rand.Read(info.Pieces)
	data := make([]byte, pieceSize)
	b.SetBytes(int64(numPieces) * pieceSize)
	oneIter := func() {
		for pieceIndex := range iter.N(numPieces) {
			pi := ti.Piece(info.Piece(pieceIndex))
			if check {
				rand.Read(data)
			}
			var wg sync.WaitGroup
			for off := int64(0); off < int64(len(data)); off += ChunkSize {
				wg.Add(1)
				go func(off int64) {
					defer wg.Done()
					n, err := pi.WriteAt(data[off:off+ChunkSize], off)
					if err != nil {
						panic(err)
					}
					if n != ChunkSize {
						panic(n)
					}
				}(off)
			}
			wg.Wait()
			b.StopTimer()
			if capacity == 0 {
				pi.MarkNotComplete()
			}
			b.StartTimer()
			// This might not apply if users of this benchmark don't cache with the expected capacity.
			c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
			c.Assert(pi.MarkComplete(), qt.IsNil)
			c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
			if check {
				readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
				c.Check(err, qt.IsNil)
				c.Assert(len(readData), qt.Equals, len(data))
				c.Assert(bytes.Equal(readData, data), qt.IsTrue)
			}
		}
	}
	// Fill the cache
	if capacity > 0 {
		for range iter.N(int((capacity + info.TotalLength() - 1) / info.TotalLength())) {
			oneIter()
		}
	}
	b.ResetTimer()
	for range iter.N(b.N) {
		oneIter()
	}
}
<commit_msg>Use bench timer more effectively, and use Piece.WriteTo<commit_after>package test_storage

import (
	"bytes"
	"math/rand"
	"sync"
	"testing"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
	"github.com/bradfitz/iter"
	qt "github.com/frankban/quicktest"
)

const (
	ChunkSize        = 1 << 14
	DefaultPieceSize = 2 << 20
	DefaultCapacity  = 0
	DefaultNumPieces = 16
)

func BenchmarkPieceMarkComplete(
	b *testing.B, ci storage.ClientImpl,
	pieceSize int64, numPieces int,
	// This drives any special handling around capacity that may be configured into the storage
	// implementation.
	capacity int64,
) {
	c := qt.New(b)
	info := &metainfo.Info{
		Pieces:      make([]byte, numPieces*metainfo.HashSize),
		PieceLength: pieceSize,
		Length:      pieceSize * int64(numPieces),
		Name:        "TorrentName",
	}
	ti, err := ci.OpenTorrent(info, metainfo.Hash{})
	c.Assert(err, qt.IsNil)
	tw := storage.Torrent{ti}
	defer tw.Close()
	rand.Read(info.Pieces)
	data := make([]byte, pieceSize)
	readData := make([]byte, pieceSize)
	b.SetBytes(int64(numPieces) * pieceSize)
	oneIter := func() {
		for pieceIndex := range iter.N(numPieces) {
			pi := tw.Piece(info.Piece(pieceIndex))
			rand.Read(data)
			b.StartTimer()
			var wg sync.WaitGroup
			for off := int64(0); off < int64(len(data)); off += ChunkSize {
				wg.Add(1)
				go func(off int64) {
					defer wg.Done()
					n, err := pi.WriteAt(data[off:off+ChunkSize], off)
					if err != nil {
						panic(err)
					}
					if n != ChunkSize {
						panic(n)
					}
				}(off)
			}
			wg.Wait()
			if capacity == 0 {
				pi.MarkNotComplete()
			}
			// This might not apply if users of this benchmark don't cache with the expected capacity.
			c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
			c.Assert(pi.MarkComplete(), qt.IsNil)
			c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
			n, err := pi.WriteTo(bytes.NewBuffer(readData[:0]))
			b.StopTimer()
			c.Assert(err, qt.IsNil)
			c.Assert(n, qt.Equals, int64(len(data)))
			c.Assert(bytes.Equal(readData[:n], data), qt.IsTrue)
		}
	}
	// Fill the cache
	if capacity > 0 {
		for range iter.N(int((capacity + info.TotalLength() - 1) / info.TotalLength())) {
			oneIter()
		}
	}
	b.ResetTimer()
	for range iter.N(b.N) {
		oneIter()
	}
}

{'fraction_non_alphanumeric': 0.14911903738719381, 'fraction_numerical': 0.013321873657069187, 'mean_word_length': 3.3181818181818183, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4614699', 'n_tokens_mistral': 2022, 'n_tokens_neox': 1769, 'n_words': 381}