Merge pull request #17 from gomicro/grpc_health

Grpc health

Author: dan9186

cmd/root.go

@@ -1,27 +1,42 @@
 package cmd
 
 import (
+	"context"
 	"crypto/tls"
 	"errors"
+	"fmt"
 	"net/http"
 	"os"
 
 	"github.com/certifi/gocertifi"
 	"github.com/spf13/cobra"
+	"google.golang.org/grpc"
+	health_pb "google.golang.org/grpc/health/grpc_health_v1"
 
-	"github.com/gomicro/probe/fmt"
+	ffmt "github.com/gomicro/probe/fmt"
 )
 
 var (
 	verbose    bool
 	skipVerify bool
+	grpcFlag   bool
+
+	// ErrGrpcBadStatus is the error returned when a non serving status is returend by a grpc service
+	ErrGrpcBadStatus = errors.New("grpc status: not ok")
+	// ErrGrpcConnFailure is the error returned when the grpc dial fails to connect to the specified host
+	ErrGrpcConnFailure = errors.New("grpc conn: failure")
+	// ErrHTTPBadStatus is the error returned when a non ok status is returned by a http service
+	ErrHTTPBadStatus = errors.New("http status: not ok")
+	// ErrHTTPGet is the errro returned when an http client encounters an error while performing a GET
+	ErrHTTPGet = errors.New("http client: get")
 )
 
 func init() {
 	cobra.OnInitialize(initEnvs)
 
 	rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "show more verbose output")
 	rootCmd.Flags().BoolVarP(&skipVerify, "insecure", "k", false, "permit operations for servers otherwise considered insecure")
+	rootCmd.Flags().BoolVarP(&grpcFlag, "grpc", "g", false, "use http2 transport for grpc health check")
 }
 
 func initEnvs() {
@@ -50,6 +65,26 @@ func Execute() {
 }
 
 func probe(cmd *cobra.Command, args []string) {
+	var err error
+
+	if grpcFlag {
+		err = probeGrpc(args[0])
+	} else {
+		err = probeHttp(args[0])
+	}
+
+	if err != nil {
+		if errors.Is(err, ErrHTTPGet) || errors.Is(err, ErrGrpcConnFailure) {
+			ffmt.Printf("%v", err.Error())
+			os.Exit(2)
+		}
+
+		ffmt.Verbosef("%v", err.Error())
+		os.Exit(1)
+	}
+}
+
+func probeHttp(host string) error {
 	pool, err := gocertifi.CACerts()
 	if err != nil {
 		fmt.Printf("Error: failed to create cert pool: %v\n", err.Error())
@@ -69,14 +104,40 @@ func probe(cmd *cobra.Command, args []string) {
 
 	client := &http.Client{Transport: transport}
 
-	resp, err := client.Get(args[0])
+	resp, err := client.Get(host)
 	if err != nil {
-		fmt.Verbosef("error: http get: %v", err.Error())
-		os.Exit(1)
+		return fmt.Errorf("%w: %v", ErrHTTPGet, err.Error())
 	}
 
 	if resp.StatusCode != http.StatusOK {
-		fmt.Verbosef("error: status %v", resp.StatusCode)
-		os.Exit(1)
+		return fmt.Errorf("%w: status %v", ErrHTTPBadStatus, resp.StatusCode)
+	}
+
+	return nil
+}
+
+func probeGrpc(host string) error {
+	ctx := context.Background()
+
+	opts := []grpc.DialOption{
+		grpc.WithBlock(),
+		grpc.WithInsecure(),
+	}
+
+	conn, err := grpc.DialContext(ctx, host, opts...)
+	if err != nil {
+		return fmt.Errorf("%w: %v", ErrGrpcConnFailure, err.Error())
 	}
+	defer conn.Close()
+
+	resp, err := health_pb.NewHealthClient(conn).Check(ctx, &health_pb.HealthCheckRequest{})
+	if err != nil {
+		return ErrGrpcBadStatus
+	}
+
+	if resp.GetStatus() != health_pb.HealthCheckResponse_SERVING {
+		return fmt.Errorf("%w: status %v", err, resp.GetStatus())
+	}
+
+	return nil
 }

cmd/root_test.go

@@ -1,7 +1,9 @@
 package cmd
 
 import (
+	"errors"
 	"fmt"
+	"net/http"
 	"testing"
 
 	"github.com/franela/goblin"
@@ -14,15 +16,37 @@ func TestProbe(t *testing.T) {
 	RegisterFailHandler(func(m string, _ ...int) { g.Fail(m) })
 
 	g.Describe("Probe", func() {
+		var server *bogus.Bogus
+		var host string
+
+		g.BeforeEach(func() {
+			server = bogus.New()
+			h, p := server.HostPort()
+			host = fmt.Sprintf("http://%v:%v/v1/status", h, p)
+		})
+
 		g.It("should probe an endpoint", func() {
-			server := bogus.New()
 			server.AddPath("/v1/status").
 				SetMethods("GET").
 				SetPayload([]byte("ok"))
-			h, p := server.HostPort()
 
-			args := []string{fmt.Sprintf("http://%v:%v/v1/status", h, p)}
-			probe(nil, args)
+			err := probeHttp(host)
+			Expect(err).NotTo(HaveOccurred())
+
+			hr := server.HitRecords()
+			Expect(len(hr)).To(Equal(1))
+			Expect(hr[0].Verb).To(Equal("GET"))
+		})
+
+		g.It("should error when a bad status is returned", func() {
+			server.AddPath("/v1/status").
+				SetMethods("GET").
+				SetStatus(http.StatusNotFound).
+				SetPayload([]byte("bad"))
+
+			err := probeHttp(host)
+			Expect(err).To(HaveOccurred())
+			Expect(errors.Is(err, ErrHTTPBadStatus)).To(BeTrue())
 
 			hr := server.HitRecords()
 			Expect(len(hr)).To(Equal(1))

go.mod

@@ -17,9 +17,7 @@ require (
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/spf13/viper v1.6.3
-	golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e // indirect
-	golang.org/x/sys v0.0.0-20200409092240-59c9f1ba88fa // indirect
-	golang.org/x/text v0.3.2 // indirect
-	golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 // indirect
+	golang.org/x/net v0.0.0-20210324205630-d1beb07c2056 // indirect
+	google.golang.org/grpc v1.36.1
 	gopkg.in/ini.v1 v1.55.0 // indirect
 )