Close connection before exiting

os.Exit() calls in main() are causing `defer` statements to be ignored.
This surfaced in several bug reports (such as Java gRPC servers) as health
RPCs causing stack traces or "broken pipe" errors.

Signed-off-by: Ahmet Alp Balkan <[email protected]>

main.go

@@ -49,7 +49,6 @@ var (
 
 const (
 	// StatusInvalidArguments indicates specified invalid arguments.
-	StatusInvalidArguments = 1
 	// StatusConnectionFailure indicates connection failed.
 	StatusConnectionFailure = 2
 	// StatusRPCFailure indicates rpc failed.
@@ -59,6 +58,8 @@ const (
 )
 
 func init() {
+
+
 	log.SetFlags(0)
 	flag.StringVar(&flAddr, "addr", "", "(required) tcp host:port to connect")
 	flag.StringVar(&flService, "service", "", "service name to check (default: \"\")")
@@ -162,6 +163,9 @@ func buildCredentials(skipVerify bool, caCerts, clientCert, clientKey, serverNam
 }
 
 func main() {
+	retcode := 0
+	defer func() { os.Exit(retcode) }()
+
 	ctx, cancel := context.WithCancel(context.Background())
 
 	c := make(chan os.Signal, 1)
@@ -182,7 +186,8 @@ func main() {
 		creds, err := buildCredentials(flTLSNoVerify, flTLSCACert, flTLSClientCert, flTLSClientKey, flTLSServerName)
 		if err != nil {
 			log.Printf("failed to initialize tls credentials. error=%v", err)
-			os.Exit(StatusInvalidArguments)
+			retcode = StatusInvalidArguments
+			return
 		}
 		opts = append(opts, grpc.WithTransportCredentials(creds))
 	} else {
@@ -202,7 +207,8 @@ func main() {
 		} else {
 			log.Printf("error: failed to connect service at %q: %+v", flAddr, err)
 		}
-		os.Exit(StatusConnectionFailure)
+		retcode = StatusConnectionFailure
+		return
 	}
 	connDuration := time.Since(connStart)
 	defer conn.Close()
@@ -222,13 +228,15 @@ func main() {
 		} else {
 			log.Printf("error: health rpc failed: %+v", err)
 		}
-		os.Exit(StatusRPCFailure)
+		retcode = StatusRPCFailure
+		return
 	}
 	rpcDuration := time.Since(rpcStart)
 
 	if resp.GetStatus() != healthpb.HealthCheckResponse_SERVING {
 		log.Printf("service unhealthy (responded with %q)", resp.GetStatus().String())
-		os.Exit(StatusUnhealthy)
+		retcode = StatusUnhealthy
+		return
 	}
 	if flVerbose {
 		log.Printf("time elapsed: connect=%v rpc=%v", connDuration, rpcDuration)