Repository: firecracker-microvm/firecracker-containerd
Pull request #207 (2019-06-12 19:56:38)

Upgrade to v2 runc service in VM Agent.
The new runc service implementation supports managing multiple tasks in a single
service. This upgrade thus allows us to use a single task service instantiation
instead of 1-per-task, which in turn helps simplify our internal task manager
implementation.

Signed-off-by: Erik Sipsma <sipsma@amazon.com>

**NOTE:** This change stands on its own, but the primary motivation for doing it now is that it ends up simplifying a lot of the code for exec support, which I'm in the process of finalizing.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.



Patch:
diff --git a/agent/main.go b/agent/main.go
index fb2b493..ea5915b 100644
--- a/agent/main.go
+++ b/agent/main.go
@@ -78,7 +78,10 @@ func main() {
 	log.G(shimCtx).Info("creating task service")
 
 	eventExchange := &event.ExchangeCloser{Exchange: exchange.NewExchange()}
-	taskService := NewTaskService(shimCtx, shimCancel, eventExchange)
+	taskService, err := NewTaskService(shimCtx, shimCancel, eventExchange)
+	if err != nil {
+		log.G(shimCtx).WithError(err).Fatal("failed to create task service")
+	}
 
 	server, err := ttrpc.NewServer()
 	if err != nil {
diff --git a/agent/service.go b/agent/service.go
index fea7853..b3e3f40 100644
--- a/agent/service.go
+++ b/agent/service.go
@@ -22,7 +22,7 @@ import (
 
 	"github.com/containerd/containerd/cio"
 	"github.com/containerd/containerd/log"
-	runc "github.com/containerd/containerd/runtime/v2/runc/v1"
+	runc "github.com/containerd/containerd/runtime/v2/runc/v2"
 	"github.com/containerd/containerd/runtime/v2/shim"
 	taskAPI "github.com/containerd/containerd/runtime/v2/task"
 	"github.com/gogo/protobuf/types"
@@ -67,14 +67,23 @@ type TaskService struct {
 }
 
 // NewTaskService creates new runc shim wrapper
-func NewTaskService(shimCtx context.Context, shimCancel context.CancelFunc, publisher shim.Publisher) taskAPI.TaskService {
+func NewTaskService(shimCtx context.Context, shimCancel context.CancelFunc, publisher shim.Publisher) (taskAPI.TaskService, error) {
+	// We provide an empty string for "id" as the service manages multiple tasks; there is no single
+	// "id" being managed. As noted in the comments of the called code, the "id" arg is only used by
+	// the Cleanup function, so it will never be invoked as part of the task service API, which is all
+	// we need.
+	runcService, err := runc.New(shimCtx, "", publisher, shimCancel)
+	if err != nil {
+		return nil, err
+	}
+
 	return &TaskService{
-		taskManager: vm.NewTaskManager(log.G(shimCtx)),
+		taskManager: vm.NewTaskManager(log.G(shimCtx), runcService),
 
 		publisher:  publisher,
 		shimCtx:    shimCtx,
 		shimCancel: shimCancel,
-	}
+	}, nil
 }
 
 func logPanicAndDie(logger *logrus.Entry) {
@@ -128,11 +137,6 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 	// TODO if we update to the v2 runc implementation in containerd, we can use a single
 	// runc service instance to manage all tasks instead of creating a new one for each
 	taskCtx, taskCancel := context.WithCancel(ts.shimCtx)
-	runcService, err := runc.New(taskCtx, req.ID, ts.publisher, taskCancel)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to create runc shim for task")
-	}
-
 	defer func() {
 		if err != nil {
 			taskCancel()
@@ -159,7 +163,7 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 		fifoSet.Stderr = ""
 	}
 
-	task, err := ts.taskManager.AddTask(req.ID, runcService, bundleDir, extraData, fifoSet, taskCtx.Done(), taskCancel)
+	task, err := ts.taskManager.AddTask(req.ID, bundleDir, extraData, fifoSet, taskCtx.Done(), taskCancel)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to add task")
 	}
diff --git a/internal/vm/task.go b/internal/vm/task.go
index 977363c..33bf710 100644
--- a/internal/vm/task.go
+++ b/internal/vm/task.go
@@ -62,7 +62,7 @@ type VSockConnector func(ctx context.Context, port uint32) (net.Conn, error)
 // being executed via a firecracker-containerd runtime. It's intended to be
 // abstracted over whether it's being executed on the Host or inside a VM Guest.
 type TaskManager interface {
-	AddTask(string, taskAPI.TaskService, bundle.Dir, *proto.ExtraData, *cio.FIFOSet, <-chan struct{}, context.CancelFunc) (*Task, error)
+	AddTask(string, bundle.Dir, *proto.ExtraData, *cio.FIFOSet, <-chan struct{}, context.CancelFunc) (*Task, error)
 	Task(string) (*Task, error)
 	TaskCount() uint
 	Remove(string)
@@ -70,24 +70,25 @@ type TaskManager interface {
 }
 
 // NewTaskManager initializes a new TaskManager
-func NewTaskManager(logger *logrus.Entry) TaskManager {
+func NewTaskManager(logger *logrus.Entry, taskService taskAPI.TaskService) TaskManager {
 	return &taskManager{
-		tasks:  make(map[string]*Task),
-		logger: logger,
+		tasks:       make(map[string]*Task),
+		logger:      logger,
+		taskService: taskService,
 	}
 }
 
 type taskManager struct {
-	mu     sync.RWMutex
-	tasks  map[string]*Task
-	logger *logrus.Entry
+	mu          sync.RWMutex
+	tasks       map[string]*Task
+	logger      *logrus.Entry
+	taskService taskAPI.TaskService
 }
 
 // AddTask registers a task with the provided metadata with the taskManager.
 // taskService should implement the TaskService API for the task (i.e. Create, Kill, Exec, etc.).
 func (m *taskManager) AddTask(
 	containerID string,
-	taskService taskAPI.TaskService,
 	bundleDir bundle.Dir,
 	extraData *proto.ExtraData,
 	fifoSet *cio.FIFOSet,
@@ -107,7 +108,7 @@ func (m *taskManager) AddTask(
 	}
 
 	task := &Task{
-		TaskService: taskService,
+		TaskService: m.taskService,
 		ID:          containerID,
 		logger:      m.logger.WithField("id", containerID),
 
diff --git a/runtime/service.go b/runtime/service.go
index 6e08ef3..0b05f71 100644
--- a/runtime/service.go
+++ b/runtime/service.go
@@ -92,8 +92,6 @@ var (
 
 // implements shimapi
 type service struct {
-	taskManager vm.TaskManager
-
 	eventExchange *exchange.Exchange
 	namespace     string
 
@@ -128,6 +126,10 @@ type service struct {
 	agentClient       taskAPI.TaskService
 	eventBridgeClient eventbridge.Getter
 
+	// taskManager is only instantiated after CreateVM has run successfully.
+	// Any use of it must be guarded by service.waitVMReady()
+	taskManager vm.TaskManager
+
 	machine          *firecracker.Machine
 	machineConfig    *firecracker.Config
 	machineCID       uint32
@@ -172,8 +174,6 @@ func NewService(shimCtx context.Context, id string, remotePublisher shim.Publish
 	}
 
 	s := &service{
-		taskManager: vm.NewTaskManager(logger),
-
 		eventExchange: exchange.NewExchange(),
 		namespace:     namespace,
 
@@ -494,6 +494,7 @@ func (s *service) createVM(requestCtx context.Context, request *proto.CreateVMRe
 	rpcClient := ttrpc.NewClient(conn, ttrpc.WithOnClose(func() { _ = conn.Close() }))
 	s.agentClient = taskAPI.NewTaskClient(rpcClient)
 	s.eventBridgeClient = eventbridge.NewGetterClient(rpcClient)
+	s.taskManager = vm.NewTaskManager(s.logger, s.agentClient)
 
 	s.logger.Info("successfully started the VM")
 
@@ -687,7 +688,7 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 	}
 
 	taskCtx, taskCancel := context.WithCancel(s.shimCtx)
-	task, err := s.taskManager.AddTask(request.ID, s.agentClient, bundleDir, extraData, cio.NewFIFOSet(cio.Config{
+	task, err := s.taskManager.AddTask(request.ID, bundleDir, extraData, cio.NewFIFOSet(cio.Config{
 		Stdin:    request.Stdin,
 		Stdout:   request.Stdout,
 		Stderr:   request.Stderr,


Test patch:
diff --git a/internal/vm/task_test.go b/internal/vm/task_test.go
index 4571d29..957360e 100644
--- a/internal/vm/task_test.go
+++ b/internal/vm/task_test.go
@@ -55,7 +55,6 @@ type mockTaskService struct {
 func defaultMockTaskArgs(id string) addTaskArgs {
 	return addTaskArgs{
 		id:        fmt.Sprintf("container-%s", id),
-		ts:        &mockTaskService{},
 		bundleDir: mockBundleDir,
 		extraData: mockExtraData,
 		fifoSet:   mockFIFOSet,
@@ -64,26 +63,24 @@ func defaultMockTaskArgs(id string) addTaskArgs {
 
 type addTaskArgs struct {
 	id        string
-	ts        task.TaskService
 	bundleDir bundle.Dir
 	extraData *proto.ExtraData
 	fifoSet   *cio.FIFOSet
 }
 
 func addTaskFromArgs(tm TaskManager, args addTaskArgs) (*Task, error) {
-	return tm.AddTask(args.id, args.ts, args.bundleDir, args.extraData, args.fifoSet, context.Background().Done(), func() {})
+	return tm.AddTask(args.id, args.bundleDir, args.extraData, args.fifoSet, context.Background().Done(), func() {})
 }
 
 func TestTaskManager_AddRemoveTask(t *testing.T) {
 	logger, _ := test.NewNullLogger()
-	tm := NewTaskManager(logger.WithField("test", t.Name()))
+	tm := NewTaskManager(logger.WithField("test", t.Name()), &mockTaskService{})
 
 	taskAArgs := defaultMockTaskArgs("A")
 	taskBArgs := defaultMockTaskArgs("B")
 	taskCArgs := defaultMockTaskArgs("C")
 
 	assertExpectedTask := func(args addTaskArgs, createdTask *Task) {
-		require.Equal(t, args.ts, createdTask.TaskService, "AddTask sets expected task service")
 		require.Equal(t, args.id, createdTask.ID, "AddTask sets expected container ID")
 		require.Equal(t, args.extraData, createdTask.extraData, "AddTask sets expected container ID")
 		require.Equal(t, args.bundleDir, createdTask.bundleDir, "AddTask sets expected container ID")
@@ -180,7 +177,7 @@ func TestTaskManager_StartStdioProxy_VSockToFIFO(t *testing.T) {
 
 func testStdioProxy(t *testing.T, inputDirection IODirection) {
 	logger, _ := test.NewNullLogger()
-	tm := NewTaskManager(logger.WithField("test", t.Name()))
+	tm := NewTaskManager(logger.WithField("test", t.Name()), &mockTaskService{})
 
 	taskArgs := defaultMockTaskArgs("A")