shim: signal vminitd to shut down gracefully before stopping the VM

internal/shim/task/service.go

@@ -152,14 +152,21 @@ func (s *service) shutdown(ctx context.Context) error {
 		// to flush ext4 journals and dirty pages to the virtio-blk devices.
 		// Best-effort with a short retry for transient EBUSY.
 		if vmc, err := s.sb.Client(); err != nil {
-			log.G(ctx).WithError(err).Warn("failed to get VM client; skipping unmount of block volumes before VM shutdown")
+			log.G(ctx).WithError(err).Warn("failed to get VM client; skipping unmount and guest shutdown before VM stop")
 		} else {
 			unmountCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
 			err := unmountAllWithRetry(unmountCtx, mountAPI.NewTTRPCMountClient(vmc))
 			cancel()
 			if err != nil {
 				log.G(ctx).WithError(err).Warn("failed to unmount all block volumes before VM shutdown")
 			}
+
+			shutdownCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+			_, err = taskAPI.NewTTRPCTaskClient(vmc).Shutdown(shutdownCtx, &taskAPI.ShutdownRequest{})
+			cancel()
+			if err != nil {
+				log.G(ctx).WithError(err).Warn("failed to wait for guest shutdown before VM stop")
+			}
 		}
 
 		if err := s.sb.Stop(ctx); err != nil {
@@ -709,10 +716,6 @@ func (s *service) Connect(ctx context.Context, r *taskAPI.ConnectRequest) (*task
 func (s *service) Shutdown(ctx context.Context, r *taskAPI.ShutdownRequest) (*ptypes.Empty, error) {
 	log.G(ctx).WithFields(log.Fields{"id": r.ID}).Info("shutdown")
 
-	// TODO: Should we forward this to VM?
-	// tc := taskAPI.NewTTRPCTaskClient(s.vm.Client())
-	// return tc.Shutdown(ctx, r)
-
 	s.initiateShutdownOnce.Do(s.initiateShutdown)
 
 	select {

internal/vminit/task/service.go

@@ -649,18 +649,21 @@ func (s *service) Connect(ctx context.Context, r *taskAPI.ConnectRequest) (*task
 
 func (s *service) Shutdown(ctx context.Context, r *taskAPI.ShutdownRequest) (*ptypes.Empty, error) {
 	s.mu.Lock()
-	defer s.mu.Unlock()
-
 	// return out if the shim is still servicing containers
 	if len(s.containers) > 0 {
+		s.mu.Unlock()
 		return empty, nil
 	}
+	s.mu.Unlock()
 
-	// please make sure that temporary resource has been cleanup or registered
-	// for cleanup before calling shutdown
 	s.shutdown.Shutdown()
 
-	return empty, nil
+	select {
+	case <-s.shutdown.Done():
+		return empty, nil
+	case <-ctx.Done():
+		return nil, errgrpc.ToGRPC(ctx.Err())
+	}
 }
 
 func (s *service) Stats(ctx context.Context, r *taskAPI.StatsRequest) (*taskAPI.StatsResponse, error) {

pkg/vminit/initd/initd.go

@@ -307,7 +307,6 @@ func newService(ctx context.Context, config Config, shutdownSvc shutdown.Service
 	if err != nil {
 		return nil, err
 	}
-	shutdownSvc.RegisterCallback(ts.Shutdown)
 
 	registry.Register(&plugin.Registration{
 		Type: cplugins.InternalPlugin,