From 48b237827e841d32a7911f14d2bb6e5d98827c1e Mon Sep 17 00:00:00 2001 From: r3inbowari Date: Fri, 17 Mar 2023 19:16:42 +0800 Subject: [PATCH] feat: new ui and fix bugs (#118) --- example/main.go | 6 +- go.mod | 9 +- go.sum | 24 ++- speedtest.go | 354 ++++++++++++-------------------------- speedtest/data_manager.go | 18 +- speedtest/output.go | 30 ++++ speedtest/request.go | 8 +- speedtest/server.go | 58 ++++--- speedtest/server_test.go | 8 +- speedtest/speedtest.go | 14 +- speedtest/user.go | 7 +- task.go | 99 +++++++++++ 12 files changed, 323 insertions(+), 312 deletions(-) create mode 100644 speedtest/output.go create mode 100644 task.go diff --git a/example/main.go b/example/main.go index 7f8c7fb..6f50508 100644 --- a/example/main.go +++ b/example/main.go @@ -7,7 +7,7 @@ import ( ) func main() { - user, _ := speedtest.FetchUserInfo() + _, _ = speedtest.FetchUserInfo() // Get a list of servers near a specified location // user.SetLocationByCity("Tokyo") // user.SetLocation("Osaka", 34.6952, 135.5006) @@ -15,14 +15,14 @@ func main() { // Select a network card as the data interface. // speedtest.WithUserConfig(&speedtest.UserConfig{Source: "192.168.1.101"})(speedtestClient) - serverList, _ := speedtest.FetchServers(user) + serverList, _ := speedtest.FetchServers() targets, _ := serverList.FindServer([]int{}) for _, s := range targets { // Please make sure your host can access this test server, // otherwise you will get an error. // It is recommended to replace a server at this time - checkError(s.PingTest()) + checkError(s.PingTest(nil)) checkError(s.DownloadTest()) checkError(s.UploadTest()) diff --git a/go.mod b/go.mod index 3d5100d..872d3b4 100644 --- a/go.mod +++ b/go.mod @@ -2,9 +2,16 @@ module github.com/showwin/speedtest-go go 1.18 -require gopkg.in/alecthomas/kingpin.v2 v2.2.6 +require ( + github.com/chelnak/ysmrr v0.2.1 + gopkg.in/alecthomas/kingpin.v2 v2.2.6 +) require ( github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect github.com/alecthomas/units v0.0.0-20210208195552-ff826a37aa15 // indirect + github.com/fatih/color v1.13.0 // indirect + github.com/mattn/go-colorable v0.1.13 // indirect + github.com/mattn/go-isatty v0.0.16 // indirect + golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab // indirect ) diff --git a/go.sum b/go.sum index 8b434e6..59cfbf4 100644 --- a/go.sum +++ b/go.sum @@ -1,18 +1,32 @@ -github.com/LyricTian/queue v1.3.0 h1:1xEFZlteW6iu5Qbrz7ZsiSKMKaxY1bQHsbx0jrB1pDA= -github.com/LyricTian/queue v1.3.0/go.mod h1:pbkoplz/zRToCay3pRjz75P8fQAgvkRKJdEzVUQYhXY= github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafoB+tBA3gMyHYHrpOtNuDiK/uB5uXxq5wM= github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc= github.com/alecthomas/units v0.0.0-20210208195552-ff826a37aa15 h1:AUNCr9CiJuwrRYS3XieqF+Z9B9gNxo/eANAJCF2eiN4= github.com/alecthomas/units v0.0.0-20210208195552-ff826a37aa15/go.mod h1:OMCwj8VM1Kc9e19TLln2VL61YJF0x1XFtfdL4JdbSyE= -github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/chelnak/ysmrr v0.2.1 h1:9xLbVcrgnvEFovFAPnDiTCtxHiuLmz03xCg5OUgdOfc= +github.com/chelnak/ysmrr v0.2.1/go.mod h1:9TEgLy2xDMGN62zJm9XZrEWY/fHoGoBslSVEkEpRCXk= 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/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w= +github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk= +github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU= +github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= +github.com/mattn/go-isatty v0.0.16 h1:bq3VjFmv/sOjHtdEhmkEV4x1AJtvUvOJ2PFAZ5+peKQ= +github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= 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/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= +github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= +golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab h1:2QkjZIsXupsJbJIdSjjUOgWK3aEtzyuh2mPt3l/CkeU= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= gopkg.in/alecthomas/kingpin.v2 v2.2.6 h1:jMFz6MfLP0/4fUyZle81rXUoxOBFi19VUFKVDOQfozc= gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= -gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/speedtest.go b/speedtest.go index 1e53fcd..7806fd7 100644 --- a/speedtest.go +++ b/speedtest.go @@ -2,11 +2,8 @@ package main import ( "context" - "encoding/json" "fmt" - "log" "os" - "runtime" "time" "github.com/showwin/speedtest-go/speedtest" @@ -32,224 +29,125 @@ var ( debug = kingpin.Flag("debug", "Enable debug mode.").Short('d').Bool() ) -type fullOutput struct { - Timestamp outputTime `json:"timestamp"` - UserInfo *speedtest.User `json:"user_info"` - Servers speedtest.Servers `json:"servers"` -} - -type outputTime time.Time - func main() { + kingpin.Version(speedtest.Version()) kingpin.Parse() - - var speedtestClient = speedtest.New() + AppInfo() + + // 0. speed test setting + var speedtestClient = speedtest.New(speedtest.WithUserConfig( + &speedtest.UserConfig{ + UserAgent: speedtest.DefaultUserAgent, + Proxy: *proxy, + Source: *source, + Debug: *debug, + ICMP: (os.Geteuid() == 0 || os.Geteuid() == -1) && len(*proxy) == 0, // proxy may not support ICMP + SavingMode: *savingMode, + CityFlag: *city, + LocationFlag: *location, + Keyword: *search, + NoDownload: *noDownload, + NoUpload: *noUpload, + })) speedtestClient.SetNThread(*thread) - config := &speedtest.UserConfig{ - UserAgent: speedtest.DefaultUserAgent, - Proxy: *proxy, - Source: *source, - Debug: *debug, - ICMP: os.Geteuid() == 0 || runtime.GOOS == "windows", - SavingMode: *savingMode, - CityFlag: *city, - LocationFlag: *location, - Keyword: *search, - NoDownload: *noDownload, - NoUpload: *noUpload, - } - speedtest.WithUserConfig(config)(speedtestClient) if *showCityList { speedtest.PrintCityList() return } - user, err := speedtestClient.FetchUserInfo() - if err != nil { - fmt.Printf("Fatal: can not fetch user information. err: %v\n", err.Error()) - return - } + // 1. retrieving user information + taskManager := InitTaskManager(!*jsonOutput) + taskManager.AsyncRun("Retrieving User Information", func(task *Task) { + u, err := speedtestClient.FetchUserInfo() + task.CheckError(err) + task.Printf("ISP: %s", u.String()) + task.Complete() + }) - if !*jsonOutput { - showUser(user) - } - - servers, err := speedtestClient.FetchServers(user) - checkError(err) + // 2. retrieving servers + var err error + var servers speedtest.Servers var targets speedtest.Servers - if *customURL == "" { - if *showList { - showServerList(servers) - return - } - - targets, err = servers.FindServer(*serverIds) - checkError(err) - - } else { - target, err := speedtestClient.CustomServer(*customURL) - checkError(err) - targets = []*speedtest.Server{target} - } - - if *multi { - startMultiTest(targets[0], servers, *savingMode, *jsonOutput) - } else { - startTest(targets, *savingMode, *jsonOutput) - } - - if *jsonOutput { - jsonBytes, err := json.Marshal( - fullOutput{ - Timestamp: outputTime(time.Now()), - UserInfo: user, - Servers: targets, - }, - ) - checkError(err) - - fmt.Println(string(jsonBytes)) - } -} - -func startMultiTest(s *speedtest.Server, servers speedtest.Servers, savingMode bool, jsonOutput bool) { - // Reset DataManager counters, avoid measurement of multiple server result mixing. - s.Context.Reset() - if !jsonOutput { - showServer(s) - } - - err := s.PingTest() - checkError(err) - - if jsonOutput { - err = s.MultiDownloadTestContext(context.Background(), servers) - checkError(err) - s.Context.Wait() - err = s.MultiUploadTestContext(context.Background(), servers) - checkError(err) - return - } - - showLatencyResult(s) - err = testDownloadM(s, servers) - checkError(err) - // It is necessary to wait for the release of the last test resource, - // otherwise the overload will cause excessive data deviation - s.Context.Wait() - err = testUploadM(s, servers) - checkError(err) - showServerResult(s) -} - -func startTest(servers speedtest.Servers, savingMode bool, jsonOutput bool) { - for _, s := range servers { - // Reset DataManager counters, avoid measurement of multiple server result mixing. - s.Context.Reset() - if !jsonOutput { - showServer(s) + taskManager.Run("Retrieving Servers", func(task *Task) { + if len(*customURL) > 0 { + var target *speedtest.Server + target, err = speedtestClient.CustomServer(*customURL) + task.CheckError(err) + targets = []*speedtest.Server{target} + task.Println("Skip: Using Custom Server") + } else { + servers, err = speedtestClient.FetchServers() + task.CheckError(err) + task.Printf("Found %d Public Servers", len(servers)) + if *showList { + task.Complete() + task.manager.Reset() + showServerList(servers) + os.Exit(1) + } + targets, err = servers.FindServer(*serverIds) + task.CheckError(err) } - - err := s.PingTest() - checkError(err) - - if jsonOutput { - err = s.DownloadTest() - checkError(err) - s.Context.Wait() - err = s.UploadTest() - checkError(err) - continue + task.Complete() + }) + taskManager.Reset() + + // 3. test each selected server with ping, download and upload. + for _, server := range targets { + if !*jsonOutput { + fmt.Println() } + taskManager.Println("Test Server: " + server.String()) + taskManager.Run("Latency: ", func(task *Task) { + server.PingTest(func(latency time.Duration) { + task.Printf("Latency: %v", latency) + }) + task.Printf("Latency: %v Jitter: %v Min: %v Max: %v", server.Latency, server.Jitter, server.MinLatency, server.MaxLatency) + task.Complete() + }) + + taskManager.Run("Download", func(task *Task) { + ticker := speedtestClient.CallbackDownloadRate(func(downRate float64) { + task.Printf("Download: %.2fMbps", downRate) + }) + if *multi { + task.CheckError(server.MultiDownloadTestContext(context.Background(), servers)) + } else { + task.CheckError(server.DownloadTest()) + } + ticker.Stop() + task.Printf("Download: %.2fMbps (used: %.2fMB)", server.DLSpeed, float64(server.Context.Manager.GetTotalDownload())/1024/1024) + task.Complete() + }) + + taskManager.Run("Upload", func(task *Task) { + ticker := speedtestClient.CallbackUploadRate(func(upRate float64) { + task.Printf("Upload: %.2fMbps", upRate) + }) + if *multi { + task.CheckError(server.MultiUploadTestContext(context.Background(), servers)) + } else { + task.CheckError(server.UploadTest()) + } + ticker.Stop() + task.Printf("Upload: %.2fMbps (used: %.2fMB)", server.ULSpeed, float64(server.Context.Manager.GetTotalUpload())/1024/1024) + task.Complete() + }) + taskManager.Reset() + speedtestClient.Manager.Reset() + } + + taskManager.Stop() - showLatencyResult(s) - - err = testDownload(s) - checkError(err) - // It is necessary to wait for the release of the last test resource, - // otherwise the overload will cause excessive data deviation - s.Context.Wait() - err = testUpload(s) - checkError(err) - showServerResult(s) - } - - if !jsonOutput && len(servers) > 1 { - showAverageServerResult(servers) - } -} - -func testDownloadM(server *speedtest.Server, servers speedtest.Servers) error { - quit := make(chan bool) - fmt.Printf("Download Test: ") - go dots(quit) - err := server.MultiDownloadTestContext(context.Background(), servers) - checkError(err) - quit <- true - if err != nil { - return err - } - fmt.Println() - return err -} - -func testUploadM(server *speedtest.Server, servers speedtest.Servers) error { - quit := make(chan bool) - fmt.Printf("Upload Test: ") - go dots(quit) - err := server.MultiUploadTestContext(context.Background(), servers) - checkError(err) - quit <- true - if err != nil { - return err - } - fmt.Println() - return nil -} - -func testDownload(server *speedtest.Server) error { - quit := make(chan bool) - fmt.Printf("Download Test: ") - go dots(quit) - err := server.DownloadTest() - quit <- true - if err != nil { - return err - } - fmt.Println() - return err -} - -func testUpload(server *speedtest.Server) error { - quit := make(chan bool) - fmt.Printf("Upload Test: ") - go dots(quit) - err := server.UploadTest() - quit <- true - if err != nil { - return err - } - fmt.Println() - return nil -} - -func dots(quit chan bool) { - for { - select { - case <-quit: + if *jsonOutput { + json, errMarshal := speedtestClient.JSON(targets) + if errMarshal != nil { + fmt.Printf(errMarshal.Error()) return - default: - time.Sleep(time.Second) - fmt.Print(".") } - } -} - -func showUser(user *speedtest.User) { - if user.IP != "" { - fmt.Printf("Testing From IP: %s\n", user.String()) + fmt.Printf(string(json)) } } @@ -266,50 +164,10 @@ func showServerList(servers speedtest.Servers) { } } -func showServer(s *speedtest.Server) { - fmt.Println() - fmt.Printf("Target Server: [%4s] %8.2fkm %s", s.ID, s.Distance, s.Name) - if s.ID == "Custom" { +func AppInfo() { + if !*jsonOutput { + fmt.Println() + fmt.Printf(" 😊 speedtest-go v%s @showwin 😊\n", speedtest.Version()) fmt.Println() - return - } - fmt.Printf(" (" + s.Country + ") by " + s.Sponsor + "\n") -} - -func showLatencyResult(server *speedtest.Server) { - fmt.Printf("Latency: %v\nJitter: %v\nMin: %v\nMax: %v\n", server.Latency, server.Jitter, server.MinLatency, server.MaxLatency) -} - -// ShowResult : show testing result -func showServerResult(server *speedtest.Server) { - fmt.Printf(" \n") - - fmt.Printf("Download: %5.2f Mbit/s\n", server.DLSpeed) - fmt.Printf("Upload: %5.2f Mbit/s\n\n", server.ULSpeed) - valid := server.CheckResultValid() - if !valid { - fmt.Println("Warning: Result seems to be wrong. Please speedtest again.") - } -} - -func showAverageServerResult(servers speedtest.Servers) { - avgDL := 0.0 - avgUL := 0.0 - for _, s := range servers { - avgDL = avgDL + s.DLSpeed - avgUL = avgUL + s.ULSpeed - } - fmt.Printf("Download Avg: %5.2f Mbit/s\n", avgDL/float64(len(servers))) - fmt.Printf("Upload Avg: %5.2f Mbit/s\n", avgUL/float64(len(servers))) -} - -func checkError(err error) { - if err != nil { - log.Fatal(err) } } - -func (t outputTime) MarshalJSON() ([]byte, error) { - stamp := fmt.Sprintf("\"%s\"", time.Time(t).Format("2006-01-02 15:04:05.000")) - return []byte(stamp), nil -} diff --git a/speedtest/data_manager.go b/speedtest/data_manager.go index 997a1af..ce7c528 100644 --- a/speedtest/data_manager.go +++ b/speedtest/data_manager.go @@ -102,15 +102,10 @@ func NewDataManager() *DataManager { func (dm *DataManager) CallbackDownloadRate(callback func(downRate float64)) *time.Ticker { ticker := time.NewTicker(dm.rateCaptureFrequency) - oldDownTotal := dm.GetTotalDownload() - unit := float64(time.Second / dm.rateCaptureFrequency) - go func() { + sTime := time.Now() for range ticker.C { - newDownTotal := dm.GetTotalDownload() - delta := newDownTotal - oldDownTotal - oldDownTotal = newDownTotal - callback(float64(delta) * 8 / 1000000 * unit) + callback((float64(dm.GetTotalDownload()) * 8 / 1000000) / float64(time.Now().Sub(sTime).Milliseconds()) * 1000) } }() return ticker @@ -118,15 +113,10 @@ func (dm *DataManager) CallbackDownloadRate(callback func(downRate float64)) *ti func (dm *DataManager) CallbackUploadRate(callback func(upRate float64)) *time.Ticker { ticker := time.NewTicker(dm.rateCaptureFrequency) - oldUpTotal := dm.GetTotalUpload() - unit := float64(time.Second / dm.rateCaptureFrequency) - go func() { + sTime := time.Now() for range ticker.C { - newUpTotal := dm.GetTotalUpload() - delta := newUpTotal - oldUpTotal - oldUpTotal = newUpTotal - callback(float64(delta) * 8 / 1000000 * unit) + callback((float64(dm.GetTotalUpload()) * 8 / 1000000) / float64(time.Now().Sub(sTime).Milliseconds()) * 1000) } }() return ticker diff --git a/speedtest/output.go b/speedtest/output.go new file mode 100644 index 0000000..733f9e5 --- /dev/null +++ b/speedtest/output.go @@ -0,0 +1,30 @@ +package speedtest + +import ( + "encoding/json" + "fmt" + "time" +) + +type fullOutput struct { + Timestamp outputTime `json:"timestamp"` + UserInfo *User `json:"user_info"` + Servers Servers `json:"servers"` +} + +type outputTime time.Time + +func (t outputTime) MarshalJSON() ([]byte, error) { + stamp := fmt.Sprintf("\"%s\"", time.Time(t).Format("2006-01-02 15:04:05.000")) + return []byte(stamp), nil +} + +func (s *Speedtest) JSON(servers Servers) ([]byte, error) { + return json.Marshal( + fullOutput{ + Timestamp: outputTime(time.Now()), + UserInfo: s.User, + Servers: servers, + }, + ) +} diff --git a/speedtest/request.go b/speedtest/request.go index 67cf2e2..15455d0 100644 --- a/speedtest/request.go +++ b/speedtest/request.go @@ -159,15 +159,15 @@ func uploadRequest(ctx context.Context, s *Server, w int) error { } // PingTest executes test to measure latency -func (s *Server) PingTest() error { - return s.PingTestContext(context.Background()) +func (s *Server) PingTest(callback func(latency time.Duration)) error { + return s.PingTestContext(context.Background(), callback) } // PingTestContext executes test to measure latency, observing the given context. -func (s *Server) PingTestContext(ctx context.Context) (err error) { +func (s *Server) PingTestContext(ctx context.Context, callback func(latency time.Duration)) (err error) { var vectorPingResult []int64 if s.Context.config.ICMP { - vectorPingResult, err = s.ICMPPing(ctx, time.Second*4, 10, time.Millisecond*200, nil) + vectorPingResult, err = s.ICMPPing(ctx, time.Second*4, 10, time.Millisecond*200, callback) } else { vectorPingResult, err = s.HTTPPing(ctx, 10, time.Millisecond*200, nil) } diff --git a/speedtest/server.go b/speedtest/server.go index 7540718..a7be3e1 100644 --- a/speedtest/server.go +++ b/speedtest/server.go @@ -125,17 +125,17 @@ func (b ByDistance) Less(i, j int) bool { } // FetchServers retrieves a list of available servers -func (s *Speedtest) FetchServers(user *User) (Servers, error) { - return s.FetchServerListContext(context.Background(), user) +func (s *Speedtest) FetchServers() (Servers, error) { + return s.FetchServerListContext(context.Background()) } // FetchServers retrieves a list of available servers -func FetchServers(user *User) (Servers, error) { - return defaultClient.FetchServers(user) +func FetchServers() (Servers, error) { + return defaultClient.FetchServers() } // FetchServerListContext retrieves a list of available servers, observing the given context. -func (s *Speedtest) FetchServerListContext(ctx context.Context, user *User) (Servers, error) { +func (s *Speedtest) FetchServerListContext(ctx context.Context) (Servers, error) { u, err := url.Parse(speedTestServersUrl) if err != nil { return Servers{}, err @@ -210,24 +210,9 @@ func (s *Speedtest) FetchServerListContext(ctx context.Context, user *User) (Ser server.Context = s } - // Calculate distance - for _, server := range servers { - sLat, _ := strconv.ParseFloat(server.Lat, 64) - sLon, _ := strconv.ParseFloat(server.Lon, 64) - uLat, _ := strconv.ParseFloat(user.Lat, 64) - uLon, _ := strconv.ParseFloat(user.Lon, 64) - server.Distance = distance(sLat, sLon, uLat, uLon) - } - - // Sort by distance - sort.Sort(ByDistance{servers}) - - if len(servers) <= 0 { - return servers, errors.New("unable to retrieve server list") - } - + // ping once var wg sync.WaitGroup - pCtx, fc := context.WithTimeout(context.Background(), time.Second*5) + pCtx, fc := context.WithTimeout(context.Background(), time.Second*4) dbg.Println("Echo each server...") for _, server := range servers { wg.Add(1) @@ -246,15 +231,33 @@ func (s *Speedtest) FetchServerListContext(ctx context.Context, user *User) (Ser wg.Done() }(server) } - wg.Wait() fc() + + // Calculate distance + if s.User == nil { + s.asyncFetchUser.Wait() + } + for _, server := range servers { + sLat, _ := strconv.ParseFloat(server.Lat, 64) + sLon, _ := strconv.ParseFloat(server.Lon, 64) + uLat, _ := strconv.ParseFloat(s.User.Lat, 64) + uLon, _ := strconv.ParseFloat(s.User.Lon, 64) + server.Distance = distance(sLat, sLon, uLat, uLon) + } + + // Sort by distance + sort.Sort(ByDistance{servers}) + + if len(servers) <= 0 { + return servers, errors.New("unable to retrieve server list") + } return servers, nil } // FetchServerListContext retrieves a list of available servers, observing the given context. -func FetchServerListContext(ctx context.Context, user *User) (Servers, error) { - return defaultClient.FetchServerListContext(ctx, user) +func FetchServerListContext(ctx context.Context) (Servers, error) { + return defaultClient.FetchServerListContext(ctx) } func distance(lat1 float64, lon1 float64, lat2 float64, lon2 float64) float64 { @@ -324,7 +327,10 @@ func (servers Servers) String() string { // String representation of Server func (s *Server) String() string { - return fmt.Sprintf("[%4s] %8.2fkm \n%s (%s) by %s\n", s.ID, s.Distance, s.Name, s.Country, s.Sponsor) + if s.Sponsor == "?" { + return fmt.Sprintf("[%4s] %s", s.ID, s.Name) + } + return fmt.Sprintf("[%4s] %.2fkm %s (%s) by %s", s.ID, s.Distance, s.Name, s.Country, s.Sponsor) } // CheckResultValid checks that results are logical given UL and DL speeds diff --git a/speedtest/server_test.go b/speedtest/server_test.go index fa41064..9bf099c 100644 --- a/speedtest/server_test.go +++ b/speedtest/server_test.go @@ -3,16 +3,14 @@ package speedtest import "testing" func TestFetchServerList(t *testing.T) { - user := User{ + client := New() + client.User = &User{ IP: "111.111.111.111", Lat: "35.22", Lon: "138.44", Isp: "Hello", } - - client := New() - - servers, err := client.FetchServers(&user) + servers, err := client.FetchServers() if err != nil { t.Errorf(err.Error()) } diff --git a/speedtest/speedtest.go b/speedtest/speedtest.go index 2b56d89..2a316cc 100644 --- a/speedtest/speedtest.go +++ b/speedtest/speedtest.go @@ -9,21 +9,26 @@ import ( "net/url" "runtime" "strings" + "sync" "time" ) var ( - version = "1.5.2" + version = "1.6.0" DefaultUserAgent = fmt.Sprintf("showwin/speedtest-go %s", version) ) // Speedtest is a speedtest client. type Speedtest struct { + User *User + Manager + doer *http.Client config *UserConfig tcpDialer *net.Dialer ipDialer *net.Dialer - Manager + + asyncFetchUser *sync.WaitGroup } type UserConfig struct { @@ -168,8 +173,9 @@ func WithUserConfig(userConfig *UserConfig) Option { func New(opts ...Option) *Speedtest { log.SetOutput(io.Discard) s := &Speedtest{ - doer: http.DefaultClient, - Manager: GlobalDataManager, + doer: http.DefaultClient, + asyncFetchUser: &sync.WaitGroup{}, + Manager: GlobalDataManager, } // load default config s.NewUserConfig(&UserConfig{UserAgent: DefaultUserAgent}) diff --git a/speedtest/user.go b/speedtest/user.go index 93a37a0..f01a27b 100644 --- a/speedtest/user.go +++ b/speedtest/user.go @@ -35,6 +35,8 @@ func FetchUserInfo() (*User, error) { // FetchUserInfoContext returns information about caller determined by speedtest.net, observing the given context. func (s *Speedtest) FetchUserInfoContext(ctx context.Context) (*User, error) { + s.asyncFetchUser.Add(1) + defer s.asyncFetchUser.Done() dbg.Printf("Retrieving user info: %s\n", speedTestConfigUrl) req, err := http.NewRequestWithContext(ctx, http.MethodGet, speedTestConfigUrl, nil) if err != nil { @@ -60,7 +62,8 @@ func (s *Speedtest) FetchUserInfoContext(ctx context.Context) (*User, error) { return nil, errors.New("failed to fetch user information") } - return &users.Users[0], nil + s.User = &users.Users[0] + return s.User, nil } // FetchUserInfoContext returns information about caller determined by speedtest.net, observing the given context. @@ -71,5 +74,5 @@ func FetchUserInfoContext(ctx context.Context) (*User, error) { // String representation of User func (u *User) String() string { extInfo := "" - return fmt.Sprintf("%s, (%s) [%s, %s] %s", u.IP, u.Isp, u.Lat, u.Lon, extInfo) + return fmt.Sprintf("%s (%s) [%s, %s] %s", u.IP, u.Isp, u.Lat, u.Lon, extInfo) } diff --git a/task.go b/task.go new file mode 100644 index 0000000..948382b --- /dev/null +++ b/task.go @@ -0,0 +1,99 @@ +package main + +import ( + "fmt" + "github.com/chelnak/ysmrr" + "os" + "strings" +) + +type TaskManager struct { + sm ysmrr.SpinnerManager + isOut bool +} + +type Task struct { + spinner *ysmrr.Spinner + manager *TaskManager + title string +} + +func InitTaskManager(isOut bool) *TaskManager { + tm := &TaskManager{sm: ysmrr.NewSpinnerManager(), isOut: isOut} + if isOut { + tm.sm.Start() + } + return tm +} + +func (tm *TaskManager) Reset() { + if tm.isOut { + tm.sm.Stop() + tm.sm = ysmrr.NewSpinnerManager() + tm.sm.Start() + } +} + +func (tm *TaskManager) Stop() { + if tm.isOut { + tm.sm.Stop() + } +} + +func (tm *TaskManager) Println(message string) { + if tm.isOut { + context := &Task{manager: tm} + context.spinner = tm.sm.AddSpinner(message) + context.Complete() + } +} + +func (tm *TaskManager) Run(title string, callback func(task *Task)) { + context := &Task{manager: tm, title: title} + if tm.isOut { + context.spinner = tm.sm.AddSpinner(title) + } + callback(context) +} + +func (tm *TaskManager) AsyncRun(title string, callback func(task *Task)) { + context := &Task{manager: tm, title: title} + if tm.isOut { + context.spinner = tm.sm.AddSpinner(title) + } + go callback(context) +} + +func (t *Task) Complete() { + if t.spinner == nil { + return + } + t.spinner.Complete() +} + +func (t *Task) Println(message string) { + if t.spinner == nil { + return + } + t.spinner.UpdateMessage(message) +} + +func (t *Task) Printf(format string, a ...interface{}) { + if t.spinner == nil { + return + } + t.spinner.UpdateMessagef(format, a...) +} + +func (t *Task) CheckError(err error) { + if err != nil { + if t.spinner != nil { + t.Printf("Fatal: %s, err: %v", strings.ToLower(t.title), err) + t.spinner.Error() + t.manager.Stop() + } else { + fmt.Printf("Fatal: %s, err: %v", strings.ToLower(t.title), err) + } + os.Exit(1) + } +}