Add CLI check tcp command

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 {

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()