Improve the interface with client and download strategies

diff --git a/client.go b/client.go
index 741416272435286f44e8b4ec3a09a1423da1dbf9..2be26443233e1aa1e3122d4c7adbce0b917c12d0 100644 (file)
--- a/client.go
+++ b/client.go
@@ -107,6 +107,12 @@ func (me *Client) PrioritizeDataRegion(ih InfoHash, off, len_ int64) error {
        if !t.haveInfo() {
                return errors.New("missing metadata")
        }
+       i := int(off / int64(t.UsualPieceSize()))
+       me.queueFirstHash(t, i)
+       i++
+       if i < t.NumPieces() {
+               me.queueFirstHash(t, i)
+       }
        me.downloadStrategy.TorrentPrioritize(t, off, len_)
        for _, cn := range t.Conns {
                me.replenishConnRequests(t, cn)
@@ -1430,6 +1436,18 @@ type Torrent struct {
        *torrent
 }
 
+func (t Torrent) AddPeers(pp []Peer) error {
+       return t.cl.AddPeers(t.torrent.InfoHash, pp)
+}
+
+func (t Torrent) DownloadAll() {
+       t.cl.mu.Lock()
+       for i := 0; i < t.NumPieces(); i++ {
+               t.cl.queueFirstHash(t.torrent, i)
+       }
+       t.cl.mu.Unlock()
+}
+
 func (me Torrent) ReadAt(p []byte, off int64) (n int, err error) {
        err = me.cl.PrioritizeDataRegion(me.InfoHash, off, int64(len(p)))
        if err != nil {
@@ -1525,27 +1543,28 @@ func (me *Client) addTorrent(t *torrent) (err error) {
 }
 
 // Adds the torrent to the client.
-func (me *Client) AddTorrent(metaInfo *metainfo.MetaInfo) (err error) {
+func (me *Client) AddTorrent(metaInfo *metainfo.MetaInfo) (t Torrent, err error) {
+       t.cl = me
        var ih InfoHash
        CopyExact(&ih, metaInfo.Info.Hash)
-       t, err := newTorrent(ih, metaInfo.AnnounceList, me.halfOpenLimit)
+       t.torrent, err = newTorrent(ih, metaInfo.AnnounceList, me.halfOpenLimit)
        if err != nil {
                return
        }
        me.mu.Lock()
        defer me.mu.Unlock()
-       err = me.addTorrent(t)
+       err = me.addTorrent(t.torrent)
        if err != nil {
                return
        }
-       err = me.setMetaData(t, metaInfo.Info.Info, metaInfo.Info.Bytes)
+       err = me.setMetaData(t.torrent, metaInfo.Info.Info, metaInfo.Info.Bytes)
        if err != nil {
                return
        }
        return
 }
 
-func (me *Client) AddTorrentFromFile(name string) (err error) {
+func (me *Client) AddTorrentFromFile(name string) (t Torrent, err error) {
        mi, err := metainfo.LoadFromFile(name)
        if err != nil {
                err = fmt.Errorf("error loading metainfo from file: %s", err)

diff --git a/cmd/torrent/main.go b/cmd/torrent/main.go
index af65f227d542f433ada8a1d8e398de90b3ce78da..121a7b837671a59717f43edadb7b0863149481f1 100644 (file)
--- a/cmd/torrent/main.go
+++ b/cmd/torrent/main.go
@@ -51,30 +51,27 @@ func main() {
                return
        }
        for _, arg := range flag.Args() {
-               var ih torrent.InfoHash
-               if strings.HasPrefix(arg, "magnet:") {
-                       m, err := torrent.ParseMagnetURI(arg)
-                       if err != nil {
-                               log.Fatalf("error parsing magnet uri: %s", err)
-                       }
-                       ih = m.InfoHash
-                       _, err = client.AddMagnet(arg)
-                       if err != nil {
-                               log.Fatalf("error adding magnet: %s", err)
-                       }
-               } else {
-                       metaInfo, err := metainfo.LoadFromFile(arg)
-                       if err != nil {
-                               log.Fatal(err)
+               t := func() torrent.Torrent {
+                       if strings.HasPrefix(arg, "magnet:") {
+                               t, err := client.AddMagnet(arg)
+                               if err != nil {
+                                       log.Fatalf("error adding magnet: %s", err)
+                               }
+                               return t
+                       } else {
+                               metaInfo, err := metainfo.LoadFromFile(arg)
+                               if err != nil {
+                                       log.Fatal(err)
+                               }
+                               t, err := client.AddTorrent(metaInfo)
+                               if err != nil {
+                                       log.Fatal(err)
+                               }
+                               return t
                        }
-                       err = client.AddTorrent(metaInfo)
-                       if err != nil {
-                               log.Fatal(err)
-                       }
-                       util.CopyExact(&ih, metaInfo.Info.Hash)
-               }
+               }()
                // client.PrioritizeDataRegion(ih, 0, 999999999)
-               err := client.AddPeers(ih, func() []torrent.Peer {
+               err := t.AddPeers(func() []torrent.Peer {
                        if *testPeer == "" {
                                return nil
                        }
@@ -90,6 +87,10 @@ func main() {
                if err != nil {
                        log.Fatal(err)
                }
+               go func() {
+                       <-t.GotMetainfo
+                       t.DownloadAll()
+               }()
        }
        if *seed {
                select {}