Lets example servers gracefully shutdown

examples/BUILD.bazel

@@ -36,10 +36,12 @@ go_test(
         "//examples/server:go_default_library",
         "//examples/sub:go_default_library",
         "//runtime:go_default_library",
+        "@com_github_golang_glog//:go_default_library",
         "@com_github_golang_protobuf//jsonpb:go_default_library",
         "@com_github_golang_protobuf//proto:go_default_library",
         "@com_github_golang_protobuf//ptypes/empty:go_default_library",
         "@org_golang_google_genproto//googleapis/rpc/status:go_default_library",
         "@org_golang_google_grpc//codes:go_default_library",
+        "@org_golang_x_net//context:go_default_library",
     ],
 )

examples/integration_test.go

@@ -42,8 +42,13 @@ func TestEcho(t *testing.T) {
 }
 
 func TestForwardResponseOption(t *testing.T) {
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	go func() {
 		if err := Run(
+			ctx,
 			":8081",
 			runtime.WithForwardResponseOption(
 				func(_ context.Context, w http.ResponseWriter, _ proto.Message) error {

examples/main.go

@@ -82,8 +82,7 @@ func preflightHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 // Run starts a HTTP server and blocks forever if successful.
-func Run(address string, opts ...runtime.ServeMuxOption) error {
-	ctx := context.Background()
+func Run(ctx context.Context, address string, opts ...runtime.ServeMuxOption) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -96,14 +95,32 @@ func Run(address string, opts ...runtime.ServeMuxOption) error {
 	}
 	mux.Handle("/", gw)
 
-	return http.ListenAndServe(address, allowCORS(mux))
+	s := &http.Server{
+		Addr:    address,
+		Handler: allowCORS(mux),
+	}
+	go func() {
+		<-ctx.Done()
+		glog.Infof("Shutting down the http server")
+		if err := s.Shutdown(context.Background()); err != nil {
+			glog.Errorf("Failed to shutdown http server: %v", err)
+		}
+	}()
+
+	glog.Infof("Starting listening at %s", address)
+	if err := s.ListenAndServe(); err != http.ErrServerClosed {
+		glog.Errorf("Failed to listen and serve: %v", err)
+		return err
+	}
+	return nil
 }
 
 func main() {
 	flag.Parse()
 	defer glog.Flush()
 
-	if err := Run(":8080"); err != nil {
+	ctx := context.Background()
+	if err := Run(ctx, ":8080"); err != nil {
 		glog.Fatal(err)
 	}
 }

examples/main_test.go

@@ -7,18 +7,20 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang/glog"
 	server "github.com/grpc-ecosystem/grpc-gateway/examples/server"
+	"golang.org/x/net/context"
 )
 
-func runServers() <-chan error {
+func runServers(ctx context.Context) <-chan error {
 	ch := make(chan error, 2)
 	go func() {
-		if err := server.Run(); err != nil {
+		if err := server.Run(ctx); err != nil {
 			ch <- fmt.Errorf("cannot run grpc service: %v", err)
 		}
 	}()
 	go func() {
-		if err := Run(":8080"); err != nil {
+		if err := Run(ctx, ":8080"); err != nil {
 			ch <- fmt.Errorf("cannot run gateway service: %v", err)
 		}
 	}()
@@ -27,7 +29,11 @@ func runServers() <-chan error {
 
 func TestMain(m *testing.M) {
 	flag.Parse()
-	errCh := runServers()
+	defer glog.Flush()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	errCh := runServers(ctx)
 
 	ch := make(chan int, 1)
 	go func() {
@@ -40,6 +46,7 @@ func TestMain(m *testing.M) {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	case status := <-ch:
+		cancel()
 		os.Exit(status)
 	}
 }

examples/proto_error_test.go

@@ -10,24 +10,32 @@ import (
 
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
+	"golang.org/x/net/context"
 	spb "google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc/codes"
 )
 
+func runServer(ctx context.Context, t *testing.T, port uint16) {
+	opt := runtime.WithProtoErrorHandler(runtime.DefaultHTTPProtoErrorHandler)
+	if err := Run(ctx, fmt.Sprintf(":%d", port), opt); err != nil {
+		t.Errorf("gw.Run() failed with %v; want success", err)
+	}
+}
+
 func TestWithProtoErrorHandler(t *testing.T) {
-	go func() {
-		if err := Run(
-			":8082",
-			runtime.WithProtoErrorHandler(runtime.DefaultHTTPProtoErrorHandler),
-		); err != nil {
-			t.Errorf("gw.Run() failed with %v; want success", err)
-			return
-		}
-	}()
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	const port = 8082
+	go runServer(ctx, t, port)
 
+	// Waiting for the server's getting available.
+	// TODO(yugui) find a better way to wait
 	time.Sleep(100 * time.Millisecond)
-	testEcho(t, 8082, "application/json")
-	testEchoBody(t, 8082)
+
+	testEcho(t, port, "application/json")
+	testEchoBody(t, port)
 }
 
 func TestABEWithProtoErrorHandler(t *testing.T) {
@@ -36,19 +44,29 @@ func TestABEWithProtoErrorHandler(t *testing.T) {
 		return
 	}
 
-	testABECreate(t, 8082)
-	testABECreateBody(t, 8082)
-	testABEBulkCreate(t, 8082)
-	testABELookup(t, 8082)
-	testABELookupNotFoundWithProtoError(t)
-	testABEList(t, 8082)
-	testABEBulkEcho(t, 8082)
-	testABEBulkEchoZeroLength(t, 8082)
-	testAdditionalBindings(t, 8082)
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	const port = 8083
+	go runServer(ctx, t, port)
+	// Waiting for the server's getting available.
+	// TODO(yugui) find a better way to wait
+	time.Sleep(100 * time.Millisecond)
+
+	testABECreate(t, port)
+	testABECreateBody(t, port)
+	testABEBulkCreate(t, port)
+	testABELookup(t, port)
+	testABELookupNotFoundWithProtoError(t, port)
+	testABEList(t, port)
+	testABEBulkEcho(t, port)
+	testABEBulkEchoZeroLength(t, port)
+	testAdditionalBindings(t, port)
 }
 
-func testABELookupNotFoundWithProtoError(t *testing.T) {
-	url := "http://localhost:8082/v1/example/a_bit_of_everything"
+func testABELookupNotFoundWithProtoError(t *testing.T, port uint16) {
+	url := fmt.Sprintf("http://localhost:%d/v1/example/a_bit_of_everything", port)
 	uuid := "not_exist"
 	url = fmt.Sprintf("%s/%s", url, uuid)
 	resp, err := http.Get(url)
@@ -98,7 +116,18 @@ func testABELookupNotFoundWithProtoError(t *testing.T) {
 }
 
 func TestUnknownPathWithProtoError(t *testing.T) {
-	url := "http://localhost:8082"
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	const port = 8084
+	go runServer(ctx, t, port)
+
+	// Waiting for the server's getting available.
+	// TODO(yugui) find a better way to wait
+	time.Sleep(100 * time.Millisecond)
+
+	url := fmt.Sprintf("http://localhost:%d", port)
 	resp, err := http.Post(url, "application/json", strings.NewReader("{}"))
 	if err != nil {
 		t.Errorf("http.Post(%q) failed with %v; want success", url, err)
@@ -134,7 +163,18 @@ func TestUnknownPathWithProtoError(t *testing.T) {
 }
 
 func TestMethodNotAllowedWithProtoError(t *testing.T) {
-	url := "http://localhost:8082/v1/example/echo/myid"
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	const port = 8085
+	go runServer(ctx, t, port)
+
+	// Waiting for the server's getting available.
+	// TODO(yugui) find a better way to wait
+	time.Sleep(100 * time.Millisecond)
+
+	url := fmt.Sprintf("http://localhost:%d/v1/example/echo/myid", port)
 	resp, err := http.Get(url)
 	if err != nil {
 		t.Errorf("http.Post(%q) failed with %v; want success", url, err)

examples/server/BUILD.bazel

@@ -25,5 +25,6 @@ go_library(
         "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//metadata:go_default_library",
         "@org_golang_google_grpc//status:go_default_library",
+        "@org_golang_x_net//context:go_default_library",
     ],
 )

examples/server/cmd/example-server/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 
 	"github.com/golang/glog"
@@ -11,7 +12,8 @@ func main() {
 	flag.Parse()
 	defer glog.Flush()
 
-	if err := server.Run(); err != nil {
+	ctx := context.Background()
+	if err := server.Run(ctx); err != nil {
 		glog.Fatal(err)
 	}
 }

examples/server/main.go

@@ -3,15 +3,24 @@ package server
 import (
 	"net"
 
+	"github.com/golang/glog"
 	examples "github.com/grpc-ecosystem/grpc-gateway/examples/examplepb"
+	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 )
 
-func Run() error {
+// Run starts the example gRPC service.
+func Run(ctx context.Context) error {
 	l, err := net.Listen("tcp", ":9090")
 	if err != nil {
 		return err
 	}
+	defer func() {
+		if err := l.Close(); err != nil {
+			glog.Errorf("Failed to close tcp :9090: %v", err)
+		}
+	}()
+
 	s := grpc.NewServer()
 	examples.RegisterEchoServiceServer(s, newEchoServer())
 	examples.RegisterFlowCombinationServer(s, newFlowCombinationServer())
@@ -20,6 +29,9 @@ func Run() error {
 	examples.RegisterABitOfEverythingServiceServer(s, abe)
 	examples.RegisterStreamServiceServer(s, abe)
 
-	s.Serve(l)
-	return nil
+	go func() {
+		defer s.GracefulStop()
+		<-ctx.Done()
+	}()
+	return s.Serve(l)
 }