Don't call s.jailer.Close() from multiple places

All cleanup logic is moved to cleanup(). StopVM() will call the method
to make sure all resouces are deleted at the end of the request.
monitorVMExit() will call the method in case the VM is terminated
without having a StopVM request.

Signed-off-by: Kazuyoshi Kato <[email protected]>

Author: kzys

runtime/service.go

@@ -121,10 +121,7 @@ type service struct {
 	config *config.Config
 
 	// vmReady is closed once CreateVM has been successfully called
-	vmReady chan struct{}
-	// vmExitErr is set with the error returned by machine.Wait(). It should
-	// only be read after shimCtx.Done() is closed
-	vmExitErr                error
+	vmReady                  chan struct{}
 	vmStartOnce              sync.Once
 	agentClient              taskAPI.TaskService
 	eventBridgeClient        eventbridge.Getter
@@ -134,6 +131,9 @@ type service struct {
 	driveMountStubs          []MountableStubDrive
 	exitAfterAllTasksDeleted bool // exit the VM and shim when all tasks are deleted
 
+	cleanupErr  error
+	cleanupOnce sync.Once
+
 	machine          *firecracker.Machine
 	machineConfig    *firecracker.Config
 	vsockIOPortCount uint32
@@ -583,7 +583,6 @@ func (s *service) StopVM(requestCtx context.Context, request *proto.StopVMReques
 		timeout = time.Duration(request.TimeoutSeconds) * time.Second
 	}
 
-	defer s.shimCancel()
 	err = s.waitVMReady()
 	if err != nil {
 		return nil, err
@@ -1178,12 +1177,9 @@ func (s *service) shutdown(
 	if err := s.machine.Wait(context.Background()); err != nil {
 		result = multierror.Append(result, err)
 	}
-	if err := s.jailer.Close(); err != nil {
+	if err := s.cleanup(); err != nil {
 		result = multierror.Append(result, err)
 	}
-	if result == nil {
-		return nil
-	}
 
 	if err := result.ErrorOrNil(); err != nil {
 		return status.Error(codes.Internal, fmt.Sprintf("the VMM was killed forcibly: %v", err))
@@ -1223,13 +1219,6 @@ func (s *service) shutdownLoop(
 			},
 			timeout: jailerStopTimeout,
 		},
-		{
-			name: "cancel the context",
-			shutdown: func() error {
-				s.shimCancel()
-				return nil
-			},
-		},
 	}
 
 	for _, action := range actions {
@@ -1305,26 +1294,38 @@ func (s *service) Cleanup(requestCtx context.Context) (*taskAPI.DeleteResponse,
 	}, nil
 }
 
-func (s *service) monitorVMExit() {
-	defer func() {
+// cleanup resources
+func (s *service) cleanup() error {
+	s.cleanupOnce.Do(func() {
+		var result *multierror.Error
 		// we ignore the error here due to cleanup will only succeed if the jailing
 		// process was killed via SIGKILL
 		if err := s.jailer.Close(); err != nil {
+			result = multierror.Append(result, err)
 			s.logger.WithError(err).Error("failed to close jailer")
 		}
 
+		if err := s.publishVMStop(); err != nil {
+			result = multierror.Append(result, err)
+			s.logger.WithError(err).Error("failed to publish stop VM event")
+		}
+
 		// once the VM shuts down, the shim should too
 		s.shimCancel()
-	}()
 
+		s.cleanupErr = result.ErrorOrNil()
+	})
+	return s.cleanupErr
+}
+
+// monitorVMExit watches the VM and cleanup resources when it terminates.
+func (s *service) monitorVMExit() {
 	// Block until the VM exits
-	s.vmExitErr = s.machine.Wait(s.shimCtx)
-	if s.vmExitErr != nil && s.vmExitErr != context.Canceled {
-		s.logger.WithError(s.vmExitErr).Error("error returned from VM wait")
+	if err := s.machine.Wait(s.shimCtx); err != nil && err != context.Canceled {
+		s.logger.WithError(err).Error("error returned from VM wait")
 	}
 
-	publishErr := s.publishVMStop()
-	if publishErr != nil {
-		s.logger.WithError(publishErr).Error("failed to publish stop VM event")
+	if err := s.cleanup(); err != nil {
+		s.logger.WithError(err).Error("failed to clean up the VM")
 	}
 }