github.com/anacrolix/torrent@v1.61.0/tests/peers-bootstrapping/main.go (about)

     1  package main
     2  
     3  import (
     4  	"crypto/rand"
     5  	"fmt"
     6  	"io"
     7  	"log/slog"
     8  	"net/http"
     9  	"os"
    10  	"path/filepath"
    11  	"time"
    12  
    13  	"github.com/anacrolix/envpprof"
    14  	"github.com/anacrolix/log"
    15  	"github.com/anacrolix/sync"
    16  	"github.com/dustin/go-humanize"
    17  
    18  	"github.com/anacrolix/torrent"
    19  	"github.com/anacrolix/torrent/bencode"
    20  	"github.com/anacrolix/torrent/metainfo"
    21  	"github.com/anacrolix/torrent/storage"
    22  )
    23  
    24  func assertNil(x any) {
    25  	if x != nil {
    26  		panic(x)
    27  	}
    28  }
    29  
    30  func newClientConfig() *torrent.ClientConfig {
    31  	cfg := torrent.NewDefaultClientConfig()
    32  	cfg.ListenPort = 0
    33  	cfg.NoDHT = true
    34  	cfg.NoDefaultPortForwarding = true
    35  	cfg.Seed = true
    36  	cfg.Debug = false
    37  	return cfg
    38  }
    39  
    40  func main() {
    41  	defer envpprof.Stop()
    42  	tmpDir, err := os.MkdirTemp("", "peers-bootstrapping")
    43  	assertNil(err)
    44  	slog.Info("made temp dir", slog.String("tmpDir", tmpDir))
    45  	sourceDir := filepath.Join(tmpDir, "source")
    46  	assertNil(os.Mkdir(sourceDir, 0o700))
    47  	f, err := os.Create(filepath.Join(sourceDir, "file"))
    48  	assertNil(err)
    49  	_, err = io.CopyN(f, rand.Reader, 1<<30)
    50  	assertNil(err)
    51  	assertNil(f.Close())
    52  	var info metainfo.Info
    53  	err = info.BuildFromFilePath(f.Name())
    54  	assertNil(err)
    55  	var mi metainfo.MetaInfo
    56  	mi.InfoBytes, err = bencode.Marshal(info)
    57  	assertNil(err)
    58  	var clients []*torrent.Client
    59  	var torrents []*torrent.Torrent
    60  	clientConfig := newClientConfig()
    61  	clientConfig.DefaultStorage = storage.NewMMap(sourceDir)
    62  	initialClient, err := torrent.NewClient(clientConfig)
    63  	assertNil(err)
    64  	clientIndex := 0
    65  	addClientAndTorrent := func(cl *torrent.Client, t *torrent.Torrent) int {
    66  		clients = append(clients, cl)
    67  		torrents = append(torrents, t)
    68  		ret := clientIndex
    69  		http.HandleFunc(
    70  			fmt.Sprintf("/%v", ret),
    71  			func(w http.ResponseWriter, r *http.Request) {
    72  				cl.WriteStatus(w)
    73  			})
    74  		clientIndex++
    75  		return ret
    76  	}
    77  	initialTorrent, err := initialClient.AddTorrent(&mi)
    78  	assertNil(err)
    79  	addClientAndTorrent(initialClient, initialTorrent)
    80  	//initialTorrent.VerifyData()
    81  	<-initialTorrent.Complete().On()
    82  	var allDownloaded sync.WaitGroup
    83  	var notCompleted sync.Map
    84  	http.HandleFunc(
    85  		"/notCompleted",
    86  		func(w http.ResponseWriter, r *http.Request) {
    87  			notCompleted.Range(func(key, value any) bool {
    88  				fmt.Fprintln(w, key)
    89  				return true
    90  			})
    91  		})
    92  	for range 5 {
    93  		clientIndex := clientIndex
    94  		storageDir := filepath.Join(tmpDir, fmt.Sprintf("client%v", clientIndex))
    95  		clientConfig := newClientConfig()
    96  		clientConfig.DefaultStorage = storage.NewMMap(storageDir)
    97  		clientConfig.Logger = log.Default.WithValues(slog.Int("clientIndex", clientIndex))
    98  		//clientConfig.Logger.Levelf(log.Critical, "test")
    99  		client, err := torrent.NewClient(clientConfig)
   100  		assertNil(err)
   101  		t, _ := client.AddTorrentInfoHash(mi.HashInfoBytes())
   102  		addClientAndTorrent(client, t)
   103  		allDownloaded.Add(1)
   104  		notCompleted.Store(clientIndex, nil)
   105  		go func() {
   106  			<-t.GotInfo()
   107  			t.DownloadAll()
   108  			<-t.Complete().On()
   109  			notCompleted.Delete(clientIndex)
   110  			slog.Info("leecher completed", slog.Int("clientIndex", clientIndex))
   111  			allDownloaded.Done()
   112  		}()
   113  		t.AddClientPeer(initialClient)
   114  	}
   115  	go func() {
   116  		for range time.Tick(time.Second) {
   117  			for _, t := range torrents {
   118  				for _, cl := range clients {
   119  					t.AddClientPeer(cl)
   120  				}
   121  			}
   122  		}
   123  	}()
   124  	allDownloaded.Wait()
   125  	slog.Info("all leechers downloaded")
   126  	for clientIndex, cl := range clients {
   127  		stats := cl.Stats()
   128  		written := stats.BytesWritten
   129  		read := stats.BytesRead
   130  		fmt.Printf(
   131  			"client %v wrote %v read %v\n",
   132  			clientIndex,
   133  			humanize.Bytes(uint64(written.Int64())),
   134  			humanize.Bytes(uint64(read.Int64())),
   135  		)
   136  	}
   137  	assertNil(os.RemoveAll(tmpDir))
   138  }