diff --git a/cmd/torrent/main.go b/cmd/torrent/main.go index 0bbd0ec4..4adda93e 100644 --- a/cmd/torrent/main.go +++ b/cmd/torrent/main.go @@ -14,6 +14,7 @@ import ( _ "github.com/anacrolix/envpprof" "github.com/anacrolix/tagflag" "github.com/dustin/go-humanize" + "github.com/gosuri/uiprogress" "github.com/anacrolix/torrent" "github.com/anacrolix/torrent/data/mmap" @@ -35,73 +36,39 @@ func resolvedPeerAddrs(ss []string) (ret []torrent.Peer, err error) { return } -func bytesCompleted(tc *torrent.Client) (ret int64) { - for _, t := range tc.Torrents() { - if t.Info() != nil { - ret += t.BytesCompleted() +func torrentBar(t torrent.Torrent) { + bar := uiprogress.AddBar(1) + bar.AppendCompleted() + bar.AppendFunc(func(*uiprogress.Bar) (ret string) { + select { + case <-t.GotInfo(): + default: + return "getting info" } - } - return -} - -// Returns an estimate of the total bytes for all torrents. -func totalBytesEstimate(tc *torrent.Client) (ret int64) { - var noInfo, hadInfo int64 - for _, t := range tc.Torrents() { - info := t.Info() - if info == nil { - noInfo++ - continue + if t.Seeding() { + return "seeding" + } else if t.BytesCompleted() == t.Info().TotalLength() { + return "completed" + } else { + return fmt.Sprintf("downloading (%s/%s)", humanize.Bytes(uint64(t.BytesCompleted())), humanize.Bytes(uint64(t.Info().TotalLength()))) } - ret += info.TotalLength() - hadInfo++ - } - if hadInfo != 0 { - // Treat each torrent without info as the average of those with, - // rounded up. - ret += (noInfo*ret + hadInfo - 1) / hadInfo - } - return -} - -func progressLine(tc *torrent.Client) string { - return fmt.Sprintf("\033[K%s / %s\r", humanize.Bytes(uint64(bytesCompleted(tc))), humanize.Bytes(uint64(totalBytesEstimate(tc)))) -} - -func main() { - log.SetFlags(log.LstdFlags | log.Lshortfile) - var opts struct { - torrent.Config `name:"Client"` - Mmap bool `help:"memory-map torrent data"` - TestPeer []*net.TCPAddr `short:"p" help:"addresses of some starting peers"` - Torrent []string `type:"pos" arity:"+" help:"torrent file path or magnet uri"` - } - tagflag.Parse(&opts, tagflag.SkipBadTypes()) - clientConfig := opts.Config - if opts.Mmap { - clientConfig.TorrentDataOpener = func(info *metainfo.Info) torrent.Data { - ret, err := mmap.TorrentData(info, "") - if err != nil { - log.Fatalf("error opening torrent data for %q: %s", info.Name, err) - } - return ret - } - } - - torrents := opts.Torrent - if len(torrents) == 0 { - fmt.Fprintf(os.Stderr, "no torrents specified\n") - return - } - client, err := torrent.NewClient(&clientConfig) - if err != nil { - log.Fatalf("error creating client: %s", err) - } - http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { - client.WriteStatus(w) }) - defer client.Close() - for _, arg := range torrents { + bar.PrependFunc(func(*uiprogress.Bar) string { + return t.Name() + }) + go func() { + <-t.GotInfo() + bar.Total = int(t.Info().TotalLength()) + for { + bc := t.BytesCompleted() + bar.Set(int(bc)) + time.Sleep(time.Second) + } + }() +} + +func addTorrents(client *torrent.Client) { + for _, arg := range opts.Torrent { t := func() torrent.Torrent { if strings.HasPrefix(arg, "magnet:") { t, err := client.AddMagnet(arg) @@ -122,6 +89,7 @@ func main() { return t } }() + torrentBar(t) err := t.AddPeers(func() (ret []torrent.Peer) { for _, ta := range opts.TestPeer { ret = append(ret, torrent.Peer{ @@ -139,25 +107,44 @@ func main() { t.DownloadAll() }() } - done := make(chan struct{}) - go func() { - defer close(done) - if client.WaitAll() { - log.Print("downloaded ALL the torrents") - } else { - log.Fatal("y u no complete torrents?!") - } - }() - ticker := time.NewTicker(time.Second) -waitDone: - for { - select { - case <-done: - break waitDone - case <-ticker.C: - os.Stdout.WriteString(progressLine(client)) +} + +var opts struct { + torrent.Config `name:"Client"` + Mmap bool `help:"memory-map torrent data"` + TestPeer []*net.TCPAddr `short:"p" help:"addresses of some starting peers"` + Torrent []string `type:"pos" arity:"+" help:"torrent file path or magnet uri"` +} + +func main() { + log.SetFlags(log.LstdFlags | log.Lshortfile) + tagflag.Parse(&opts, tagflag.SkipBadTypes()) + clientConfig := opts.Config + if opts.Mmap { + clientConfig.TorrentDataOpener = func(info *metainfo.Info) torrent.Data { + ret, err := mmap.TorrentData(info, "") + if err != nil { + log.Fatalf("error opening torrent data for %q: %s", info.Name, err) + } + return ret } } + + client, err := torrent.NewClient(&clientConfig) + if err != nil { + log.Fatalf("error creating client: %s", err) + } + defer client.Close() + http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { + client.WriteStatus(w) + }) + uiprogress.Start() + addTorrents(client) + if client.WaitAll() { + log.Print("downloaded ALL the torrents") + } else { + log.Fatal("y u no complete torrents?!") + } if opts.Seed { select {} }