diff --git a/embed/config.go b/embed/config.go index e1c6df2db..ed647a063 100644 --- a/embed/config.go +++ b/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 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)