Get filesizes from DHT metadata

main.go

package main;

import (
	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/ewhal/nyaa/config"
	"github.com/ewhal/nyaa/db"
	"github.com/ewhal/nyaa/model"
	serviceBase "github.com/ewhal/nyaa/service"
	torrentService "github.com/ewhal/nyaa/service/torrent"
	"github.com/ewhal/nyaa/util"
	"errors"
	"fmt"
	"flag"
	"time"
	"strings"
)

const (
	QUEUE_SIZE = 10
	TIMEOUT = 60
)

var client *torrent.Client

type MetadataFetcher struct {
	dbEntry model.Torrent
	torrent *torrent.Torrent
	done bool
}

type Result struct {
	fetcher *MetadataFetcher
	err error
	info *metainfo.Info
}

func NewFetcher(dbEntry model.Torrent) (f *MetadataFetcher) {
	f = &MetadataFetcher{dbEntry, nil, false}
	return
}

func (f *MetadataFetcher) Fetch(out chan Result) {
	go func(fetcher *MetadataFetcher) {
		magnet := util.InfoHashToMagnet(strings.TrimSpace(fetcher.dbEntry.Hash), fetcher.dbEntry.Name, config.Trackers...)
		T, err := client.AddMagnet(magnet)
		if err != nil {
			fetcher.done = true
			out <- Result{fetcher, err, nil}
			return
		}

		fetcher.torrent = T
		<-T.GotInfo()
		fetcher.done = true
		T.Drop()
		out <- Result{fetcher, nil, T.Info()}
	}(f)

	go func(fetcher *MetadataFetcher) {
		time.Sleep(TIMEOUT * time.Second)
		if !fetcher.done {
			if fetcher.torrent != nil {
				fetcher.torrent.Drop()
			}

			out <- Result{fetcher, errors.New("Timeout"), nil}
		}
	}(f)
}

var queue []*MetadataFetcher
var failed map[uint]struct{} // Store the torrent ID of the torrents we already fetched data

func RemoveFromQueue(fetcher *MetadataFetcher) (removed bool){
	removed = false
	for i, f := range queue {
		if fetcher == f {
			removed = true
			queue = append(queue[:i], queue[i+1:]...)
			break;
		}
	}

	return
}

func IsBeingFetchedOrHasAlreadyFailed(t model.Torrent) bool {
	for _, fetcher := range queue {
		if t.ID == fetcher.dbEntry.ID {
			return true
		}
	}

	_, ok := failed[t.ID];
	return ok
}

func FillQueue(startFetcher bool, out chan Result) (noMoreTorrents bool) {
	noMoreTorrents = false
	remainingToFill := QUEUE_SIZE - len(queue)
	if remainingToFill <= 0 {
		return
	}

	params := serviceBase.CreateWhereParams("filesize IS NULL")
	torrents, count, err := torrentService.GetTorrents(params, QUEUE_SIZE, len(failed))
	if err != nil {
		panic(fmt.Sprintf("Failed to fetch torrents, %v", err))
	}

	if count == 0 {
		noMoreTorrents = true
		return
	}

	noMoreTorrents = true
	for _, T := range torrents {
		if IsBeingFetchedOrHasAlreadyFailed(T) {
			continue;
		}

		noMoreTorrents = false
		fmt.Printf("Adding %s (ID: %d) to queue...\n", T.Name, T.ID)
		fetcher := NewFetcher(T)
		queue = append(queue, fetcher)
		if startFetcher {
			fetcher.Fetch(out)
		}
	}

	return
}

func main() {
	failed = make(map[uint]struct{})
	var r Result

	// Initialize nyaa db
	nyaaConfig := config.New()
	processFlags := nyaaConfig.BindFlags()
	flag.Parse()
	err := processFlags()
	if err != nil {
		panic(err.Error())
	}

	db.ORM, err = db.GormInit(nyaaConfig)
	if err != nil {
		panic(err.Error())
	}
	db.ORM.LogMode(false) // Disable message spam

	client, err = torrent.NewClient(nil)
	if err != nil {
		panic(fmt.Sprintf("Failed to start torrent client, %v", err))
	}

	results := make(chan Result)
	noMoreTorrents := FillQueue(true, results)
	for !noMoreTorrents {
		r = <-results
		dbEntry := r.fetcher.dbEntry
		info := r.info
		updatedSuccessfully := false
		if r.err != nil {
			fmt.Printf("[WARNING] Failed to get metadata for %s (ID: %d), error is %s\n", dbEntry.Name, dbEntry.ID, r.err.Error())
		} else if info.Length == 0 {
			fmt.Printf("[WARNING] %s (ID: %d) - Got no error, but length is 0. Bug?\n", dbEntry.Name, dbEntry.ID)
		} else {
			dbEntry.Filesize = info.Length
			_, err = torrentService.UpdateTorrent(dbEntry)
			if err == nil {
				updatedSuccessfully = true
				fmt.Printf("Torrent %s (ID: %d) updated successfully! Length is %d bytes (%s)\n", dbEntry.Name, dbEntry.ID, info.Length, util.FormatFilesize2(info.Length))
			} else {
				fmt.Printf("[WARNING] Torrent %s (ID: %d) couldn't be updated! Length is %d bytes (%s)\n", dbEntry.Name, dbEntry.ID, info.Length, util.FormatFilesize2(info.Length))
			}
		}

		if !updatedSuccessfully {
			failed[dbEntry.ID] = struct{}{}
		}

		RemoveFromQueue(r.fetcher)
		noMoreTorrents = FillQueue(true, results)
	}

	fmt.Printf("Finished!\n")
}