Merge pull request #7215 from disksing/grpc-service

embed: support user defined grpc services.

embed/config.go

@@ -28,7 +28,9 @@ 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"
 )
 
 const (
@@ -126,6 +128,14 @@ 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. 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:"-"`
 }
 
 // configYAML holds the config suitable for yaml parsing

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()
 		}

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)