From ced99b966716add1cbe40f7a139de296d2a07551 Mon Sep 17 00:00:00 2001 From: Ben Sarmiento Date: Wed, 29 Nov 2023 23:05:08 +0100 Subject: [PATCH] Refactor structure, use cobra --- cmd/zurg/main.go | 86 ++++++++++--------------------- go.mod | 5 +- go.sum | 9 ++++ internal/app.go | 69 +++++++++++++++++++++++++ internal/network.go | 7 +++ internal/torrent/manager.go | 4 -- internal/torrent/util.go | 13 ----- internal/{version => }/version.go | 4 +- pkg/realdebrid/network.go | 45 +++++++++------- pkg/utils/ensure.go | 15 ++++++ 10 files changed, 158 insertions(+), 99 deletions(-) create mode 100644 internal/app.go create mode 100644 internal/network.go rename internal/{version => }/version.go (85%) create mode 100644 pkg/utils/ensure.go diff --git a/cmd/zurg/main.go b/cmd/zurg/main.go index 3587742..1279011 100644 --- a/cmd/zurg/main.go +++ b/cmd/zurg/main.go @@ -2,79 +2,45 @@ package main import ( "fmt" - "net/http" "os" - "github.com/debridmediamanager.com/zurg/internal/config" - "github.com/debridmediamanager.com/zurg/internal/net" - "github.com/debridmediamanager.com/zurg/internal/torrent" - "github.com/debridmediamanager.com/zurg/internal/universal" - "github.com/debridmediamanager.com/zurg/internal/version" - "github.com/dgraph-io/ristretto" - "github.com/panjf2000/ants/v2" - - zurghttp "github.com/debridmediamanager.com/zurg/pkg/http" - "github.com/debridmediamanager.com/zurg/pkg/logutil" - "github.com/debridmediamanager.com/zurg/pkg/realdebrid" + "github.com/debridmediamanager.com/zurg/internal" + "github.com/spf13/cobra" ) func main() { - // special commands - if len(os.Args) > 1 { - switch os.Args[1] { - case "version": - version.Show() - case "networktest": - realdebrid.RunTest() - } - os.Exit(0) + var configPath string // Variable to hold the config path + + var rootCmd = &cobra.Command{ + Use: "zurg", + Run: func(cmd *cobra.Command, args []string) { + internal.MainApp(configPath) // Pass the configPath to MainApp + }, } - // normal startup - log := logutil.NewLogger() - zurglog := log.Named("zurg") + // Adding a flag for the config path with a default value + rootCmd.PersistentFlags().StringVarP(&configPath, "config", "c", "./config.yml", "config file path") - config, configErr := config.LoadZurgConfig("./config.yml", log.Named("config")) - if configErr != nil { - zurglog.Errorf("Config failed to load: %v", configErr) - os.Exit(1) + var versionCmd = &cobra.Command{ + Use: "version", + Short: "Prints zurg's current version", + Run: func(cmd *cobra.Command, args []string) { + internal.ShowVersion() + }, } - apiClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), config.GetRealDebridTimeout(), config, log.Named("httpclient")) - - rd := realdebrid.NewRealDebrid(apiClient, log.Named("realdebrid")) - - p, err := ants.NewPool(config.GetNumOfWorkers()) - if err != nil { - zurglog.Errorf("Failed to create worker pool: %v", err) - os.Exit(1) - } - defer p.Release() - - cache, err := ristretto.NewCache(&ristretto.Config{ - NumCounters: 1e5, // 200,000 to track frequency for 100,000 items. - MaxCost: 100 << 20, // maximum cost of cache (100MB). - BufferItems: 1 << 10, // number of keys per Get buffer, can be adjusted. - }) - if err != nil { - zurglog.Errorf("Failed to create cache: %v", err) - os.Exit(1) + var networkTestCmd = &cobra.Command{ + Use: "network-test", + Short: "Run a network test", + Run: func(cmd *cobra.Command, args []string) { + internal.NetworkTest() + }, } - torrentMgr := torrent.NewTorrentManager(config, rd, p, cache, log.Named("manager")) + rootCmd.AddCommand(versionCmd, networkTestCmd) - downloadClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), 0, config, log.Named("dlclient")) - getfile := universal.NewGetFile(downloadClient) - - mux := http.NewServeMux() - net.Router(mux, getfile, config, torrentMgr, log.Named("net")) - - addr := fmt.Sprintf("%s:%s", config.GetHost(), config.GetPort()) - server := &http.Server{Addr: addr, Handler: mux} - - zurglog.Infof("Starting server on %s", addr) - if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { - zurglog.Errorf("Failed to start server: %v", err) + if err := rootCmd.Execute(); err != nil { + fmt.Println(err) os.Exit(1) } } diff --git a/go.mod b/go.mod index d31eb87..870f141 100644 --- a/go.mod +++ b/go.mod @@ -13,10 +13,13 @@ require github.com/panjf2000/ants/v2 v2.8.2 require ( github.com/cespare/xxhash/v2 v2.1.1 // indirect - github.com/dgraph-io/ristretto v0.1.1 // indirect + github.com/dgraph-io/ristretto v0.1.1 github.com/dustin/go-humanize v1.0.0 // indirect github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b // indirect + github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/spf13/cobra v1.8.0 + github.com/spf13/pflag v1.0.5 // indirect go.uber.org/multierr v1.10.0 // indirect golang.org/x/sys v0.0.0-20221010170243-090e33056c14 // indirect ) diff --git a/go.sum b/go.sum index bab7975..a80ce4b 100644 --- a/go.sum +++ b/go.sum @@ -1,15 +1,19 @@ github.com/cespare/xxhash/v2 v2.1.1 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgraph-io/ristretto v0.1.1 h1:6CWw5tJNgpegArSHpNHJKldNeq03FQCwYvfMVWajOK8= github.com/dgraph-io/ristretto v0.1.1/go.mod h1:S1GPSBCYCIhmVNfcth17y2zZtQT6wzkzgwUve0VDWWA= +github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczCTSixgIKmwPv6+wP5DGjqLYw5SUiA= github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= +github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= +github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/orcaman/concurrent-map/v2 v2.0.1 h1:jOJ5Pg2w1oeB6PeDurIYf6k9PQ+aTITr/6lP/L/zp6c= github.com/orcaman/concurrent-map/v2 v2.0.1/go.mod h1:9Eq3TG2oBe5FirmYWQfYO5iH1q0Jv47PLaNK++uCdOM= github.com/panjf2000/ants/v2 v2.8.2 h1:D1wfANttg8uXhC9149gRt1PDQ+dLVFjNXkCEycMcvQQ= @@ -18,6 +22,11 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0= +github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho= +github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= +github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= diff --git a/internal/app.go b/internal/app.go new file mode 100644 index 0000000..e2d0790 --- /dev/null +++ b/internal/app.go @@ -0,0 +1,69 @@ +package internal + +import ( + "fmt" + "net/http" + "os" + + "github.com/debridmediamanager.com/zurg/internal/config" + "github.com/debridmediamanager.com/zurg/internal/net" + "github.com/debridmediamanager.com/zurg/internal/torrent" + "github.com/debridmediamanager.com/zurg/internal/universal" + zurghttp "github.com/debridmediamanager.com/zurg/pkg/http" + "github.com/debridmediamanager.com/zurg/pkg/logutil" + "github.com/debridmediamanager.com/zurg/pkg/realdebrid" + "github.com/debridmediamanager.com/zurg/pkg/utils" + "github.com/dgraph-io/ristretto" + "github.com/panjf2000/ants/v2" +) + +func MainApp(configPath string) { + log := logutil.NewLogger() + zurglog := log.Named("zurg") + + config, configErr := config.LoadZurgConfig(configPath, log.Named("config")) + if configErr != nil { + zurglog.Errorf("Config failed to load: %v", configErr) + os.Exit(1) + } + + apiClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), config.GetRealDebridTimeout(), config, log.Named("httpclient")) + + rd := realdebrid.NewRealDebrid(apiClient, log.Named("realdebrid")) + + p, err := ants.NewPool(config.GetNumOfWorkers()) + if err != nil { + zurglog.Errorf("Failed to create worker pool: %v", err) + os.Exit(1) + } + defer p.Release() + + utils.EnsureDirExists("data") + + cache, err := ristretto.NewCache(&ristretto.Config{ + NumCounters: 1e5, // 200,000 to track frequency for 100,000 items. + MaxCost: 100 << 20, // maximum cost of cache (100MB). + BufferItems: 1 << 10, // number of keys per Get buffer, can be adjusted. + }) + if err != nil { + zurglog.Errorf("Failed to create cache: %v", err) + os.Exit(1) + } + + torrentMgr := torrent.NewTorrentManager(config, rd, p, cache, log.Named("manager")) + + downloadClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), 0, config, log.Named("dlclient")) + getfile := universal.NewGetFile(downloadClient) + + mux := http.NewServeMux() + net.Router(mux, getfile, config, torrentMgr, log.Named("net")) + + addr := fmt.Sprintf("%s:%s", config.GetHost(), config.GetPort()) + server := &http.Server{Addr: addr, Handler: mux} + + zurglog.Infof("Starting server on %s", addr) + if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { + zurglog.Errorf("Failed to start server: %v", err) + os.Exit(1) + } +} diff --git a/internal/network.go b/internal/network.go new file mode 100644 index 0000000..d3c3e00 --- /dev/null +++ b/internal/network.go @@ -0,0 +1,7 @@ +package internal + +import "github.com/debridmediamanager.com/zurg/pkg/realdebrid" + +func NetworkTest() { + realdebrid.RunTest() +} diff --git a/internal/torrent/manager.go b/internal/torrent/manager.go index ff7c14f..9675642 100644 --- a/internal/torrent/manager.go +++ b/internal/torrent/manager.go @@ -40,7 +40,6 @@ type TorrentManager struct { antsPool *ants.Pool unrestrictPool *ants.Pool log *zap.SugaredLogger - mu *sync.Mutex } // NewTorrentManager creates a new torrent manager @@ -55,7 +54,6 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p requiredVersion: "18.11.2023", antsPool: p, log: log, - mu: &sync.Mutex{}, } unrestrictPool, err := ants.NewPool(t.Config.GetUnrestrictWorkers()) @@ -65,8 +63,6 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p t.unrestrictPool = unrestrictPool } - ensureDir(DATA_DIR) - // create internal directories t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]()) // key is AccessKey t.DirectoryMap.Set(INT_INFO_CACHE, cmap.New[*Torrent]()) // key is Torrent ID diff --git a/internal/torrent/util.go b/internal/torrent/util.go index 0bcd2d4..f186680 100644 --- a/internal/torrent/util.go +++ b/internal/torrent/util.go @@ -2,7 +2,6 @@ package torrent import ( "fmt" - "os" ) func getFileIDs(files []File) []string { @@ -14,15 +13,3 @@ func getFileIDs(files []File) []string { } return fileIDs } - -func ensureDir(dirName string) error { - if _, err := os.Stat(dirName); os.IsNotExist(err) { - err := os.Mkdir(dirName, 0755) - if err != nil { - return err - } - } else if err != nil { - return err - } - return nil -} diff --git a/internal/version/version.go b/internal/version.go similarity index 85% rename from internal/version/version.go rename to internal/version.go index 378e419..1466b48 100644 --- a/internal/version/version.go +++ b/internal/version.go @@ -1,4 +1,4 @@ -package version +package internal import ( "fmt" @@ -11,7 +11,7 @@ var ( Version string = "dev" ) -func Show() { +func ShowVersion() { fmt.Printf("zurg\nBuilt At: %s\nGo Version: %s\nCommit: %s\nVersion: %s\n", BuiltAt, GoVersion, GitCommit, Version) } diff --git a/pkg/realdebrid/network.go b/pkg/realdebrid/network.go index c33e62d..1ae8a37 100644 --- a/pkg/realdebrid/network.go +++ b/pkg/realdebrid/network.go @@ -31,7 +31,7 @@ func traceroute(ip string) (int, time.Duration, error) { var latency time.Duration if hopCount > 0 { lastLine := lines[hopCount-1] - fmt.Println(lastLine) + fmt.Printf("IP: %s, last hop: %s\n", ip, lastLine) parts := strings.Fields(lastLine) if len(parts) >= 3 { latencyValue, parseErr := strconv.ParseFloat(parts[2], 64) @@ -70,6 +70,7 @@ func RunTest() { wg.Wait() close(infoChan) + fmt.Printf("Network test complete.\n\n") var ipInfos []IPInfo for info := range infoChan { @@ -77,30 +78,36 @@ func RunTest() { } sort.Slice(ipInfos, func(i, j int) bool { - if ipInfos[i].Hops == ipInfos[j].Hops { - return ipInfos[i].Latency < ipInfos[j].Latency - } - return ipInfos[i].Hops < ipInfos[j].Hops + return ipInfos[i].Latency < ipInfos[j].Latency }) - var lowestLatency time.Duration + + const minResults = 10 + const maxResults = 20 + var okIPs []IPInfo + if len(ipInfos) > 0 { - lowestLatency = ipInfos[0].Latency - for _, info := range ipInfos { - if info.Latency < lowestLatency { - lowestLatency = info.Latency + // Start by adding the best IPs based on hops and latency up to the minResults + for i := 0; i < min(len(ipInfos), minResults); i++ { + okIPs = append(okIPs, ipInfos[i]) + } + + // Find the highest latency in the current okIPs list + highestLatency := okIPs[len(okIPs)-1].Latency + + // Add any additional IPs that have latency within a reasonable range of the highest latency + for _, info := range ipInfos[minResults:] { + if len(okIPs) >= maxResults { + break // Stop adding IPs if maxResults is reached + } + if info.Latency <= highestLatency+(highestLatency/3) { + okIPs = append(okIPs, info) } } } - latencyThreshold := lowestLatency + lowestLatency/3 - - var okIPs []IPInfo - for _, info := range ipInfos { - if info.Latency <= latencyThreshold { - okIPs = append(okIPs, info) - } - } + fmt.Printf("Here are the results, you can copy-paste the following to your config.yml:\n\n") + fmt.Println("preferred_hosts:") for _, info := range okIPs { - fmt.Printf("Host: %s, Hops: %d, Latency: %v\n", info.Address, info.Hops, info.Latency) + fmt.Printf(" - %s # hops: %d latency: %v\n", info.Address, info.Hops, info.Latency) } } diff --git a/pkg/utils/ensure.go b/pkg/utils/ensure.go new file mode 100644 index 0000000..528e7ef --- /dev/null +++ b/pkg/utils/ensure.go @@ -0,0 +1,15 @@ +package utils + +import "os" + +func EnsureDirExists(dirName string) error { + if _, err := os.Stat(dirName); os.IsNotExist(err) { + err := os.Mkdir(dirName, 0755) + if err != nil { + return err + } + } else if err != nil { + return err + } + return nil +}