From db7ab961bf9b19971f9a3a339d430c3de4363069 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Mon, 22 May 2017 12:59:13 -0700
Subject: [PATCH] embed: share grpc connection for grpc json services

---
 embed/serve.go | 32 +++++++++++++++++++++-----------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/embed/serve.go b/embed/serve.go
index e43611f77..a70fa307a 100644
--- a/embed/serve.go
+++ b/embed/serve.go
@@ -160,28 +160,38 @@ func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Ha
 	})
 }
 
-type registerHandlerFunc func(context.Context, *gw.ServeMux, string, []grpc.DialOption) error
+type registerHandlerFunc func(context.Context, *gw.ServeMux, *grpc.ClientConn) error
 
 func (sctx *serveCtx) registerGateway(opts []grpc.DialOption) (*gw.ServeMux, error) {
 	ctx := sctx.ctx
-	addr := sctx.l.Addr().String()
+	conn, err := grpc.DialContext(ctx, sctx.l.Addr().String(), opts...)
+	if err != nil {
+		return nil, err
+	}
 	gwmux := gw.NewServeMux()
 
 	handlers := []registerHandlerFunc{
-		etcdservergw.RegisterKVHandlerFromEndpoint,
-		etcdservergw.RegisterWatchHandlerFromEndpoint,
-		etcdservergw.RegisterLeaseHandlerFromEndpoint,
-		etcdservergw.RegisterClusterHandlerFromEndpoint,
-		etcdservergw.RegisterMaintenanceHandlerFromEndpoint,
-		etcdservergw.RegisterAuthHandlerFromEndpoint,
-		v3lockgw.RegisterLockHandlerFromEndpoint,
-		v3electiongw.RegisterElectionHandlerFromEndpoint,
+		etcdservergw.RegisterKVHandler,
+		etcdservergw.RegisterWatchHandler,
+		etcdservergw.RegisterLeaseHandler,
+		etcdservergw.RegisterClusterHandler,
+		etcdservergw.RegisterMaintenanceHandler,
+		etcdservergw.RegisterAuthHandler,
+		v3lockgw.RegisterLockHandler,
+		v3electiongw.RegisterElectionHandler,
 	}
 	for _, h := range handlers {
-		if err := h(ctx, gwmux, addr, opts); err != nil {
+		if err := h(ctx, gwmux, conn); err != nil {
 			return nil, err
 		}
 	}
+	go func() {
+		<-ctx.Done()
+		if cerr := conn.Close(); cerr != nil {
+			plog.Warningf("failed to close conn to %s: %v", sctx.l.Addr().String(), cerr)
+		}
+	}()
+
 	return gwmux, nil
 }