From 4374d944d484c1f59294696bd8642d72810ffc46 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Sun, 22 Jan 2017 18:21:19 +0800
Subject: [PATCH 1/2] embed: support user defined grpc services.

Fixes #7200
---
 embed/config.go | 3 +++
 embed/etcd.go   | 1 +
 embed/serve.go  | 9 ++++++++-
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/embed/config.go b/embed/config.go
index e1c6df2db..ad23b5c24 100644
--- a/embed/config.go
+++ b/embed/config.go
@@ -29,6 +29,7 @@ import (
 	"github.com/coreos/etcd/pkg/transport"
 	"github.com/coreos/etcd/pkg/types"
 	"github.com/ghodss/yaml"
+	"google.golang.org/grpc"
 )
 
 const (
@@ -126,6 +127,8 @@ type Config struct {
 	// The map key is the route path for the handler, and
 	// you must ensure it can't be conflicted with etcd's.
 	UserHandlers map[string]http.Handler `json:"-"`
+	// ServiceRegister is for registering users' gRPC services.
+	ServiceRegister func(*grpc.Server) `json:"-"`
 }
 
 // configYAML holds the config suitable for yaml parsing
diff --git a/embed/etcd.go b/embed/etcd.go
index c33efa4f8..828f3c7a9 100644
--- a/embed/etcd.go
+++ b/embed/etcd.go
@@ -289,6 +289,7 @@ func startClientListeners(cfg *Config) (sctxs map[string]*serveCtx, err error) {
 		for k := range cfg.UserHandlers {
 			sctx.userHandlers[k] = cfg.UserHandlers[k]
 		}
+		sctx.serviceRegister = cfg.ServiceRegister
 		if cfg.EnablePprof {
 			sctx.registerPprof()
 		}
diff --git a/embed/serve.go b/embed/serve.go
index d63638058..322e4b586 100644
--- a/embed/serve.go
+++ b/embed/serve.go
@@ -46,7 +46,8 @@ type serveCtx struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	userHandlers map[string]http.Handler
+	userHandlers    map[string]http.Handler
+	serviceRegister func(*grpc.Server)
 }
 
 func newServeCtx() *serveCtx {
@@ -66,6 +67,9 @@ func (sctx *serveCtx) serve(s *etcdserver.EtcdServer, tlscfg *tls.Config, handle
 
 	if sctx.insecure {
 		gs := v3rpc.Server(s, nil)
+		if sctx.serviceRegister != nil {
+			sctx.serviceRegister(gs)
+		}
 		grpcl := m.Match(cmux.HTTP2())
 		go func() { errc <- gs.Serve(grpcl) }()
 
@@ -90,6 +94,9 @@ func (sctx *serveCtx) serve(s *etcdserver.EtcdServer, tlscfg *tls.Config, handle
 
 	if sctx.secure {
 		gs := v3rpc.Server(s, tlscfg)
+		if sctx.serviceRegister != nil {
+			sctx.serviceRegister(gs)
+		}
 		handler = grpcHandlerFunc(gs, handler)
 
 		dtls := transport.ShallowCopyTLSConfig(tlscfg)

From 861cb5cfa2df26c5010f7410dc12d761a55d1528 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Mon, 23 Jan 2017 10:43:03 +0800
Subject: [PATCH 2/2] embed: add example for ServiceRegister.

---
 embed/config.go | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/embed/config.go b/embed/config.go
index ad23b5c24..ed647a063 100644
--- a/embed/config.go
+++ b/embed/config.go
@@ -28,6 +28,7 @@ import (
 	"github.com/coreos/etcd/pkg/netutil"
 	"github.com/coreos/etcd/pkg/transport"
 	"github.com/coreos/etcd/pkg/types"
+
 	"github.com/ghodss/yaml"
 	"google.golang.org/grpc"
 )
@@ -127,7 +128,13 @@ type Config struct {
 	// The map key is the route path for the handler, and
 	// you must ensure it can't be conflicted with etcd's.
 	UserHandlers map[string]http.Handler `json:"-"`
-	// ServiceRegister is for registering users' gRPC services.
+	// ServiceRegister is for registering users' gRPC services. A simple usage example:
+	//	cfg := embed.NewConfig()
+	//	cfg.ServerRegister = func(s *grpc.Server) {
+	//		pb.RegisterFooServer(s, &fooServer{})
+	//		pb.RegisterBarServer(s, &barServer{})
+	//	}
+	//	embed.StartEtcd(cfg)
 	ServiceRegister func(*grpc.Server) `json:"-"`
 }