Rework task create and cleanup flow

Signed-off-by: Maksym Pavlenko <pavlenko.maksym@gmail.com>

runtime/v2/binary.go

@@ -60,7 +60,7 @@ type binary struct {
 	bundle                 *Bundle
 }
 
-func (b *binary) Start(ctx context.Context, opts *types.Any, onClose func()) (_ *shimTask, err error) {
+func (b *binary) Start(ctx context.Context, opts *types.Any, onClose func()) (_ *shim, err error) {
 	args := []string{"-id", b.bundle.ID}
 	switch logrus.GetLevel() {
 	case logrus.DebugLevel, logrus.TraceLevel:
@@ -128,12 +128,9 @@ func (b *binary) Start(ctx context.Context, opts *types.Any, onClose func()) (_
 		f.Close()
 	}
 	client := ttrpc.NewClient(conn, ttrpc.WithOnClose(onCloseWithShimLog))
-	return &shimTask{
+	return &shim{
-		shim: &shim{
 		bundle: b.bundle,
 		client: client,
-		},
-		task: task.NewTaskClient(client),
 	}, nil
 }
 

runtime/v2/manager.go

@@ -33,6 +33,7 @@ import (
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/containerd/plugin"
 	"github.com/containerd/containerd/runtime"
+	"github.com/containerd/containerd/runtime/v2/task"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 )
@@ -163,8 +164,8 @@ func (m *ShimManager) ID() string {
 	return fmt.Sprintf("%s.%s", plugin.RuntimePluginV2, "shim")
 }
 
-// Create a new task
+// Start launches a new shim instance
-func (m *ShimManager) Create(ctx context.Context, id string, opts runtime.CreateOpts) (_ runtime.Task, retErr error) {
+func (m *ShimManager) Start(ctx context.Context, id string, opts runtime.CreateOpts) (_ *shimTask, retErr error) {
 	bundle, err := NewBundle(ctx, m.root, m.state, id, opts.Spec.Value)
 	if err != nil {
 		return nil, err
@@ -185,19 +186,21 @@ func (m *ShimManager) Create(ctx context.Context, id string, opts runtime.Create
 		}
 	}()
 
-	t, err := shim.Create(ctx, opts)
+	// NOTE: temporarily keep this wrapper around until containerd's task service depends on it.
-	if err != nil {
+	// This will no longer be required once we migrate to client side task management.
-		return nil, errors.Wrap(err, "failed to create shim")
+	shimTask := &shimTask{
+		shim: shim,
+		task: task.NewTaskClient(shim.client),
 	}
 
-	if err := m.list.Add(ctx, t); err != nil {
+	if err := m.list.Add(ctx, shimTask); err != nil {
 		return nil, errors.Wrap(err, "failed to add task")
 	}
 
-	return t, nil
+	return shimTask, nil
 }
 
-func (m *ShimManager) startShim(ctx context.Context, bundle *Bundle, id string, opts runtime.CreateOpts) (*shimTask, error) {
+func (m *ShimManager) startShim(ctx context.Context, bundle *Bundle, id string, opts runtime.CreateOpts) (*shim, error) {
 	ns, err := namespaces.NamespaceRequired(ctx)
 	if err != nil {
 		return nil, err
@@ -232,50 +235,35 @@ func (m *ShimManager) startShim(ctx context.Context, bundle *Bundle, id string,
 }
 
 // cleanupShim attempts to properly delete and cleanup shim after error
-func (m *ShimManager) cleanupShim(shim *shimTask) {
+func (m *ShimManager) cleanupShim(shim *shim) {
 	dctx, cancel := timeout.WithContext(context.Background(), cleanupTimeout)
 	defer cancel()
 
-	_, errShim := shim.delete(dctx, m.list.Delete)
+	_ = shim.delete(dctx)
-	if errShim != nil {
+	m.list.Delete(dctx, shim.ID())
-		if errdefs.IsDeadlineExceeded(errShim) {
-			dctx, cancel = timeout.WithContext(context.Background(), cleanupTimeout)
-			defer cancel()
-		}
-		shim.Shutdown(dctx)
-		shim.Close()
-	}
 }
 
-// Get a specific task
+func (m *ShimManager) Get(ctx context.Context, id string) (*shim, error) {
-func (m *ShimManager) Get(ctx context.Context, id string) (runtime.Task, error) {
+	item, err := m.list.Get(ctx, id)
-	return m.list.Get(ctx, id)
+	if err != nil {
-}
+		return nil, err
+	}
 
-// Add a runtime task
+	shimTask := item.(*shimTask)
-func (m *ShimManager) Add(ctx context.Context, task runtime.Task) error {
+	return shimTask.shim, nil
-	return m.list.Add(ctx, task)
 }
 
 // Delete a runtime task
-func (m *ShimManager) Delete(ctx context.Context, id string) (*runtime.Exit, error) {
+func (m *ShimManager) Delete(ctx context.Context, id string) error {
-	task, err := m.list.Get(ctx, id)
+	shim, err := m.Get(ctx, id)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	shim := task.(*shimTask)
+	err = shim.delete(ctx)
-	exit, err := shim.delete(ctx, m.list.Delete)
+	m.list.Delete(ctx, id)
-	if err != nil {
-		return nil, err
-	}
 
-	return exit, err
+	return err
-}
-
-// Tasks lists all tasks
-func (m *ShimManager) Tasks(ctx context.Context, all bool) ([]runtime.Task, error) {
-	return m.list.GetAll(ctx, all)
 }
 
 func (m *ShimManager) loadExistingTasks(ctx context.Context) error {
@@ -434,20 +422,61 @@ func (m *TaskManager) ID() string {
 
 // Create launches new shim instance and creates new task
 func (m *TaskManager) Create(ctx context.Context, taskID string, opts runtime.CreateOpts) (runtime.Task, error) {
-	return m.shims.Create(ctx, taskID, opts)
+	shim, err := m.shims.Start(ctx, taskID, opts)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to start shim")
+	}
+
+	t, err := shim.Create(ctx, opts)
+	if err != nil {
+		dctx, cancel := timeout.WithContext(context.Background(), cleanupTimeout)
+		defer cancel()
+
+		_, errShim := shim.delete(dctx, func(ctx context.Context, id string) {
+			m.shims.list.Delete(ctx, id)
+		})
+
+		if errShim != nil {
+			if errdefs.IsDeadlineExceeded(errShim) {
+				dctx, cancel = timeout.WithContext(context.Background(), cleanupTimeout)
+				defer cancel()
+			}
+
+			shim.Shutdown(dctx)
+			shim.Close()
+		}
+
+		return nil, errors.Wrap(err, "failed to create shim task")
+	}
+
+	return t, nil
 }
 
 // Get a specific task
 func (m *TaskManager) Get(ctx context.Context, id string) (runtime.Task, error) {
-	return m.shims.Get(ctx, id)
+	return m.shims.list.Get(ctx, id)
 }
 
 // Tasks lists all tasks
 func (m *TaskManager) Tasks(ctx context.Context, all bool) ([]runtime.Task, error) {
-	return m.shims.Tasks(ctx, all)
+	return m.shims.list.GetAll(ctx, all)
 }
 
 // Delete deletes the task and shim instance
 func (m *TaskManager) Delete(ctx context.Context, taskID string) (*runtime.Exit, error) {
-	return m.shims.Delete(ctx, taskID)
+	item, err := m.shims.list.Get(ctx, taskID)
+	if err != nil {
+		return nil, err
+	}
+
+	shimTask := item.(*shimTask)
+	exit, err := shimTask.delete(ctx, func(ctx context.Context, id string) {
+		m.shims.list.Delete(ctx, id)
+	})
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to delete task: %w", err)
+	}
+
+	return exit, nil
 }

runtime/v2/shim.go

@@ -18,6 +18,7 @@ package v2
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"os"
 	"path/filepath"
@@ -37,6 +38,7 @@ import (
 	"github.com/containerd/containerd/runtime/v2/task"
 	"github.com/containerd/ttrpc"
 	ptypes "github.com/gogo/protobuf/types"
+	"github.com/hashicorp/go-multierror"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
@@ -202,6 +204,27 @@ func (s *shim) Close() error {
 	return s.client.Close()
 }
 
+func (s *shim) delete(ctx context.Context) error {
+	var (
+		result *multierror.Error
+	)
+
+	if err := s.Close(); err != nil {
+		result = multierror.Append(result, fmt.Errorf("failed to close ttrpc client: %w", err))
+	}
+
+	if err := s.client.UserOnCloseWait(ctx); err != nil {
+		result = multierror.Append(result, fmt.Errorf("close wait error: %w", err))
+	}
+
+	if err := s.bundle.Delete(); err != nil {
+		log.G(ctx).WithField("id", s.ID()).WithError(err).Error("failed to delete bundle")
+		result = multierror.Append(result, fmt.Errorf("failed to delete bundle: %w", err))
+	}
+
+	return result.ErrorOrNil()
+}
+
 var _ runtime.Task = &shimTask{}
 
 // shimTask wraps shim process and adds task service client for compatibility with existing shim manager.
@@ -267,20 +290,21 @@ func (s *shimTask) delete(ctx context.Context, removeTask func(ctx context.Conte
 	}
 
 	if err := s.waitShutdown(ctx); err != nil {
-		log.G(ctx).WithField("id", s.ID()).WithError(err).Error("failed to shutdown shim")
+		log.G(ctx).WithField("id", s.ID()).WithError(err).Error("failed to shutdown shim task")
+	}
+
+	if err := s.shim.delete(ctx); err != nil {
+		log.G(ctx).WithField("id", s.ID()).WithError(err).Errorf("failed to delete shim")
 	}
-	s.Close()
-	s.client.UserOnCloseWait(ctx)
 
 	// remove self from the runtime task list
 	// this seems dirty but it cleans up the API across runtimes, tasks, and the service
 	removeTask(ctx, s.ID())
-	if err := s.bundle.Delete(); err != nil {
+
-		log.G(ctx).WithField("id", s.ID()).WithError(err).Error("failed to delete bundle")
-	}
 	if shimErr != nil {
 		return nil, shimErr
 	}
+
 	return &runtime.Exit{
 		Status:    response.ExitStatus,
 		Timestamp: response.ExitedAt,