From 63d6aa151dd4edb2bd1bf44b5a35df6ab26fe4b4 Mon Sep 17 00:00:00 2001
From: Yuki Yugui Sonoda <yugui@yugui.jp>
Date: Fri, 27 Apr 2018 23:08:10 +0900
Subject: [PATCH] Also wait for the grpc server to get ready

---
 examples/gateway/BUILD.bazel |  1 +
 examples/gateway/gateway.go  | 13 +------------
 examples/gateway/handlers.go | 14 +++++++++++---
 examples/gateway/main.go     | 15 +++++++++++++--
 4 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/examples/gateway/BUILD.bazel b/examples/gateway/BUILD.bazel
index d076e613055..393a7c41f9c 100644
--- a/examples/gateway/BUILD.bazel
+++ b/examples/gateway/BUILD.bazel
@@ -15,5 +15,6 @@ go_library(
         "//runtime:go_default_library",
         "@com_github_golang_glog//:go_default_library",
         "@org_golang_google_grpc//:go_default_library",
+        "@org_golang_google_grpc//connectivity:go_default_library",
     ],
 )
diff --git a/examples/gateway/gateway.go b/examples/gateway/gateway.go
index 8430331bc6e..ffdbe2ca246 100644
--- a/examples/gateway/gateway.go
+++ b/examples/gateway/gateway.go
@@ -7,24 +7,13 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/golang/glog"
 	"github.com/grpc-ecosystem/grpc-gateway/examples/proto/examplepb"
 	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"google.golang.org/grpc"
 )
 
 // newGateway returns a new gateway server which translates HTTP into gRPC.
-func newGateway(ctx context.Context, network, addr string, opts []gwruntime.ServeMuxOption) (http.Handler, error) {
-	conn, err := dial(ctx, network, addr)
-	if err != nil {
-		return nil, err
-	}
-	go func() {
-		<-ctx.Done()
-		if err := conn.Close(); err != nil {
-			glog.Errorf("Failed to close a client connection to the gRPC server: %v", err)
-		}
-	}()
+func newGateway(ctx context.Context, conn *grpc.ClientConn, opts []gwruntime.ServeMuxOption) (http.Handler, error) {
 
 	mux := gwruntime.NewServeMux(opts...)
 
diff --git a/examples/gateway/handlers.go b/examples/gateway/handlers.go
index b0099bc3f60..d52580f52df 100644
--- a/examples/gateway/handlers.go
+++ b/examples/gateway/handlers.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 
 	"github.com/golang/glog"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 )
 
 func swaggerServer(dir string) http.HandlerFunc {
@@ -47,7 +49,13 @@ func preflightHandler(w http.ResponseWriter, r *http.Request) {
 	glog.Infof("preflight request for %s", r.URL.Path)
 }
 
-func healthzHandler(w http.ResponseWriter, r *http.Request) {
-	w.Header().Set("Content-Type", "text/plain")
-	fmt.Fprintln(w, "ok")
+func healthzServer(conn *grpc.ClientConn) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/plain")
+		if s := conn.GetState(); s != connectivity.Ready {
+			http.Error(w, fmt.Sprintf("grpc server is %s", s), http.StatusBadGateway)
+			return
+		}
+		fmt.Fprintln(w, "ok")
+	}
 }
diff --git a/examples/gateway/main.go b/examples/gateway/main.go
index d21cbcb4a0b..ea8c6836e0a 100644
--- a/examples/gateway/main.go
+++ b/examples/gateway/main.go
@@ -35,11 +35,22 @@ func Run(ctx context.Context, opts Options) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	conn, err := dial(ctx, opts.GRPCServer.Network, opts.GRPCServer.Addr)
+	if err != nil {
+		return err
+	}
+	go func() {
+		<-ctx.Done()
+		if err := conn.Close(); err != nil {
+			glog.Errorf("Failed to close a client connection to the gRPC server: %v", err)
+		}
+	}()
+
 	mux := http.NewServeMux()
 	mux.HandleFunc("/swagger/", swaggerServer(opts.SwaggerDir))
-	mux.HandleFunc("/healthz", healthzHandler)
+	mux.HandleFunc("/healthz", healthzServer(conn))
 
-	gw, err := newGateway(ctx, opts.GRPCServer.Network, opts.GRPCServer.Addr, opts.Mux)
+	gw, err := newGateway(ctx, conn, opts.Mux)
 	if err != nil {
 		return err
 	}