Skip to content

Commit

Permalink
Add peer bootstrapping integration test
Browse files Browse the repository at this point in the history
  • Loading branch information
anacrolix committed Jun 27, 2024
1 parent 1d30609 commit dcfb61f
Show file tree
Hide file tree
Showing 3 changed files with 136 additions and 1 deletion.
1 change: 1 addition & 0 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,7 @@ func NewClient(cfg *ClientConfig) (cl *Client, err error) {
cl.init(cfg)
go cl.acceptLimitClearer()
cl.initLogger()
//cl.logger.Levelf(log.Critical, "test after init")
defer func() {
if err != nil {
cl.Close()
Expand Down
133 changes: 133 additions & 0 deletions tests/peers-bootstrapping/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
package main

import (
"crypto/rand"
"fmt"
_ "github.com/anacrolix/envpprof"
"github.com/anacrolix/log"
"github.com/anacrolix/sync"
"github.com/anacrolix/torrent"
"github.com/anacrolix/torrent/bencode"
"github.com/anacrolix/torrent/metainfo"
"github.com/dustin/go-humanize"
"golang.org/x/exp/slog"
"io"
"net/http"
"os"
"path/filepath"
"time"
)

func assertNil(x any) {
if x != nil {
panic(x)
}
}

func newClientConfig() *torrent.ClientConfig {
cfg := torrent.NewDefaultClientConfig()
cfg.ListenPort = 0
cfg.NoDHT = true
cfg.NoDefaultPortForwarding = true
cfg.Seed = true
cfg.Debug = false
return cfg
}

func main() {
tmpDir, err := os.MkdirTemp("", "peers-bootstrapping")
assertNil(err)
slog.Info("made temp dir", slog.String("tmpDir", tmpDir))
sourceDir := filepath.Join(tmpDir, "source")
assertNil(os.Mkdir(sourceDir, 0o700))
f, err := os.Create(filepath.Join(sourceDir, "file"))
assertNil(err)
_, err = io.CopyN(f, rand.Reader, 1<<30)
assertNil(err)
assertNil(f.Close())
var info metainfo.Info
err = info.BuildFromFilePath(f.Name())
assertNil(err)
var mi metainfo.MetaInfo
mi.InfoBytes, err = bencode.Marshal(info)
assertNil(err)
var clients []*torrent.Client
var torrents []*torrent.Torrent
clientConfig := newClientConfig()
clientConfig.DataDir = sourceDir
initialClient, err := torrent.NewClient(clientConfig)
assertNil(err)
clientIndex := 0
addClientAndTorrent := func(cl *torrent.Client, t *torrent.Torrent) int {
clients = append(clients, cl)
torrents = append(torrents, t)
ret := clientIndex
http.HandleFunc(
fmt.Sprintf("/%v", ret),
func(w http.ResponseWriter, r *http.Request) {
cl.WriteStatus(w)
})
clientIndex++
return ret
}
initialTorrent, err := initialClient.AddTorrent(&mi)
assertNil(err)
addClientAndTorrent(initialClient, initialTorrent)
//initialTorrent.VerifyData()
<-initialTorrent.Complete().On()
var allDownloaded sync.WaitGroup
var notCompleted sync.Map
http.HandleFunc(
"/notCompleted",
func(w http.ResponseWriter, r *http.Request) {
notCompleted.Range(func(key, value any) bool {
fmt.Fprintln(w, key)
return true
})
})
for range 5 {
clientIndex := clientIndex
storageDir := filepath.Join(tmpDir, fmt.Sprintf("client%v", clientIndex))
clientConfig := newClientConfig()
clientConfig.DataDir = storageDir
clientConfig.Logger = log.Default.WithValues(slog.Int("clientIndex", clientIndex))
//clientConfig.Logger.Levelf(log.Critical, "test")
client, err := torrent.NewClient(clientConfig)
assertNil(err)
t, _ := client.AddTorrentInfoHash(mi.HashInfoBytes())
addClientAndTorrent(client, t)
allDownloaded.Add(1)
notCompleted.Store(clientIndex, nil)
go func() {
<-t.GotInfo()
t.DownloadAll()
<-t.Complete().On()
notCompleted.Delete(clientIndex)
slog.Info("leecher completed", slog.Int("clientIndex", clientIndex))
allDownloaded.Done()
}()
t.AddClientPeer(initialClient)
}
go func() {
for range time.Tick(time.Second) {
for _, t := range torrents {
for _, cl := range clients {
t.AddClientPeer(cl)
}
}
}
}()
allDownloaded.Wait()
slog.Info("all leechers downloaded")
for clientIndex, cl := range clients {
stats := cl.Stats()
written := stats.BytesWritten
read := stats.BytesRead
fmt.Printf(
"client %v wrote %v read %v\n",
clientIndex,
humanize.Bytes(uint64(written.Int64())),
humanize.Bytes(uint64(read.Int64())),
)
}
}
3 changes: 2 additions & 1 deletion torrent.go
Original file line number Diff line number Diff line change
Expand Up @@ -520,7 +520,8 @@ func (t *Torrent) setInfo(info *metainfo.Info) error {
t.nameMu.Lock()
t.info = info
t.nameMu.Unlock()
t._chunksPerRegularPiece = chunkIndexType((pp.Integer(t.usualPieceSize()) + t.chunkSize - 1) / t.chunkSize)
t._chunksPerRegularPiece = chunkIndexType(
(pp.Integer(t.usualPieceSize()) + t.chunkSize - 1) / t.chunkSize)
t.updateComplete()
t.displayName = "" // Save a few bytes lol.
t.initFiles()
Expand Down

0 comments on commit dcfb61f

Please sign in to comment.