diff --git a/pow/pow.go b/pow/pow.go index 0b823fd..7f372f1 100644 --- a/pow/pow.go +++ b/pow/pow.go @@ -6,6 +6,9 @@ import ( "github.com/tsosunchia/powclient" "net/url" "os" + "os/signal" + "syscall" + "time" ) const ( @@ -13,6 +16,9 @@ const ( ) func GetToken(fastIp string, host string, port string) (string, error) { + // 捕获中断信号 + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) getTokenParams := powclient.NewGetTokenParams() u := url.URL{Scheme: "https", Host: fastIp + ":" + port, Path: baseURL} getTokenParams.BaseUrl = u.String() @@ -26,20 +32,30 @@ func GetToken(fastIp string, host string, port string) (string, error) { var ( token string err error + done = make(chan bool) ) - // 尝试三次RetToken,如果都失败了,异常退出 - for i := 0; i < 3; i++ { - token, err = powclient.RetToken(getTokenParams) - if err != nil { - continue + // 在 goroutine 中处理阻塞调用 + go func() { + for i := 0; i < 3; i++ { + token, err = powclient.RetToken(getTokenParams) + if err != nil { + continue // 如果失败则重试 + } + done <- true // 成功后通知主线程 + return } - //fmt.Println("GetToken success", token, getTokenParams.UserAgent) - return token, nil + done <- false // 失败后通知主线程 + }() + + select { + case <-sigChan: // 监听中断信号 + return "", fmt.Errorf("Program interrupted by user ") // 添加返回值 + case success := <-done: // 等待 goroutine 完成 + if success { + return token, nil + } + return "", fmt.Errorf("RetToken failed 3 times, please try again later") + case <-time.After(10 * time.Second): // 超时处理 + return "", fmt.Errorf("RetToken timed out(10s), please check your network") // 添加返回值 } - if err != nil { - fmt.Println(err) - } - fmt.Println("RetToken failed 3 times, please try again after a while, exit") - os.Exit(1) - return "", err } diff --git a/util/latency.go b/util/latency.go index e6c4145..51aae07 100644 --- a/util/latency.go +++ b/util/latency.go @@ -7,6 +7,8 @@ import ( "log" "net" "net/http" + "os" + "os/signal" "strings" "time" @@ -57,12 +59,17 @@ func GetFastIP(domain string, port string, enableOutput bool) string { var result ResponseInfo + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt) + select { case result = <-results: - //等待5s没有结果 视为连不上API了 + // 正常返回结果 case <-time.After(timeout): - log.Println("IP connection has been timeout, please check your network") - + log.Println("IP connection has been timeout(5s), please check your network") + case <-sigChan: // 响应中断信号 + log.Println("Program interrupted by user") + os.Exit(0) } //if len(ips) > 0 {