diff --git a/cmd/podcli/check.go b/cmd/podcli/check.go index 7ea04e5..0d1aad8 100644 --- a/cmd/podcli/check.go +++ b/cmd/podcli/check.go @@ -6,10 +6,12 @@ import ( "fmt" "net/http" "os" + "strings" "time" "github.com/spf13/cobra" "go.uber.org/zap" + "net" ) var ( @@ -27,19 +29,32 @@ var checkCmd = &cobra.Command{ } var checkUrlCmd = &cobra.Command{ - Use: `http [--method POST] [--retry 2] [--delay 1s] [--timeout 1s] [--body '{"test"="testing"}']`, + Use: `http [URL]`, Short: "HTTP/S health check", - Example: ` check http https://httpbin.org/anything --retry=2 --delay=2s --timeout=1s --body='{"test"=1}'`, + Example: ` check http https://httpbin.org/anything --method=POST --retry=2 --delay=2s --timeout=1s --body='{"test"=1}'`, RunE: runCheck, } +var checkTcpCmd = &cobra.Command{ + Use: `tcp [address]`, + Short: "HTTP/S health check", + Example: ` check http https://httpbin.org/anything --retry=2 --delay=2s --timeout=1s`, + RunE: runCheckTCP, +} + func init() { checkUrlCmd.Flags().StringVar(&method, "method", "GET", "HTTP method") checkUrlCmd.Flags().StringVar(&body, "body", "", "HTTP POST/PUT content") - checkUrlCmd.Flags().IntVar(&retryCount, "retry", 1, "times to retry the HTTP call") + checkUrlCmd.Flags().IntVar(&retryCount, "retry", 0, "times to retry the HTTP call") checkUrlCmd.Flags().DurationVar(&retryDelay, "delay", 1*time.Second, "wait duration between retries") - checkUrlCmd.Flags().DurationVar(&timeout, "timeout", 2*time.Second, "timeout") + checkUrlCmd.Flags().DurationVar(&timeout, "timeout", 5*time.Second, "timeout") checkCmd.AddCommand(checkUrlCmd) + + checkTcpCmd.Flags().IntVar(&retryCount, "retry", 0, "times to retry the TCP check") + checkTcpCmd.Flags().DurationVar(&retryDelay, "delay", 1*time.Second, "wait duration between retries") + checkTcpCmd.Flags().DurationVar(&timeout, "timeout", 5*time.Second, "timeout") + checkCmd.AddCommand(checkTcpCmd) + rootCmd.AddCommand(checkCmd) } @@ -48,9 +63,13 @@ func runCheck(cmd *cobra.Command, args []string) error { return fmt.Errorf("--retry is required") } if len(args) < 1 { - return fmt.Errorf("address is required") + return fmt.Errorf("address is required! example: check http https://httpbin.org") } + url := args[0] + if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") { + url = fmt.Sprintf("http://%s", url) + } for n := 0; n <= retryCount; n++ { if n != 1 { @@ -60,7 +79,7 @@ func runCheck(cmd *cobra.Command, args []string) error { req, err := http.NewRequest(method, url, bytes.NewBuffer([]byte(body))) if err != nil { logger.Info("check failed", - zap.String("url", url), + zap.String("address", url), zap.Error(err)) os.Exit(1) } @@ -70,7 +89,7 @@ func runCheck(cmd *cobra.Command, args []string) error { cancel() if err != nil { logger.Info("check failed", - zap.String("url", url), + zap.String("address", url), zap.Error(err)) continue } @@ -81,13 +100,13 @@ func runCheck(cmd *cobra.Command, args []string) error { if resp.StatusCode >= 200 && resp.StatusCode < 400 { logger.Info("check succeed", - zap.String("url", url), + zap.String("address", url), zap.Int("status code", resp.StatusCode), zap.String("response size", fmtContentLength(resp.ContentLength))) os.Exit(0) } else { logger.Info("check failed", - zap.String("url", url), + zap.String("address", url), zap.Int("status code", resp.StatusCode)) continue } @@ -97,6 +116,39 @@ func runCheck(cmd *cobra.Command, args []string) error { return nil } +func runCheckTCP(cmd *cobra.Command, args []string) error { + if retryCount < 0 { + return fmt.Errorf("--retry is required") + } + if len(args) < 1 { + return fmt.Errorf("address is required! example: check tcp httpbin.org:80") + } + host := args[0] + + for n := 0; n <= retryCount; n++ { + if n != 1 { + time.Sleep(retryDelay) + } + + conn, err := net.DialTimeout("tcp", host, timeout) + + if err != nil { + logger.Info("check failed", + zap.String("address", host), + zap.Error(err)) + continue + } + + conn.Close() + logger.Info("check succeed", zap.String("address", host)) + os.Exit(0) + + } + + os.Exit(1) + return nil +} + func fmtContentLength(b int64) string { const unit = 1000 if b < unit { diff --git a/cmd/podcli/main.go b/cmd/podcli/main.go index 195d876..f59c22b 100644 --- a/cmd/podcli/main.go +++ b/cmd/podcli/main.go @@ -2,12 +2,12 @@ package main import ( "fmt" + "log" "os" "strings" "github.com/spf13/cobra" "go.uber.org/zap" - "log" ) var rootCmd = &cobra.Command{ @@ -18,7 +18,6 @@ podinfo command line utilities`, } var ( - configFile string logger *zap.Logger ) @@ -31,7 +30,6 @@ func main() { } defer logger.Sync() - rootCmd.PersistentFlags().StringVarP(&configFile, "config", "f", "", "path to config file") rootCmd.SetArgs(os.Args[1:]) if err := rootCmd.Execute(); err != nil { e := err.Error()