shimv2: Improve shim shutdown logic

Gopkg.toml

@@ -52,7 +52,7 @@
 
 [[constraint]]
   name = "github.com/kata-containers/agent"
-  revision = "48dd1c031530fce9bf16b0f6a7305979cedd8fc9"
+  revision = "714bf625cc68505f31d802f98b84bf2b8006e146"
 
 [[constraint]]
   name = "github.com/containerd/cri-containerd"
@@ -64,15 +64,19 @@
 
 [[constraint]]
  name = "github.com/containerd/containerd"
- revision = "f05672357f56f26751a521175c5a96fc21fa8603"
+ revision = "25daa7355cdc1334718ce42e887f1c367aff9d0e"
 
 [[constraint]]
   name = "github.com/firecracker-microvm/firecracker-go-sdk"
   revision = "961461227bddf7e40a1d690634e866c343910f86"
 
 [[constraint]]
   name = "github.com/gogo/protobuf"
-  revision = "342cbe0a04158f6dcb03ca0079991a51a4248c02"
+  revision = "4cbf7e384e768b4e01799441fdf2a706a5635ae7"
+
+[[override]]
+  name = "github.com/containerd/ttrpc"
+  revision = "f82148331ad2181edea8f3f649a1f7add6c3f9c2"
 
 [[override]]
   branch = "master"

containerd-shim-v2/service.go

@@ -16,7 +16,6 @@ import (
 
 	eventstypes "github.com/containerd/containerd/api/events"
 	"github.com/containerd/containerd/errdefs"
-	"github.com/containerd/containerd/events"
 	"github.com/containerd/containerd/namespaces"
 	cdruntime "github.com/containerd/containerd/runtime"
 	cdshim "github.com/containerd/containerd/runtime/v2/shim"
@@ -58,7 +57,7 @@ var (
 var vci vc.VC = &vc.VCImpl{}
 
 // New returns a new shim service that can be used via GRPC
-func New(ctx context.Context, id string, publisher events.Publisher) (cdshim.Shim, error) {
+func New(ctx context.Context, id string, publisher cdshim.Publisher, shutdown func()) (cdshim.Shim, error) {
 	logger := logrus.WithField("ID", id)
 	// Discard the log before shim init its log output. Otherwise
 	// it will output into stdio, from which containerd would like
@@ -67,22 +66,20 @@ func New(ctx context.Context, id string, publisher events.Publisher) (cdshim.Shi
 	vci.SetLogger(ctx, logger)
 	katautils.SetLogger(ctx, logger, logger.Logger.Level)
 
-	ctx, cancel := context.WithCancel(ctx)
-
 	s := &service{
 		id:         id,
 		pid:        uint32(os.Getpid()),
 		ctx:        ctx,
 		containers: make(map[string]*container),
 		events:     make(chan interface{}, chSize),
 		ec:         make(chan exit, bufferSize),
-		cancel:     cancel,
+		cancel:     shutdown,
 		mount:      false,
 	}
 
 	go s.processExits()
 
-	go s.forward(publisher)
+	go s.forward(ctx, publisher)
 
 	return s, nil
 }
@@ -216,15 +213,20 @@ func (s *service) StartShim(ctx context.Context, id, containerdBinary, container
 	return address, nil
 }
 
-func (s *service) forward(publisher events.Publisher) {
+func (s *service) forward(ctx context.Context, publisher cdshim.Publisher) {
+	ns, _ := namespaces.Namespace(ctx)
+	ctx = namespaces.WithNamespace(context.Background(), ns)
+
 	for e := range s.events {
-		ctx, cancel := context.WithTimeout(s.ctx, timeOut)
+		ctx, cancel := context.WithTimeout(ctx, timeOut)
 		err := publisher.Publish(ctx, getTopic(e), e)
 		cancel()
 		if err != nil {
 			logrus.WithError(err).Error("post event")
 		}
 	}
+
+	publisher.Close()
 }
 
 func (s *service) send(evt interface{}) {
@@ -769,11 +771,8 @@ func (s *service) Shutdown(ctx context.Context, r *taskAPI.ShutdownRequest) (_ *
 	s.mu.Unlock()
 
 	s.cancel()
+	close(s.events)
 
-	os.Exit(0)
-
-	// This will never be called, but this is only there to make sure the
-	// program can compile.
 	return empty, nil
 }