notify readiness when registered plugins are ready

henry118 · henry118 · commit bccaf68b7749 · 2023-05-27T00:32:25.000Z
Signed-off-by: Henry Wang <henwang@amazon.com> (cherry picked from commit 4bfcac8)
diff --git a/cmd/containerd/command/main.go b/cmd/containerd/command/main.go
@@ -271,12 +271,21 @@ can be used and modified as necessary as a custom configuration.`
 		}
 		serve(ctx, l, server.ServeGRPC)
 
-		if err := notifyReady(ctx); err != nil {
-			log.G(ctx).WithError(err).Warn("notify ready failed")
-		}
+		readyC := make(chan struct{})
+		go func() {
+			server.Wait()
+			close(readyC)
+		}()
 
-		log.G(ctx).Infof("containerd successfully booted in %fs", time.Since(start).Seconds())
-		<-done
+		select {
+		case <-readyC:
+			if err := notifyReady(ctx); err != nil {
+				log.G(ctx).WithError(err).Warn("notify ready failed")
+			}
+			log.G(ctx).Infof("containerd successfully booted in %fs", time.Since(start).Seconds())
+			<-done
+		case <-done:
+		}
 		return nil
 	}
 	return app
diff --git a/pkg/cri/cri.go b/pkg/cri/cri.go
@@ -60,6 +60,7 @@ func init() {
 }
 
 func initCRIService(ic *plugin.InitContext) (interface{}, error) {
+	ready := ic.RegisterReadiness()
 	ic.Meta.Platforms = []imagespec.Platform{platforms.DefaultSpec()}
 	ic.Meta.Exports = map[string]string{"CRIVersion": constants.CRIVersion, "CRIVersionAlpha": constants.CRIVersionAlpha}
 	ctx := ic.Context
@@ -103,7 +104,7 @@ func initCRIService(ic *plugin.InitContext) (interface{}, error) {
 	}
 
 	go func() {
-		if err := s.Run(); err != nil {
+		if err := s.Run(ready); err != nil {
 			log.G(ctx).WithError(err).Fatal("Failed to run CRI service")
 		}
 		// TODO(random-liu): Whether and how we can stop containerd.
diff --git a/pkg/cri/server/service.go b/pkg/cri/server/service.go
@@ -66,7 +66,7 @@ type grpcAlphaServices interface {
 
 // CRIService is the interface implement CRI remote service server.
 type CRIService interface {
-	Run() error
+	Run(ready func()) error
 	// io.Closer is used by containerd to gracefully stop cri service.
 	io.Closer
 	Register(*grpc.Server) error
@@ -200,7 +200,7 @@ func (c *criService) RegisterTCP(s *grpc.Server) error {
 }
 
 // Run starts the CRI service.
-func (c *criService) Run() error {
+func (c *criService) Run(ready func()) error {
 	logrus.Info("Start subscribing containerd event")
 	c.eventMonitor.subscribe(c.client)
 
@@ -251,6 +251,7 @@ func (c *criService) Run() error {
 
 	// Set the server as initialized. GRPC services could start serving traffic.
 	c.initialized.Set()
+	ready()
 
 	var eventMonitorErr, streamServerErr, cniNetConfMonitorErr error
 	// Stop the whole CRI service if any of the critical service exits.
diff --git a/plugin/context.go b/plugin/context.go
@@ -28,12 +28,13 @@ import (
 
 // InitContext is used for plugin initialization
 type InitContext struct {
-	Context      context.Context
-	Root         string
-	State        string
-	Config       interface{}
-	Address      string
-	TTRPCAddress string
+	Context           context.Context
+	Root              string
+	State             string
+	Config            interface{}
+	Address           string
+	TTRPCAddress      string
+	RegisterReadiness func() func()
 
 	// deprecated: will be removed in 2.0, use plugin.EventType
 	Events *exchange.Exchange
diff --git a/services/server/server.go b/services/server/server.go
@@ -231,6 +231,7 @@ func New(ctx context.Context, config *srvconfig.Config) (*Server, error) {
 		initContext.Events = events
 		initContext.Address = config.GRPC.Address
 		initContext.TTRPCAddress = config.TTRPC.Address
+		initContext.RegisterReadiness = s.RegisterReadiness
 
 		// load the plugin specific configuration if it is provided
 		if p.Config != nil {
@@ -306,6 +307,7 @@ type Server struct {
 	tcpServer   *grpc.Server
 	config      *srvconfig.Config
 	plugins     []*plugin.Plugin
+	ready       sync.WaitGroup
 }
 
 // ServeGRPC provides the containerd grpc APIs on the provided listener
@@ -383,6 +385,17 @@ func (s *Server) Stop() {
 	}
 }
 
+func (s *Server) RegisterReadiness() func() {
+	s.ready.Add(1)
+	return func() {
+		s.ready.Done()
+	}
+}
+
+func (s *Server) Wait() {
+	s.ready.Wait()
+}
+
 // LoadPlugins loads all plugins into containerd and generates an ordered graph
 // of all plugins.
 func LoadPlugins(ctx context.Context, config *srvconfig.Config) ([]*plugin.Registration, error) {

Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,7 @@ func init() {`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`func initCRIService(ic *plugin.InitContext) (interface{}, error) {`
	`63`	`+ ready := ic.RegisterReadiness()`
`63`	`64`	`ic.Meta.Platforms = []imagespec.Platform{platforms.DefaultSpec()}`
`64`	`65`	`ic.Meta.Exports = map[string]string{"CRIVersion": constants.CRIVersion, "CRIVersionAlpha": constants.CRIVersionAlpha}`
`65`	`66`	`ctx := ic.Context`
`@@ -103,7 +104,7 @@ func initCRIService(ic *plugin.InitContext) (interface{}, error) {`
`103`	`104`	`}`
`104`	`105`
`105`	`106`	`go func() {`
`106`		`- if err := s.Run(); err != nil {`
	`107`	`+ if err := s.Run(ready); err != nil {`
`107`	`108`	`log.G(ctx).WithError(err).Fatal("Failed to run CRI service")`
`108`	`109`	`}`
`109`	`110`	`// TODO(random-liu): Whether and how we can stop containerd.`