Repository: firecracker-microvm/firecracker-containerd
Pull request #179 (2019-05-08 00:57:37)

Cleanup use of context objects throughout the code.
There are a few related changes being made here:
1. Give each ctx object a more specific name that indicates the scope of its
   lifetime (i.e. one that will be canceled after a request or one that shuts
   down the process).
2. Remove ctx when it's not needed for anything but logging; just use a pre-made
   logger object instead
3. Use ctx to cancel requests where previously it was just ignored.

This is all spurred by the recently imported updates to TTRPC that resulted in
service method context objects to be canceled after the request is finished, but
the changes here seem beneficial besides that as its intended to make our use
of context more explicit and consistent.

Signed-off-by: Erik Sipsma <sipsma@amazon.com>

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 649014e..001c8ea 100644
--- a/agent/main.go
+++ b/agent/main.go
@@ -65,56 +65,54 @@ func main() {
 	signals := make(chan os.Signal, 32)
 	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, unix.SIGCHLD)
 
-	ctx, cancel := context.WithCancel(namespaces.WithNamespace(context.Background(), defaultNamespace))
-	defer cancel()
-
-	group, ctx := errgroup.WithContext(ctx)
+	shimCtx, shimCancel := context.WithCancel(namespaces.WithNamespace(context.Background(), defaultNamespace))
+	group, shimCtx := errgroup.WithContext(shimCtx)
 
 	// verify arguments, id is required
 	if id == "" {
-		log.G(ctx).WithError(errors.New("invalid argument")).Fatal("id not set")
+		log.G(shimCtx).WithError(errors.New("invalid argument")).Fatal("id not set")
 	}
 
 	// Ensure this process is a subreaper or else containers created via runc will
 	// not be its children.
 	if err := sys.SetSubreaper(enableSubreaper); err != nil {
-		log.G(ctx).WithError(err).Fatal("failed to set shim as subreaper")
+		log.G(shimCtx).WithError(err).Fatal("failed to set shim as subreaper")
 	}
 
 	// Create a runc task service that can be used via GRPC.
 	// This can be wrapped to add missing functionality (like
 	// running multiple containers inside one Firecracker VM)
 
-	log.G(ctx).WithField("id", id).Info("creating runc shim")
+	log.G(shimCtx).WithField("id", id).Info("creating runc shim")
 
 	eventExchange := &event.ExchangeCloser{Exchange: exchange.NewExchange()}
-	taskService := NewTaskService(eventExchange, cancel)
+	taskService := NewTaskService(shimCtx, shimCancel, eventExchange)
 
 	server, err := ttrpc.NewServer()
 	if err != nil {
-		log.G(ctx).WithError(err).Fatal("failed to create ttrpc server")
+		log.G(shimCtx).WithError(err).Fatal("failed to create ttrpc server")
 	}
 
 	taskAPI.RegisterTaskService(server, taskService)
-	eventbridge.RegisterGetterService(server, eventbridge.NewGetterService(ctx, eventExchange))
+	eventbridge.RegisterGetterService(server, eventbridge.NewGetterService(shimCtx, eventExchange))
 
 	// Run ttrpc over vsock
 
-	log.G(ctx).WithField("port", port).Info("listening to vsock")
+	log.G(shimCtx).WithField("port", port).Info("listening to vsock")
 	listener, err := vsock.Listen(uint32(port))
 	if err != nil {
-		log.G(ctx).WithError(err).Fatalf("failed to listen to vsock on port %d", port)
+		log.G(shimCtx).WithError(err).Fatalf("failed to listen to vsock on port %d", port)
 	}
 
 	group.Go(func() error {
-		return server.Serve(ctx, listener)
+		return server.Serve(shimCtx, listener)
 	})
 
 	group.Go(func() error {
 		defer func() {
-			log.G(ctx).Info("stopping ttrpc server")
-			if err := server.Shutdown(ctx); err != nil {
-				log.G(ctx).WithError(err).Errorf("failed to close ttrpc server")
+			log.G(shimCtx).Info("stopping ttrpc server")
+			if err := server.Shutdown(shimCtx); err != nil {
+				log.G(shimCtx).WithError(err).Errorf("failed to close ttrpc server")
 			}
 		}()
 
@@ -124,23 +122,23 @@ func main() {
 				switch s {
 				case unix.SIGCHLD:
 					if err := shim.Reap(); err != nil {
-						log.G(ctx).WithError(err).Error("reap error")
+						log.G(shimCtx).WithError(err).Error("reap error")
 					}
 				case syscall.SIGINT, syscall.SIGTERM:
-					cancel()
+					shimCancel()
 					return nil
 				}
-			case <-ctx.Done():
-				return ctx.Err()
+			case <-shimCtx.Done():
+				return shimCtx.Err()
 			}
 		}
 	})
 
 	err = group.Wait()
-	log.G(ctx).Info("shutting down agent")
+	log.G(shimCtx).Info("shutting down agent")
 
-	if err != nil {
-		log.G(ctx).WithError(err).Error("shim error")
+	if err != nil && err != context.Canceled {
+		log.G(shimCtx).WithError(err).Error("shim error")
 		panic(err)
 	}
 }
diff --git a/agent/service.go b/agent/service.go
index a83d143..19df5c1 100644
--- a/agent/service.go
+++ b/agent/service.go
@@ -22,7 +22,6 @@ import (
 
 	"github.com/containerd/containerd/cio"
 	"github.com/containerd/containerd/log"
-	"github.com/containerd/containerd/namespaces"
 	runc "github.com/containerd/containerd/runtime/v2/runc/v1"
 	"github.com/containerd/containerd/runtime/v2/shim"
 	taskAPI "github.com/containerd/containerd/runtime/v2/task"
@@ -48,16 +47,33 @@ type TaskService struct {
 	taskManager vm.TaskManager
 
 	publisher shim.Publisher
-	cancel    func()
+
+	// Normally, it's ill-advised to store a context object in a struct. However,
+	// in this case there appears to be little choice.
+	//
+	// The problem is that sometimes service methods, such as CreateVM, require
+	// a context scoped to the lifetime of the agent process but they are only provided
+	// by the TTRPC server a context scoped to the lifetime of the individual request,
+	// not the shimCtx.
+	//
+	// shimCtx should thus only be used by service methods that need to provide
+	// a context that will be canceled only when agent is shutting down and
+	// cleanup should happen.
+	//
+	// This approach is also taken by containerd's current reference runc shim
+	// v2 implementation
+	shimCtx    context.Context
+	shimCancel context.CancelFunc
 }
 
 // NewTaskService creates new runc shim wrapper
-func NewTaskService(publisher shim.Publisher, cancel context.CancelFunc) taskAPI.TaskService {
+func NewTaskService(shimCtx context.Context, shimCancel context.CancelFunc, publisher shim.Publisher) taskAPI.TaskService {
 	return &TaskService{
-		taskManager: vm.NewTaskManager(),
+		taskManager: vm.NewTaskManager(log.G(shimCtx)),
 
-		publisher: publisher,
-		cancel:    cancel,
+		publisher:  publisher,
+		shimCtx:    shimCtx,
+		shimCancel: shimCancel,
 	}
 }
 
@@ -79,11 +95,10 @@ func unmarshalExtraData(marshalled *types.Any) (*proto.ExtraData, error) {
 }
 
 // Create creates a new initial process and container using runc
-func (ts *TaskService) Create(ctx context.Context, req *taskAPI.CreateTaskRequest) (*taskAPI.CreateTaskResponse, error) {
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
+func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTaskRequest) (*taskAPI.CreateTaskResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	defer logPanicAndDie(log.G(ctx))
-	logger := log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "bundle": req.Bundle})
+	logger := log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "bundle": req.Bundle})
 	logger.Info("create")
 
 	extraData, err := unmarshalExtraData(req.Options)
@@ -112,15 +127,22 @@ func (ts *TaskService) Create(ctx context.Context, req *taskAPI.CreateTaskReques
 	// Create a runc shim to manage this task
 	// 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
-	runcService, err := runc.New(ctx, req.ID, ts.publisher, func() {})
+	taskCtx, taskCancel := context.WithCancel(ts.shimCtx)
+	runcService, err := runc.New(taskCtx, req.ID, ts.publisher, taskCancel)
 	if err != nil {
 		return nil, err
 	}
 
+	defer func() {
+		if err != nil {
+			taskCancel()
+		}
+	}()
+
 	// Override the incoming stdio FIFOs, which have paths from the host that we can't use
 	fifoSet, err := cio.NewFIFOSetInDir(bundleDir.RootPath(), req.ID, req.Terminal)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("failed opening stdio FIFOs")
+		logger.WithError(err).Error("failed opening stdio FIFOs")
 		return nil, err
 	}
 
@@ -137,7 +159,7 @@ func (ts *TaskService) Create(ctx context.Context, req *taskAPI.CreateTaskReques
 		fifoSet.Stderr = ""
 	}
 
-	task, err := ts.taskManager.AddTask(ctx, req.ID, runcService, bundleDir, extraData, fifoSet)
+	task, err := ts.taskManager.AddTask(req.ID, runcService, bundleDir, extraData, fifoSet, taskCtx.Done(), taskCancel)
 	if err != nil {
 		return nil, err
 	}
@@ -156,12 +178,12 @@ func (ts *TaskService) Create(ctx context.Context, req *taskAPI.CreateTaskReques
 
 	// Start the io proxy and wait for initialization to complete before starting
 	// the task to ensure we capture all task output
-	err = <-task.StartStdioProxy(ctx, vm.VSockToFIFO, acceptVSock)
+	err = <-task.StartStdioProxy(taskCtx, vm.VSockToFIFO, acceptVSock)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Create(ctx, req)
+	resp, err := task.Create(requestCtx, req)
 	if err != nil {
 		logger.WithError(err).Error("error creating container")
 		return nil, err
@@ -181,49 +203,52 @@ func isTemporaryNetErr(err error) bool {
 
 var _ vm.VSockConnector = acceptVSock
 
-func acceptVSock(ctx context.Context, port uint32) (net.Conn, error) {
+func acceptVSock(requestCtx context.Context, port uint32) (net.Conn, error) {
 	listener, err := vsock.Listen(port)
 	if err != nil {
 		return nil, errors.Wrap(err, "unable to listen on vsock")
 	}
 
 	for range time.Tick(10 * time.Millisecond) {
-		// accept is non-blocking so try to accept until we get a connection
-		conn, err := listener.Accept()
-		if err == nil {
-			return conn, nil
-		}
-
-		if isTemporaryNetErr(err) {
-			log.G(ctx).WithError(err).Debug("temporary stdio vsock accept failure")
-			continue
+		select {
+		case <-requestCtx.Done():
+			return nil, requestCtx.Err()
+		default:
+			// accept is non-blocking so try to accept until we get a connection
+			conn, err := listener.Accept()
+			if err == nil {
+				return conn, nil
+			}
+
+			if isTemporaryNetErr(err) {
+				log.G(requestCtx).WithError(err).Debug("temporary stdio vsock accept failure")
+			} else {
+				log.G(requestCtx).WithError(err).Error("non-temporary stdio vsock accept failure")
+				return nil, err
+			}
 		}
-
-		log.G(ctx).WithError(err).Error("non-temporary stdio vsock accept failure")
-		return nil, err
 	}
 
 	panic("unreachable code") // appeases the compiler, which doesn't know the for loop is infinite
 }
 
 // State returns process state information
-func (ts *TaskService) State(ctx context.Context, req *taskAPI.StateRequest) (*taskAPI.StateResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) State(requestCtx context.Context, req *taskAPI.StateRequest) (*taskAPI.StateResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("state")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("state")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.State(ctx, req)
+	resp, err := task.State(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("state failed")
+		log.G(requestCtx).WithError(err).Error("state failed")
 		return nil, err
 	}
 
-	log.G(ctx).WithFields(logrus.Fields{
+	log.G(requestCtx).WithFields(logrus.Fields{
 		"id":     resp.ID,
 		"bundle": resp.Bundle,
 		"pid":    resp.Pid,
@@ -233,44 +258,44 @@ func (ts *TaskService) State(ctx context.Context, req *taskAPI.StateRequest) (*t
 }
 
 // Start starts a process
-func (ts *TaskService) Start(ctx context.Context, req *taskAPI.StartRequest) (*taskAPI.StartResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Start(requestCtx context.Context, req *taskAPI.StartRequest) (*taskAPI.StartResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("start")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("start")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Start(ctx, req)
+	resp, err := task.Start(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("start failed")
+		log.G(requestCtx).WithError(err).Error("start failed")
 		return nil, err
 	}
 
-	log.G(ctx).WithField("pid", resp.Pid).Debug("start succeeded")
+	log.G(requestCtx).WithField("pid", resp.Pid).Debug("start succeeded")
 	return resp, nil
 }
 
 // Delete deletes the initial process and container
-func (ts *TaskService) Delete(ctx context.Context, req *taskAPI.DeleteRequest) (*taskAPI.DeleteResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Delete(requestCtx context.Context, req *taskAPI.DeleteRequest) (*taskAPI.DeleteResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("delete")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("delete")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Delete(ctx, req)
+	resp, err := task.Delete(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("delete failed")
+		log.G(requestCtx).WithError(err).Error("delete failed")
 		return nil, err
 	}
 
-	log.G(ctx).WithFields(logrus.Fields{
+	ts.taskManager.Remove(req.ID)
+
+	log.G(requestCtx).WithFields(logrus.Fields{
 		"pid":         resp.Pid,
 		"exit_status": resp.ExitStatus,
 	}).Debug("delete succeeded")
@@ -278,254 +303,242 @@ func (ts *TaskService) Delete(ctx context.Context, req *taskAPI.DeleteRequest) (
 }
 
 // Pids returns all pids inside the container
-func (ts *TaskService) Pids(ctx context.Context, req *taskAPI.PidsRequest) (*taskAPI.PidsResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Pids(requestCtx context.Context, req *taskAPI.PidsRequest) (*taskAPI.PidsResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("pids")
+	log.G(requestCtx).WithField("id", req.ID).Debug("pids")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Pids(ctx, req)
+	resp, err := task.Pids(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("pids failed")
+		log.G(requestCtx).WithError(err).Error("pids failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("pids succeeded")
+	log.G(requestCtx).Debug("pids succeeded")
 	return resp, nil
 }
 
 // Pause pauses the container
-func (ts *TaskService) Pause(ctx context.Context, req *taskAPI.PauseRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Pause(requestCtx context.Context, req *taskAPI.PauseRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("pause")
+	log.G(requestCtx).WithField("id", req.ID).Debug("pause")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Pause(ctx, req)
+	resp, err := task.Pause(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("pause failed")
+		log.G(requestCtx).WithError(err).Error("pause failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("pause succeeded")
+	log.G(requestCtx).Debug("pause succeeded")
 	return resp, nil
 }
 
 // Resume resumes the container
-func (ts *TaskService) Resume(ctx context.Context, req *taskAPI.ResumeRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Resume(requestCtx context.Context, req *taskAPI.ResumeRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("resume")
+	log.G(requestCtx).WithField("id", req.ID).Debug("resume")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Resume(ctx, req)
+	resp, err := task.Resume(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Debug("resume failed")
+		log.G(requestCtx).WithError(err).Debug("resume failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("resume succeeded")
+	log.G(requestCtx).Debug("resume succeeded")
 	return resp, nil
 }
 
 // Checkpoint saves the state of the container instance
-func (ts *TaskService) Checkpoint(ctx context.Context, req *taskAPI.CheckpointTaskRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Checkpoint(requestCtx context.Context, req *taskAPI.CheckpointTaskRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "path": req.Path}).Info("checkpoint")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "path": req.Path}).Info("checkpoint")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Checkpoint(ctx, req)
+	resp, err := task.Checkpoint(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("checkout failed")
+		log.G(requestCtx).WithError(err).Error("checkout failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("checkpoint sutaskAPI")
+	log.G(requestCtx).Debug("checkpoint sutaskAPI")
 	return resp, nil
 }
 
 // Kill kills a process with the provided signal
-func (ts *TaskService) Kill(ctx context.Context, req *taskAPI.KillRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Kill(requestCtx context.Context, req *taskAPI.KillRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("kill")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("kill")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Kill(ctx, req)
+	resp, err := task.Kill(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("kill failed")
+		log.G(requestCtx).WithError(err).Error("kill failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("kill succeeded")
+	log.G(requestCtx).Debug("kill succeeded")
 	return resp, nil
 }
 
 // Exec runs an additional process inside the container
-func (ts *TaskService) Exec(ctx context.Context, req *taskAPI.ExecProcessRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Exec(requestCtx context.Context, req *taskAPI.ExecProcessRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("exec")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("exec")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Exec(ctx, req)
+	resp, err := task.Exec(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("exec failed")
+		log.G(requestCtx).WithError(err).Error("exec failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("exec succeeded")
+	log.G(requestCtx).Debug("exec succeeded")
 	return resp, nil
 }
 
 // ResizePty resizes pty
-func (ts *TaskService) ResizePty(ctx context.Context, req *taskAPI.ResizePtyRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) ResizePty(requestCtx context.Context, req *taskAPI.ResizePtyRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("resize_pty")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("resize_pty")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.ResizePty(ctx, req)
+	resp, err := task.ResizePty(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("resize_pty failed")
+		log.G(requestCtx).WithError(err).Error("resize_pty failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("resize_pty succeeded")
+	log.G(requestCtx).Debug("resize_pty succeeded")
 	return resp, nil
 }
 
 // CloseIO closes all IO inside container
-func (ts *TaskService) CloseIO(ctx context.Context, req *taskAPI.CloseIORequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) CloseIO(requestCtx context.Context, req *taskAPI.CloseIORequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("close_io")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("close_io")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.CloseIO(ctx, req)
+	resp, err := task.CloseIO(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("close io failed")
+		log.G(requestCtx).WithError(err).Error("close io failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("close io succeeded")
+	log.G(requestCtx).Debug("close io succeeded")
 	return resp, nil
 }
 
 // Update updates running container
-func (ts *TaskService) Update(ctx context.Context, req *taskAPI.UpdateTaskRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Update(requestCtx context.Context, req *taskAPI.UpdateTaskRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("update")
+	log.G(requestCtx).WithField("id", req.ID).Debug("update")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Update(ctx, req)
+	resp, err := task.Update(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("update failed")
+		log.G(requestCtx).WithError(err).Error("update failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("update succeeded")
+	log.G(requestCtx).Debug("update succeeded")
 	return resp, nil
 }
 
 // Wait waits for a process to exit
-func (ts *TaskService) Wait(ctx context.Context, req *taskAPI.WaitRequest) (*taskAPI.WaitResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Wait(requestCtx context.Context, req *taskAPI.WaitRequest) (*taskAPI.WaitResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("wait")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("wait")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Wait(ctx, req)
+	resp, err := task.Wait(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("wait failed")
+		log.G(requestCtx).WithError(err).Error("wait failed")
 		return nil, err
 	}
 
-	log.G(ctx).WithField("exit_status", resp.ExitStatus).Debug("wait succeeded")
+	log.G(requestCtx).WithField("exit_status", resp.ExitStatus).Debug("wait succeeded")
 	return resp, nil
 }
 
 // Stats returns a process stats
-func (ts *TaskService) Stats(ctx context.Context, req *taskAPI.StatsRequest) (*taskAPI.StatsResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Stats(requestCtx context.Context, req *taskAPI.StatsRequest) (*taskAPI.StatsResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("stats")
+	log.G(requestCtx).WithField("id", req.ID).Debug("stats")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Stats(ctx, req)
+	resp, err := task.Stats(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("stats failed")
+		log.G(requestCtx).WithError(err).Error("stats failed")
 		return nil, err
 	}
 
-	log.G(ctx).Debug("stats succeeded")
+	log.G(requestCtx).Debug("stats succeeded")
 	return resp, nil
 }
 
 // Connect returns shim information such as the shim's pid
-func (ts *TaskService) Connect(ctx context.Context, req *taskAPI.ConnectRequest) (*taskAPI.ConnectResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Connect(requestCtx context.Context, req *taskAPI.ConnectRequest) (*taskAPI.ConnectResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("connect")
+	log.G(requestCtx).WithField("id", req.ID).Debug("connect")
 	task, err := ts.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
-	resp, err := task.Connect(ctx, req)
+	resp, err := task.Connect(requestCtx, req)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("connect failed")
+		log.G(requestCtx).WithError(err).Error("connect failed")
 		return nil, err
 	}
 
-	log.G(ctx).WithFields(logrus.Fields{
+	log.G(requestCtx).WithFields(logrus.Fields{
 		"shim_pid": resp.ShimPid,
 		"task_pid": resp.TaskPid,
 		"version":  resp.Version,
@@ -534,15 +547,17 @@ func (ts *TaskService) Connect(ctx context.Context, req *taskAPI.ConnectRequest)
 }
 
 // Shutdown cleanups runc resources ans gracefully shutdowns ttrpc server
-func (ts *TaskService) Shutdown(ctx context.Context, req *taskAPI.ShutdownRequest) (*types.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (ts *TaskService) Shutdown(requestCtx context.Context, req *taskAPI.ShutdownRequest) (*types.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
+
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "now": req.Now}).Debug("shutdown")
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "now": req.Now}).Debug("shutdown")
-	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
+	ts.taskManager.RemoveAll()
 
-	ts.taskManager.RemoveAll(ctx)
-	ts.cancel()
+	// shimCancel will result in each ctx provided to runc shims to be canceled in addition to unblocking agent's
+	// main func, which will allow it to exit gracefully.
+	defer ts.shimCancel()
 
-	log.G(ctx).Debug("going to gracefully shutdown agent")
+	log.G(requestCtx).Debug("going to gracefully shutdown agent")
 	return &types.Empty{}, nil
 }
diff --git a/go.sum b/go.sum
index 94c510a..9268e87 100644
--- a/go.sum
+++ b/go.sum
@@ -28,8 +28,6 @@ github.com/containerd/fifo v0.0.0-20180307165137-3d5202aec260 h1:XGyg7oTtD0DoRFh
 github.com/containerd/fifo v0.0.0-20180307165137-3d5202aec260/go.mod h1:ODA38xgv3Kuk8dQz2ZQXpnv/UZZUHUCL7pnLehbXgQI=
 github.com/containerd/go-runc v0.0.0-20190226155025-7d11b49dc076 h1:biYGul6kwz5/Lh6UxkYQyRKPeYA/ZL5n2pdCAcYUpLA=
 github.com/containerd/go-runc v0.0.0-20190226155025-7d11b49dc076/go.mod h1:IV7qH3hrUgRmyYrtgEeGWJfWbgcHL9CSRruz2Vqcph0=
-github.com/containerd/ttrpc v0.0.0-20181001154009-f51df4475b76 h1:vUPO9S35+FvukX2L/1mmPWgdv0y6FEKJzeWzmgVe9IA=
-github.com/containerd/ttrpc v0.0.0-20181001154009-f51df4475b76/go.mod h1:PvCDdDGpgqzQIzDW1TphrGLssLDZp2GuS+X5DkEJB8o=
 github.com/containerd/ttrpc v0.0.0-20190411181408-699c4e40d1e7 h1:SKDlsIhYxNE1LO0xwuOR+3QWj3zRibVQu5jWIMQmOfU=
 github.com/containerd/ttrpc v0.0.0-20190411181408-699c4e40d1e7/go.mod h1:PvCDdDGpgqzQIzDW1TphrGLssLDZp2GuS+X5DkEJB8o=
 github.com/containerd/typeurl v0.0.0-20181015155603-461401dc8f19 h1:gzdItdct+4eLnZxiZi1YcIXx3uo5QWa/xXKnsldEqY8=
diff --git a/internal/vm/task.go b/internal/vm/task.go
index 3b2c834..83d56e4 100644
--- a/internal/vm/task.go
+++ b/internal/vm/task.go
@@ -21,12 +21,11 @@ import (
 	"syscall"
 
 	"github.com/containerd/containerd/cio"
-	"github.com/containerd/containerd/log"
 	taskAPI "github.com/containerd/containerd/runtime/v2/task"
 	"github.com/containerd/fifo"
 	ptypes "github.com/gogo/protobuf/types"
-	"github.com/hashicorp/go-multierror"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
 
 	"github.com/firecracker-microvm/firecracker-containerd/internal"
 	"github.com/firecracker-microvm/firecracker-containerd/internal/bundle"
@@ -62,34 +61,37 @@ 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(context.Context, string, taskAPI.TaskService, bundle.Dir, *proto.ExtraData, *cio.FIFOSet) (*Task, error)
+	AddTask(string, taskAPI.TaskService, bundle.Dir, *proto.ExtraData, *cio.FIFOSet, <-chan struct{}, context.CancelFunc) (*Task, error)
 	Task(string) (*Task, error)
 	TaskCount() uint
-	Remove(context.Context, string)
-	RemoveAll(context.Context)
+	Remove(string)
+	RemoveAll()
 }
 
 // NewTaskManager initializes a new TaskManager
-func NewTaskManager() TaskManager {
+func NewTaskManager(logger *logrus.Entry) TaskManager {
 	return &taskManager{
-		tasks: make(map[string]*Task),
+		tasks:  make(map[string]*Task),
+		logger: logger,
 	}
 }
 
 type taskManager struct {
-	mu    sync.RWMutex
-	tasks map[string]*Task
+	mu     sync.RWMutex
+	tasks  map[string]*Task
+	logger *logrus.Entry
 }
 
 // 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(
-	ctx context.Context,
 	containerID string,
 	taskService taskAPI.TaskService,
 	bundleDir bundle.Dir,
 	extraData *proto.ExtraData,
 	fifoSet *cio.FIFOSet,
+	doneCh <-chan struct{},
+	cancel context.CancelFunc,
 ) (*Task, error) {
 	err := bundleDir.Create() // no-op if it already exists
 	if err != nil {
@@ -106,12 +108,14 @@ func (m *taskManager) AddTask(
 	task := &Task{
 		TaskService: taskService,
 		ID:          containerID,
+		logger:      m.logger.WithField("id", containerID),
+
+		doneCh: doneCh,
+		cancel: cancel,
 
 		extraData: extraData,
 		bundleDir: bundleDir,
 		fifoSet:   fifoSet,
-
-		cancelCh: make(chan struct{}),
 	}
 
 	m.tasks[task.ID] = task
@@ -131,12 +135,9 @@ func (m *taskManager) get(containerID string) (*Task, error) {
 }
 
 // assumes caller is holding m.mu lock
-func (m *taskManager) remove(ctx context.Context, containerID string, task *Task) {
+func (m *taskManager) remove(containerID string, task *Task) {
 	delete(m.tasks, containerID)
-	err := task.CancelIO()
-	if err != nil {
-		log.G(ctx).WithField("id", containerID).Error("failed to cancel task io")
-	}
+	task.cancel()
 }
 
 // Task returns the TaskService API for the task with the provided containerID, allowing
@@ -155,7 +156,7 @@ func (m *taskManager) TaskCount() uint {
 
 // Remove unregisters the task with the provided containerID with the taskManager and cancels
 // any ongoing I/O proxying for the task.
-func (m *taskManager) Remove(ctx context.Context, containerID string) {
+func (m *taskManager) Remove(containerID string) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -165,29 +166,31 @@ func (m *taskManager) Remove(ctx context.Context, containerID string) {
 		return
 	}
 
-	m.remove(ctx, containerID, task)
+	m.remove(containerID, task)
 }
 
 // RemoveAll calls Remove() on all tasks registered with the taskManager.
-func (m *taskManager) RemoveAll(ctx context.Context) {
+func (m *taskManager) RemoveAll() {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
 	for containerID, task := range m.tasks {
-		m.remove(ctx, containerID, task)
+		m.remove(containerID, task)
 	}
 }
 
 // Task represents a containerd Task under the control of a taskManager
 type Task struct {
 	taskAPI.TaskService
-	ID string
+	ID     string
+	logger *logrus.Entry
+
+	doneCh <-chan struct{}
+	cancel context.CancelFunc
 
 	extraData *proto.ExtraData
 	bundleDir bundle.Dir
 	fifoSet   *cio.FIFOSet
-
-	cancelCh chan struct{}
 }
 
 // ExtraData returns the extra data the VM Agent needs to execute this task
@@ -205,13 +208,6 @@ func (task *Task) BundleDir() bundle.Dir {
 	return task.bundleDir
 }
 
-// CancelIO will stop any ongoing IO proxying. It returns an error to be compatible
-// with a cio interface but at this time should only return nil in practice.
-func (task *Task) CancelIO() error {
-	close(task.cancelCh)
-	return nil
-}
-
 // StartStdioProxy sets up stdin, stdout and stderr streams for the task. It returns a channel that will be closed once
 // *initialization* of the streams has completed (or, if there was an error, the error will be published and then the chan
 // will be closed). Initialization consists of opening the FIFOs and establishing the vsock connections. After initialization,
@@ -222,8 +218,11 @@ func (task *Task) CancelIO() error {
 //
 // inputDirection is the IODirection that should be used for *stdin*. Stdout/stderr will, by definition, go in the opposite direction.
 // vsockConnector is a callback for establishing a vsock connection.
-func (task *Task) StartStdioProxy(ctx context.Context, inputDirection IODirection, vsockConnector VSockConnector) <-chan error {
-	logger := log.G(ctx).WithField("id", task.ID)
+//
+// initializeCtx is used exclusively during initialization, not during the actual IO copying. To stop IO copying, the task should be
+// removed from the TaskManager it was added to.
+func (task *Task) StartStdioProxy(initializeCtx context.Context, inputDirection IODirection, vsockConnector VSockConnector) <-chan error {
+	initializeCtx, initializeCancel := context.WithCancel(initializeCtx)
 	returnCh := make(chan error)
 
 	var wg sync.WaitGroup
@@ -234,61 +233,65 @@ func (task *Task) StartStdioProxy(ctx context.Context, inputDirection IODirectio
 	if task.fifoSet.Stdin != "" {
 		wg.Add(1)
 		go func() {
-			errCh <- task.proxyIO(ctx, task.fifoSet.Stdin, task.extraData.StdinPort, inputDirection, vsockConnector)
+			errCh <- task.proxyIO(initializeCtx, task.fifoSet.Stdin, task.extraData.StdinPort, inputDirection, vsockConnector)
 			wg.Done()
 		}()
 	} else {
-		logger.Info("skipping proxy io for unset stdin")
+		task.logger.Info("skipping proxy io for unset stdin")
 	}
 
 	if task.fifoSet.Stdout != "" {
 		wg.Add(1)
 		go func() {
-			errCh <- task.proxyIO(ctx, task.fifoSet.Stdout, task.extraData.StdoutPort, inputDirection.opposite(), vsockConnector)
+			errCh <- task.proxyIO(initializeCtx, task.fifoSet.Stdout, task.extraData.StdoutPort, inputDirection.opposite(), vsockConnector)
 			wg.Done()
 		}()
 	} else {
-		logger.Info("skipping proxy io for unset stdout")
+		task.logger.Info("skipping proxy io for unset stdout")
 	}
 
 	if task.fifoSet.Stderr != "" {
 		wg.Add(1)
 		go func() {
-			errCh <- task.proxyIO(ctx, task.fifoSet.Stderr, task.extraData.StderrPort, inputDirection.opposite(), vsockConnector)
+			errCh <- task.proxyIO(initializeCtx, task.fifoSet.Stderr, task.extraData.StderrPort, inputDirection.opposite(), vsockConnector)
 			wg.Done()
 		}()
 	} else {
-		logger.Info("skipping proxy io for unset stderr")
+		task.logger.Info("skipping proxy io for unset stderr")
 	}
 
-	// Once the streams have been initialized, gather any errors they returned and let the caller know by publishing
-	// on the return chan
+	// once each stream is initialized, close the error chan
 	go func() {
+		defer close(errCh)
 		wg.Wait()
-		close(errCh)
+	}()
 
-		var finalError *multierror.Error
+	// Once each stream has been initialized, get any error it returned and let the caller know by publishing
+	// on the return chan
+	go func() {
+		defer close(returnCh)
 		for err := range errCh {
-			finalError = multierror.Append(finalError, err)
+			if err != nil {
+				initializeCancel()
+				returnCh <- err
+				return
+			}
 		}
-
-		returnCh <- finalError.ErrorOrNil()
-		close(returnCh)
 	}()
 
 	return returnCh
 }
 
-func (task *Task) proxyIO(ctx context.Context, fifoPath string, port uint32, ioDirection IODirection, vsockConnector VSockConnector) error {
-	logger := log.G(ctx).WithField("path", fifoPath).WithField("port", port)
+func (task *Task) proxyIO(initializeCtx context.Context, fifoPath string, port uint32, ioDirection IODirection, vsockConnector VSockConnector) error {
+	logger := task.logger.WithField("path", fifoPath).WithField("port", port)
 
-	fifoFile, err := fifo.OpenFifo(ctx, fifoPath, syscall.O_RDWR|syscall.O_CREAT|syscall.O_NONBLOCK, 0700)
+	fifoFile, err := fifo.OpenFifo(initializeCtx, fifoPath, syscall.O_RDWR|syscall.O_CREAT|syscall.O_NONBLOCK, 0700)
 	if err != nil {
 		logger.WithError(err).Error("failed to open fifo")
 		return err
 	}
 
-	sock, err := vsockConnector(ctx, port)
+	sock, err := vsockConnector(initializeCtx, port)
 	if err != nil {
 		fifoFile.Close()
 		logger.WithError(err).Error("failed to establish vsock connection")
@@ -305,7 +308,7 @@ func (task *Task) proxyIO(ctx context.Context, fifoPath string, port uint32, ioD
 		switch ioDirection {
 		case FIFOtoVSock:
 			go func() {
-				<-task.cancelCh
+				<-task.doneCh
 				// If we are copying from FIFO to vsock, be sure to first close just the FIFO,
 				// allowing any buffered data to be flushed into the vsock first.
 				fifoFile.Close()
@@ -316,7 +319,7 @@ func (task *Task) proxyIO(ctx context.Context, fifoPath string, port uint32, ioD
 
 		case VSockToFIFO:
 			go func() {
-				<-task.cancelCh
+				<-task.doneCh
 				// If we are copying from vsock to FIFO, be sure to first close just the vsock,
 				// allowing any buffered data to be flushed into the FIFO first.
 				sock.Close()
diff --git a/runtime/service.go b/runtime/service.go
index b89f363..dcf8b0b 100644
--- a/runtime/service.go
+++ b/runtime/service.go
@@ -88,7 +88,27 @@ type service struct {
 	publish       shim.Publisher
 	eventExchange *exchange.Exchange
 	namespace     string
-	shutdown      func()
+
+	logger *logrus.Entry
+
+	// Normally, it's ill-advised to store a context object in a struct. However,
+	// in this case there appears to be little choice. Containerd provides the
+	// context and cancel from their shim setup code to our NewService function.
+	// It is expected to live for the lifetime of the shim process and canceled
+	// when the shim is shutting down.
+	//
+	// The problem is that sometimes service methods, such as CreateVM, require
+	// a context scoped to the lifetime of the shim process but they are only provided
+	// a context scoped to the lifetime of the individual request, not the shimCtx.
+	//
+	// shimCtx should thus only be used by service methods that need to provide
+	// a context that will be canceled only when the shim is shutting down and
+	// cleanup should happen.
+	//
+	// This approach is also taken by containerd's current reference runc shim
+	// v2 implementation
+	shimCtx    context.Context
+	shimCancel func()
 
 	vmID   string
 	config *Config
@@ -103,8 +123,8 @@ type service struct {
 	vsockPortMu      sync.Mutex
 }
 
-func shimOpts(ctx context.Context) (*shim.Opts, error) {
-	opts, ok := ctx.Value(shim.OptsKey{}).(shim.Opts)
+func shimOpts(shimCtx context.Context) (*shim.Opts, error) {
+	opts, ok := shimCtx.Value(shim.OptsKey{}).(shim.Opts)
 	if !ok {
 		return nil, errors.New("failed to parse containerd shim opts from context")
 	}
@@ -113,7 +133,7 @@ func shimOpts(ctx context.Context) (*shim.Opts, error) {
 }
 
 // NewService creates new runtime shim.
-func NewService(ctx context.Context, id string, remotePublisher shim.Publisher, shutdown func()) (shim.Shim, error) {
+func NewService(shimCtx context.Context, id string, remotePublisher shim.Publisher, shimCancel func()) (shim.Shim, error) {
 	server, err := newServer()
 	if err != nil {
 		return nil, err
@@ -125,7 +145,7 @@ func NewService(ctx context.Context, id string, remotePublisher shim.Publisher,
 	}
 
 	if !config.Debug {
-		opts, err := shimOpts(ctx)
+		opts, err := shimOpts(shimCtx)
 		if err != nil {
 			return nil, err
 		}
@@ -133,7 +153,7 @@ func NewService(ctx context.Context, id string, remotePublisher shim.Publisher,
 		config.Debug = opts.Debug
 	}
 
-	namespace, ok := namespaces.Namespace(ctx)
+	namespace, ok := namespaces.Namespace(shimCtx)
 	if !ok {
 		namespace = namespaces.Default
 	}
@@ -146,19 +166,22 @@ func NewService(ctx context.Context, id string, remotePublisher shim.Publisher,
 	// the containerd v2 runtime model only provides a shim with a publisher, not a forwarder,
 	// so we have to republish for now.
 	go func() {
-		if err := <-eventbridge.Republish(ctx, eventExchange, remotePublisher); err != nil && err != context.Canceled {
-			log.G(ctx).WithError(err).Error("error while republishing events")
+		if err := <-eventbridge.Republish(shimCtx, eventExchange, remotePublisher); err != nil && err != context.Canceled {
+			log.G(shimCtx).WithError(err).Error("error while republishing events")
 		}
 	}()
 
 	s := &service{
-		taskManager: vm.NewTaskManager(),
+		taskManager: vm.NewTaskManager(log.G(shimCtx)),
 
 		server:        server,
 		publish:       remotePublisher,
 		eventExchange: eventExchange,
 		namespace:     namespace,
-		shutdown:      shutdown,
+
+		logger:     log.G(shimCtx),
+		shimCtx:    shimCtx,
+		shimCancel: shimCancel,
 
 		vmID:   os.Getenv(vmIDEnvVarKey),
 		config: config,
@@ -173,25 +196,14 @@ func (s *service) vmDir() vm.Dir {
 	return vm.Dir(filepath.Join(varRunDir, s.namespace, s.vmID))
 }
 
-// shimSocketAddress is the abstract unix socket path at which our shim serves its API
-// It is unique per-VM and containerd namespace.
-func (s *service) shimSocketAddress(ctx context.Context) (string, error) {
-	return shim.SocketAddress(ctx, s.vmID)
-}
-
-func (s *service) newShim(ctx context.Context, containerdBinary, containerdAddress string, shimSocket *net.UnixListener) (*exec.Cmd, error) {
-	ns, err := namespaces.NamespaceRequired(ctx)
-	if err != nil {
-		return nil, err
-	}
-
+func (s *service) newShim(namespace string, containerdBinary, containerdAddress string, shimSocket *net.UnixListener) (*exec.Cmd, error) {
 	self, err := os.Executable()
 	if err != nil {
 		return nil, err
 	}
 
 	args := []string{
-		"-namespace", ns,
+		"-namespace", namespace,
 		"-address", containerdAddress,
 		"-publish-binary", containerdBinary,
 	}
@@ -224,14 +236,14 @@ func (s *service) newShim(ctx context.Context, containerdBinary, containerdAddre
 
 	// make sure to wait after start
 	go func() {
-		log.G(ctx).WithField("vmID", s.vmID).Debug("waiting on shim process")
+		s.logger.WithField("vmID", s.vmID).Debug("waiting on shim process")
 		waitErr := cmd.Wait()
-		log.G(ctx).WithError(waitErr).WithField("vmID", s.vmID).Debug("completed waiting on shim process")
+		s.logger.WithError(waitErr).WithField("vmID", s.vmID).Debug("completed waiting on shim process")
 	}()
 
 	err = shim.SetScore(cmd.Process.Pid)
 	if err != nil {
-		log.G(ctx).WithError(err).WithField("vmID", s.vmID).Error("failed to set OOM score on shim")
+		s.logger.WithError(err).WithField("vmID", s.vmID).Error("failed to set OOM score on shim")
 		return nil, errors.Wrap(err, "failed to set OOM Score on shim")
 	}
 
@@ -242,8 +254,13 @@ func isEADDRINUSE(err error) bool {
 	return err != nil && strings.Contains(err.Error(), "address already in use")
 }
 
-func (s *service) StartShim(ctx context.Context, containerID, containerdBinary, containerdAddress string) (string, error) {
-	log.G(ctx).WithField("id", containerID).Debug("StartShim")
+func (s *service) StartShim(shimCtx context.Context, containerID, containerdBinary, containerdAddress string) (string, error) {
+	s.logger.WithField("id", containerID).Debug("StartShim")
+
+	namespace, err := namespaces.NamespaceRequired(shimCtx)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to extract namespace from context")
+	}
 
 	// If we are running a shim start routine, we can safely assume our current working
 	// directory is the bundle directory
@@ -278,7 +295,7 @@ func (s *service) StartShim(ctx context.Context, containerID, containerdBinary,
 	// one. This is in line with the approach used by containerd's reference runC shim v2
 	// implementation (which is also designed to manage multiple containers from a single shim
 	// process)
-	shimSocketAddress, err := s.shimSocketAddress(ctx)
+	shimSocketAddress, err := shim.SocketAddress(shimCtx, s.vmID)
 	if err != nil {
 		return "", err
 	}
@@ -303,7 +320,7 @@ func (s *service) StartShim(ctx context.Context, containerID, containerdBinary,
 
 	// If we're here, there is no pre-existing shim for this VMID, so we spawn a new one
 	defer func() {
-		log.G(ctx).WithField("id", containerID).Debug("closing shim socket")
+		s.logger.WithField("id", containerID).Debug("closing shim socket")
 		shimSocket.Close()
 	}()
 
@@ -317,22 +334,22 @@ func (s *service) StartShim(ctx context.Context, containerID, containerdBinary,
 		return "", err
 	}
 
-	cmd, err := s.newShim(ctx, containerdBinary, containerdAddress, shimSocket)
+	cmd, err := s.newShim(namespace, containerdBinary, containerdAddress, shimSocket)
 	if err != nil {
-		log.G(ctx).WithError(err).WithField("id", containerID).Error("Failed to start new shim")
+		s.logger.WithError(err).WithField("id", containerID).Error("Failed to start new shim")
 		return "", err
 	}
 
 	defer func() {
 		if err != nil {
-			log.G(ctx).WithError(err).WithField("id", containerID).Error("killing shim process after error")
+			s.logger.WithError(err).WithField("id", containerID).Error("killing shim process after error")
 			cmd.Process.Kill()
 		}
 	}()
 
 	err = s.vmDir().WriteAddress(shimSocketAddress)
 	if err != nil {
-		log.G(ctx).WithError(err).WithField("id", containerID).Error("failed to write address")
+		s.logger.WithError(err).WithField("id", containerID).Error("failed to write address")
 		return "", err
 	}
 
@@ -355,7 +372,7 @@ func logPanicAndDie(logger *logrus.Entry) {
 	}
 }
 
-func parseCreateTaskOpts(ctx context.Context, opts *ptypes.Any) (*proto.FirecrackerConfig, *ptypes.Any, error) {
+func parseCreateTaskOpts(opts *ptypes.Any) (*proto.FirecrackerConfig, *ptypes.Any, error) {
 	cfg, err := typeurl.UnmarshalAny(opts)
 	if err != nil {
 		return nil, nil, errors.Wrapf(err, "unmarshaling task create request options")
@@ -419,11 +436,11 @@ func (s *service) nextVSockPort() uint32 {
 	return port
 }
 
-func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest) (*taskAPI.CreateTaskResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTaskRequest) (*taskAPI.CreateTaskResponse, error) {
+	logger := log.G(requestCtx).WithField("id", request.ID)
+	defer logPanicAndDie(logger)
 
-	log.G(ctx).WithFields(logrus.Fields{
-		"id":         request.ID,
+	logger.WithFields(logrus.Fields{
 		"bundle":     request.Bundle,
 		"terminal":   request.Terminal,
 		"stdin":      request.Stdin,
@@ -437,10 +454,9 @@ func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest
 	var runcOpts *ptypes.Any
 
 	if request.Options != nil {
-		firecrackerConfig, runcOpts, err = parseCreateTaskOpts(ctx, request.Options)
+		firecrackerConfig, runcOpts, err = parseCreateTaskOpts(request.Options)
 		if err != nil {
-			log.G(ctx).WithFields(logrus.Fields{
-				"id":      request.ID,
+			logger.WithFields(logrus.Fields{
 				"options": request.Options,
 				"error":   err,
 			}).Error("failed to unmarshal task create request options")
@@ -450,11 +466,11 @@ func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest
 
 	s.startVMMutex.Lock()
 	if !s.agentStarted {
-		log.G(ctx).WithField("id", request.ID).Debug("calling startVM")
-		client, err := s.startVM(ctx, request, firecrackerConfig)
+		logger.Debug("calling startVM")
+		client, err := s.startVM(requestCtx, request, firecrackerConfig)
 		if err != nil {
 			s.startVMMutex.Unlock()
-			log.G(ctx).WithError(err).WithField("id", request.ID).Error("failed to start VM")
+			logger.WithError(err).Error("failed to start VM")
 			return nil, err
 		}
 
@@ -463,7 +479,7 @@ func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest
 	}
 	s.startVMMutex.Unlock()
 
-	log.G(ctx).WithField("id", request.ID).Info("creating task")
+	logger.Info("creating task")
 
 	bundleDir := bundle.Dir(request.Bundle)
 	extraData, err := s.generateExtraData(bundleDir, runcOpts)
@@ -471,25 +487,27 @@ func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest
 		return nil, err
 	}
 
-	task, err := s.taskManager.AddTask(ctx, request.ID, s.agentClient, bundleDir, extraData, cio.NewFIFOSet(cio.Config{
+	taskCtx, taskCancel := context.WithCancel(s.shimCtx)
+	task, err := s.taskManager.AddTask(request.ID, s.agentClient, bundleDir, extraData, cio.NewFIFOSet(cio.Config{
 		Stdin:    request.Stdin,
 		Stdout:   request.Stdout,
 		Stderr:   request.Stderr,
 		Terminal: request.Terminal,
-	}, nil))
+	}, nil), taskCtx.Done(), taskCancel)
 	if err != nil {
 		return nil, err
 	}
 
 	defer func() {
 		if err != nil {
-			s.taskManager.Remove(ctx, request.ID)
+			s.taskManager.Remove(request.ID)
 		}
 	}()
 
 	// Begin initializing stdio, but don't block on the initialization so we can send the Create
-	// call (which will allow the stdio initialization to complete)
-	stdioReadyCh := task.StartStdioProxy(ctx, vm.FIFOtoVSock, cidDialer(s.machineCID))
+	// call (which will allow the stdio initialization to complete).
+	// Use requestCtx as the ctx here is only used for initializing stdio, not the actual copying of io
+	stdioReadyCh := task.StartStdioProxy(requestCtx, vm.FIFOtoVSock, cidDialer(s.machineCID))
 
 	// Override the original request options with ExtraData needed by the VM Agent before sending it off
 	request.Options, err = task.MarshalExtraData()
@@ -497,9 +515,9 @@ func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest
 		return nil, err
 	}
 
-	resp, err := task.Create(ctx, request)
+	resp, err := task.Create(requestCtx, request)
 	if err != nil {
-		log.G(ctx).WithError(err).Error("create failed")
+		logger.WithError(err).Error("create failed")
 		return nil, err
 	}
 
@@ -509,20 +527,20 @@ func (s *service) Create(ctx context.Context, request *taskAPI.CreateTaskRequest
 		return nil, err
 	}
 
-	log.G(ctx).WithField("id", request.ID).WithField("pid", resp.Pid).Info("successfully created task")
+	logger.WithField("pid", resp.Pid).Info("successfully created task")
 	return resp, nil
 }
 
-func (s *service) Start(ctx context.Context, req *taskAPI.StartRequest) (*taskAPI.StartResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Start(requestCtx context.Context, req *taskAPI.StartRequest) (*taskAPI.StartResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("start")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("start")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Start(ctx, req)
+	resp, err := task.Start(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -530,37 +548,37 @@ func (s *service) Start(ctx context.Context, req *taskAPI.StartRequest) (*taskAP
 	return resp, nil
 }
 
-func (s *service) Delete(ctx context.Context, req *taskAPI.DeleteRequest) (*taskAPI.DeleteResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Delete(requestCtx context.Context, req *taskAPI.DeleteRequest) (*taskAPI.DeleteResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("delete")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("delete")
 
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Delete(ctx, req)
+	resp, err := task.Delete(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
 
-	s.taskManager.Remove(ctx, req.ID)
+	s.taskManager.Remove(req.ID)
 
 	return resp, nil
 }
 
 // Exec an additional process inside the container
-func (s *service) Exec(ctx context.Context, req *taskAPI.ExecProcessRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Exec(requestCtx context.Context, req *taskAPI.ExecProcessRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("exec")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("exec")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Exec(ctx, req)
+	resp, err := task.Exec(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -569,16 +587,16 @@ func (s *service) Exec(ctx context.Context, req *taskAPI.ExecProcessRequest) (*p
 }
 
 // ResizePty of a process
-func (s *service) ResizePty(ctx context.Context, req *taskAPI.ResizePtyRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) ResizePty(requestCtx context.Context, req *taskAPI.ResizePtyRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("resize_pty")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("resize_pty")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.ResizePty(ctx, req)
+	resp, err := task.ResizePty(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -587,16 +605,16 @@ func (s *service) ResizePty(ctx context.Context, req *taskAPI.ResizePtyRequest)
 }
 
 // State returns runtime state information for a process
-func (s *service) State(ctx context.Context, req *taskAPI.StateRequest) (*taskAPI.StateResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) State(requestCtx context.Context, req *taskAPI.StateRequest) (*taskAPI.StateResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("state")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("state")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.State(ctx, req)
+	resp, err := task.State(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -605,16 +623,16 @@ func (s *service) State(ctx context.Context, req *taskAPI.StateRequest) (*taskAP
 }
 
 // Pause the container
-func (s *service) Pause(ctx context.Context, req *taskAPI.PauseRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Pause(requestCtx context.Context, req *taskAPI.PauseRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("pause")
+	log.G(requestCtx).WithField("id", req.ID).Debug("pause")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Pause(ctx, req)
+	resp, err := task.Pause(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -623,16 +641,16 @@ func (s *service) Pause(ctx context.Context, req *taskAPI.PauseRequest) (*ptypes
 }
 
 // Resume the container
-func (s *service) Resume(ctx context.Context, req *taskAPI.ResumeRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Resume(requestCtx context.Context, req *taskAPI.ResumeRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("resume")
+	log.G(requestCtx).WithField("id", req.ID).Debug("resume")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Resume(ctx, req)
+	resp, err := task.Resume(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -641,16 +659,16 @@ func (s *service) Resume(ctx context.Context, req *taskAPI.ResumeRequest) (*ptyp
 }
 
 // Kill a process with the provided signal
-func (s *service) Kill(ctx context.Context, req *taskAPI.KillRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Kill(requestCtx context.Context, req *taskAPI.KillRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("kill")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("kill")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Kill(ctx, req)
+	resp, err := task.Kill(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -658,16 +676,16 @@ func (s *service) Kill(ctx context.Context, req *taskAPI.KillRequest) (*ptypes.E
 }
 
 // Pids returns all pids inside the container
-func (s *service) Pids(ctx context.Context, req *taskAPI.PidsRequest) (*taskAPI.PidsResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Pids(requestCtx context.Context, req *taskAPI.PidsRequest) (*taskAPI.PidsResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("pids")
+	log.G(requestCtx).WithField("id", req.ID).Debug("pids")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Pids(ctx, req)
+	resp, err := task.Pids(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -676,16 +694,16 @@ func (s *service) Pids(ctx context.Context, req *taskAPI.PidsRequest) (*taskAPI.
 }
 
 // CloseIO of a process
-func (s *service) CloseIO(ctx context.Context, req *taskAPI.CloseIORequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) CloseIO(requestCtx context.Context, req *taskAPI.CloseIORequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("close_io")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("close_io")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.CloseIO(ctx, req)
+	resp, err := task.CloseIO(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -694,16 +712,16 @@ func (s *service) CloseIO(ctx context.Context, req *taskAPI.CloseIORequest) (*pt
 }
 
 // Checkpoint the container
-func (s *service) Checkpoint(ctx context.Context, req *taskAPI.CheckpointTaskRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Checkpoint(requestCtx context.Context, req *taskAPI.CheckpointTaskRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "path": req.Path}).Info("checkpoint")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "path": req.Path}).Info("checkpoint")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Checkpoint(ctx, req)
+	resp, err := task.Checkpoint(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -712,16 +730,16 @@ func (s *service) Checkpoint(ctx context.Context, req *taskAPI.CheckpointTaskReq
 }
 
 // Connect returns shim information such as the shim's pid
-func (s *service) Connect(ctx context.Context, req *taskAPI.ConnectRequest) (*taskAPI.ConnectResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Connect(requestCtx context.Context, req *taskAPI.ConnectRequest) (*taskAPI.ConnectResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("connect")
+	log.G(requestCtx).WithField("id", req.ID).Debug("connect")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Connect(ctx, req)
+	resp, err := task.Connect(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -729,54 +747,48 @@ func (s *service) Connect(ctx context.Context, req *taskAPI.ConnectRequest) (*ta
 	return resp, nil
 }
 
-func (s *service) Shutdown(ctx context.Context, req *taskAPI.ShutdownRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Shutdown(requestCtx context.Context, req *taskAPI.ShutdownRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "now": req.Now}).Debug("shutdown")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "now": req.Now}).Debug("shutdown")
 
 	// If we are still managing containers, don't shutdown
 	if s.taskManager.TaskCount() > 0 {
 		return &ptypes.Empty{}, nil
 	}
 
-	_, err := s.agentClient.Shutdown(ctx, req)
+	_, err := s.agentClient.Shutdown(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
 
-	log.G(ctx).Debug("stopping VM")
+	log.G(requestCtx).Debug("stopping VM")
 	if err := s.stopVM(); err != nil {
-		log.G(ctx).WithError(err).Error("failed to stop VM")
+		log.G(requestCtx).WithError(err).Error("failed to stop VM")
 		return nil, err
 	}
 
 	err = os.RemoveAll(s.vmDir().RootPath())
 	if err != nil {
-		log.G(ctx).WithField("path", s.vmDir().RootPath()).WithError(err).Error("failed to remove VM dir during shutdown")
+		log.G(requestCtx).WithField("path", s.vmDir().RootPath()).WithError(err).Error("failed to remove VM dir during shutdown")
 	}
 
-	// Exit to avoid 'zombie' shim processes
-	defer func() {
-		if s.shutdown != nil {
-			s.shutdown()
-		}
-		os.Exit(0)
-	}()
+	defer s.shimCancel()
 
-	log.G(ctx).Debug("stopping runtime")
+	log.G(requestCtx).Debug("stopping runtime")
 	return &ptypes.Empty{}, nil
 }
 
-func (s *service) Stats(ctx context.Context, req *taskAPI.StatsRequest) (*taskAPI.StatsResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Stats(requestCtx context.Context, req *taskAPI.StatsRequest) (*taskAPI.StatsResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("stats")
+	log.G(requestCtx).WithField("id", req.ID).Debug("stats")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Stats(ctx, req)
+	resp, err := task.Stats(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -785,16 +797,16 @@ func (s *service) Stats(ctx context.Context, req *taskAPI.StatsRequest) (*taskAP
 }
 
 // Update a running container
-func (s *service) Update(ctx context.Context, req *taskAPI.UpdateTaskRequest) (*ptypes.Empty, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Update(requestCtx context.Context, req *taskAPI.UpdateTaskRequest) (*ptypes.Empty, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithField("id", req.ID).Debug("update")
+	log.G(requestCtx).WithField("id", req.ID).Debug("update")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Update(ctx, req)
+	resp, err := task.Update(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -803,16 +815,16 @@ func (s *service) Update(ctx context.Context, req *taskAPI.UpdateTaskRequest) (*
 }
 
 // Wait for a process to exit
-func (s *service) Wait(ctx context.Context, req *taskAPI.WaitRequest) (*taskAPI.WaitResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Wait(requestCtx context.Context, req *taskAPI.WaitRequest) (*taskAPI.WaitResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("wait")
+	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "exec_id": req.ExecID}).Debug("wait")
 	task, err := s.taskManager.Task(req.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := task.Wait(ctx, req)
+	resp, err := task.Wait(requestCtx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -820,10 +832,10 @@ func (s *service) Wait(ctx context.Context, req *taskAPI.WaitRequest) (*taskAPI.
 	return resp, nil
 }
 
-func (s *service) Cleanup(ctx context.Context) (*taskAPI.DeleteResponse, error) {
-	defer logPanicAndDie(log.G(ctx))
+func (s *service) Cleanup(requestCtx context.Context) (*taskAPI.DeleteResponse, error) {
+	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(ctx).Debug("cleanup")
+	log.G(requestCtx).Debug("cleanup")
 	// Destroy VM/etc here?
 	// copied from runcs impl, nothing to cleanup atm
 	return &taskAPI.DeleteResponse{
@@ -836,7 +848,7 @@ func newServer() (*ttrpc.Server, error) {
 	return ttrpc.NewServer(ttrpc.WithServerHandshaker(ttrpc.UnixSocketRequireSameUser()))
 }
 
-func dialVsock(ctx context.Context, contextID, port uint32) (net.Conn, error) {
+func dialVsock(requestCtx context.Context, contextID, port uint32) (net.Conn, error) {
 	// VM should start within 200ms, vsock dial will make retries at 100ms, 200ms, 400ms, 800ms, 1.6s, 3.2s, 6.4s
 	const (
 		retryCount      = 7
@@ -846,27 +858,33 @@ func dialVsock(ctx context.Context, contextID, port uint32) (net.Conn, error) {
 
 	var lastErr error
 	var currentDelay = initialDelay
+
 	for i := 1; i <= retryCount; i++ {
-		conn, err := vsock.Dial(contextID, port)
-		if err == nil {
-			log.G(ctx).WithField("connection", conn).Debug("Dial succeeded")
-			return conn, nil
-		}
+		select {
+		case <-requestCtx.Done():
+			return nil, requestCtx.Err()
+		default:
+			conn, err := vsock.Dial(contextID, port)
+			if err == nil {
+				log.G(requestCtx).WithField("connection", conn).Debug("Dial succeeded")
+				return conn, nil
+			}
 
-		log.G(ctx).WithError(err).Warnf("vsock dial failed (attempt %d of %d), will retry in %s", i, retryCount, currentDelay)
-		time.Sleep(currentDelay)
+			log.G(requestCtx).WithError(err).Warnf("vsock dial failed (attempt %d of %d), will retry in %s", i, retryCount, currentDelay)
+			time.Sleep(currentDelay)
 
-		lastErr = err
-		currentDelay *= delayMultiplier
+			lastErr = err
+			currentDelay *= delayMultiplier
+		}
 	}
 
-	log.G(ctx).WithError(lastErr).WithFields(logrus.Fields{"context_id": contextID, "port": port}).Error("vsock dial failed")
+	log.G(requestCtx).WithError(lastErr).WithFields(logrus.Fields{"context_id": contextID, "port": port}).Error("vsock dial failed")
 	return nil, lastErr
 }
 
 func cidDialer(cid uint32) vm.VSockConnector {
-	return func(ctx context.Context, port uint32) (net.Conn, error) {
-		return dialVsock(ctx, cid, port)
+	return func(requestCtx context.Context, port uint32) (net.Conn, error) {
+		return dialVsock(requestCtx, cid, port)
 	}
 }
 
@@ -874,7 +892,7 @@ func cidDialer(cid uint32) vm.VSockConnector {
 // It uses VHOST_VSOCK_SET_GUEST_CID ioctl which allows some CID ranges to be statically reserved in advance.
 // The ioctl fails with EADDRINUSE if cid is already taken and with EINVAL if the CID is invalid.
 // Taken from https://bugzilla.redhat.com/show_bug.cgi?id=1291851
-func findNextAvailableVsockCID(ctx context.Context) (uint32, error) {
+func findNextAvailableVsockCID(requestCtx context.Context) (uint32, error) {
 	const (
 		// Corresponds to VHOST_VSOCK_SET_GUEST_CID in vhost.h
 		ioctlVsockSetGuestCID = uintptr(0x4008AF60)
@@ -897,8 +915,8 @@ func findNextAvailableVsockCID(ctx context.Context) (uint32, error) {
 	for n := 0; n < maxCID-minCID; n++ {
 		cid := minCID + ((start + n) % (maxCID - minCID))
 		select {
-		case <-ctx.Done():
-			return 0, ctx.Err()
+		case <-requestCtx.Done():
+			return 0, requestCtx.Err()
 		default:
 			_, _, err = sysCall(
 				unix.SYS_IOCTL,
@@ -923,13 +941,13 @@ func findNextAvailableVsockCID(ctx context.Context) (uint32, error) {
 }
 
 // TODO: replace startVM with calls to the FC-control plugin
-func (s *service) startVM(ctx context.Context,
+func (s *service) startVM(requestCtx context.Context,
 	request *taskAPI.CreateTaskRequest,
 	vmConfig *proto.FirecrackerConfig,
 ) (taskAPI.TaskService, error) {
-	log.G(ctx).Info("starting VM")
+	log.G(requestCtx).Info("starting VM")
 
-	cid, err := findNextAvailableVsockCID(ctx)
+	cid, err := findNextAvailableVsockCID(requestCtx)
 	if err != nil {
 		return nil, err
 	}
@@ -973,33 +991,31 @@ func (s *service) startVM(ctx context.Context,
 	cmd := firecracker.VMCommandBuilder{}.
 		WithBin(s.config.FirecrackerBinaryPath).
 		WithSocketPath(s.vmDir().FirecrackerSockPath()).
-		Build(context.Background())
+		Build(s.shimCtx) // use shimCtx to ensure the cmd is only killed once the shim shutsdown
 	machineOpts := []firecracker.Opt{
 		firecracker.WithProcessRunner(cmd),
-		firecracker.WithLogger(log.G(ctx)),
+		firecracker.WithLogger(s.logger),
 	}
 
-	vmmCtx, vmmCancel := context.WithCancel(context.Background())
-	defer vmmCancel()
-	s.machine, err = firecracker.NewMachine(vmmCtx, cfg, machineOpts...)
+	s.machine, err = firecracker.NewMachine(requestCtx, cfg, machineOpts...)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create firecracker config")
 	}
 	s.machineCID = cid
 
-	log.G(ctx).Info("starting instance")
-	if err := s.machine.Start(vmmCtx); err != nil {
+	log.G(requestCtx).Info("starting instance")
+	if err := s.machine.Start(requestCtx); err != nil {
 		return nil, errors.Wrap(err, "failed to start firecracker VM")
 	}
 
-	log.G(ctx).Info("calling agent")
-	conn, err := dialVsock(ctx, cid, defaultVsockPort)
+	log.G(requestCtx).Info("calling agent")
+	conn, err := dialVsock(requestCtx, cid, defaultVsockPort)
 	if err != nil {
 		s.stopVM()
 		return nil, errors.Wrap(err, "failed to dial vsock")
 	}
 
-	log.G(ctx).Info("creating clients")
+	log.G(requestCtx).Info("creating clients")
 	rpcClient := ttrpc.NewClient(conn, ttrpc.WithOnClose(func() { _ = conn.Close() }))
 	apiClient := taskAPI.NewTaskClient(rpcClient)
 	eventBridgeClient := eventbridge.NewGetterClient(rpcClient)
@@ -1008,12 +1024,12 @@ func (s *service) startVM(ctx context.Context,
 		// Connect the agent's event exchange to our own own event exchange
 		// using the eventbridge client. All events that are published on the
 		// agent's exchange will also be published on our own
-		if err := <-eventbridge.Attach(context.Background(), eventBridgeClient, s.eventExchange); err != nil && err != context.Canceled {
-			log.G(ctx).WithError(err).Error("error while forwarding events from VM agent")
+		if err := <-eventbridge.Attach(s.shimCtx, eventBridgeClient, s.eventExchange); err != nil && err != context.Canceled {
+			log.G(requestCtx).WithError(err).Error("error while forwarding events from VM agent")
 		}
 	}()
 
-	go s.monitorTaskExit(context.Background())
+	go s.monitorTaskExit()
 
 	return apiClient, nil
 }
@@ -1022,14 +1038,13 @@ func (s *service) stopVM() error {
 	return s.machine.StopVMM()
 }
 
-func (s *service) monitorTaskExit(ctx context.Context) {
-	logger := log.G(ctx)
-	exitEvents, exitEventErrs := s.eventExchange.Subscribe(ctx, fmt.Sprintf(`topic=="%s"`, runtime.TaskExitEventTopic))
+func (s *service) monitorTaskExit() {
+	exitEvents, exitEventErrs := s.eventExchange.Subscribe(s.shimCtx, fmt.Sprintf(`topic=="%s"`, runtime.TaskExitEventTopic))
 
 	var err error
 	defer func() {
 		if err != nil && err != context.Canceled {
-			logger.WithError(err).Error("error while waiting for task exit events")
+			s.logger.WithError(err).Error("error while waiting for task exit events")
 		}
 	}()
 
@@ -1038,35 +1053,35 @@ func (s *service) monitorTaskExit(ctx context.Context) {
 		case envelope := <-exitEvents:
 			unmarshaledEvent, err := typeurl.UnmarshalAny(envelope.Event)
 			if err != nil {
-				logger.WithError(err).Error("error unmarshaling event")
+				s.logger.WithError(err).Error("error unmarshaling event")
 				continue
 			}
 
 			switch event := unmarshaledEvent.(type) {
 			case *eventsAPI.TaskExit:
-				logger = logger.WithField("id", event.ContainerID)
+				logger := s.logger.WithField("id", event.ContainerID)
 				logger.Debug("received container exit event")
 
-				s.taskManager.Remove(ctx, event.ContainerID)
+				s.taskManager.Remove(event.ContainerID)
 
 				// If we have no more containers, shutdown. If we still have containers left,
 				// this will be a no-op
-				_, err = s.Shutdown(ctx, &taskAPI.ShutdownRequest{})
+				_, err = s.Shutdown(s.shimCtx, &taskAPI.ShutdownRequest{})
 				if err != nil {
 					logger.WithError(err).Fatal("failed to shutdown after container exit")
 				}
 
 			default:
-				logger.Error("unexpected non-exit event type published on exit event channel")
+				s.logger.Error("unexpected non-exit event type published on exit event channel")
 			}
 
 		case err = <-exitEventErrs:
 			if err != nil {
-				logger.WithError(err).Error("event error channel published to")
+				s.logger.WithError(err).Error("event error channel published to")
 			}
 
-		case <-ctx.Done():
-			err = ctx.Err()
+		case <-s.shimCtx.Done():
+			err = s.shimCtx.Err()
 			return
 
 		}


Test patch:
diff --git a/internal/vm/task_test.go b/internal/vm/task_test.go
index b20659b..4571d29 100644
--- a/internal/vm/task_test.go
+++ b/internal/vm/task_test.go
@@ -26,6 +26,7 @@ import (
 	"github.com/containerd/containerd/runtime/v2/task"
 	"github.com/firecracker-microvm/firecracker-containerd/internal/bundle"
 	"github.com/firecracker-microvm/firecracker-containerd/proto"
+	"github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
 )
 
@@ -70,11 +71,12 @@ type addTaskArgs struct {
 }
 
 func addTaskFromArgs(tm TaskManager, args addTaskArgs) (*Task, error) {
-	return tm.AddTask(context.TODO(), args.id, args.ts, args.bundleDir, args.extraData, args.fifoSet)
+	return tm.AddTask(args.id, args.ts, args.bundleDir, args.extraData, args.fifoSet, context.Background().Done(), func() {})
 }
 
 func TestTaskManager_AddRemoveTask(t *testing.T) {
-	tm := NewTaskManager()
+	logger, _ := test.NewNullLogger()
+	tm := NewTaskManager(logger.WithField("test", t.Name()))
 
 	taskAArgs := defaultMockTaskArgs("A")
 	taskBArgs := defaultMockTaskArgs("B")
@@ -123,7 +125,7 @@ func TestTaskManager_AddRemoveTask(t *testing.T) {
 	require.Error(t, err, "get non-existent task should fail")
 
 	// Removing a task should remove the task requested but have no effect on the others
-	tm.Remove(context.TODO(), taskBArgs.id)
+	tm.Remove(taskBArgs.id)
 	require.EqualValues(t, 2, tm.TaskCount(), "task count should return correct number of tasks after remove")
 
 	supposedlyTaskA, err = tm.Task(taskAArgs.id)
@@ -138,7 +140,7 @@ func TestTaskManager_AddRemoveTask(t *testing.T) {
 	require.Exactly(t, taskC, supposedlyTaskC, "get should return expected task")
 
 	// Remove all should remove all the tasks
-	tm.RemoveAll(context.TODO())
+	tm.RemoveAll()
 	require.EqualValues(t, 0, tm.TaskCount(), "task count should return correct number of tasks after remove all")
 
 	_, err = tm.Task(taskAArgs.id)
@@ -147,9 +149,9 @@ func TestTaskManager_AddRemoveTask(t *testing.T) {
 	_, err = tm.Task(taskCArgs.id)
 	require.Error(t, err, "get removed task should fail")
 
-	tm.RemoveAll(context.TODO())
+	tm.RemoveAll()
 	require.EqualValues(t, 0, tm.TaskCount(), "remove all on empty task manager should have no effect")
-	tm.Remove(context.TODO(), taskAArgs.id)
+	tm.Remove(taskAArgs.id)
 	require.EqualValues(t, 0, tm.TaskCount(), "remove on non-existent task should have no effect")
 
 	// Tasks should still be able to be added after removes
@@ -177,7 +179,8 @@ func TestTaskManager_StartStdioProxy_VSockToFIFO(t *testing.T) {
 }
 
 func testStdioProxy(t *testing.T, inputDirection IODirection) {
-	tm := NewTaskManager()
+	logger, _ := test.NewNullLogger()
+	tm := NewTaskManager(logger.WithField("test", t.Name()))
 
 	taskArgs := defaultMockTaskArgs("A")
 
diff --git a/runtime/service_test.go b/runtime/service_test.go
index 6896bac..05a64bc 100644
--- a/runtime/service_test.go
+++ b/runtime/service_test.go
@@ -63,7 +63,7 @@ func TestParseCreateTaskOptsParsesFirecrackerConfig(t *testing.T) {
 
 	protoFirecrackerConfig, err := typeurl.MarshalAny(inFirecrackerConfig)
 	require.NoError(t, err, "unable to marshal firecracker config proto message")
-	outFirecrackerConfig, outRuncOpts, err := parseCreateTaskOpts(context.TODO(), protoFirecrackerConfig)
+	outFirecrackerConfig, outRuncOpts, err := parseCreateTaskOpts(protoFirecrackerConfig)
 	require.NoError(t, err, "unable to parse firecracker config from proto message")
 	if outRuncOpts != nil {
 		// assert.Equal is insufficient here as the nil comparison fails for
@@ -82,7 +82,7 @@ func TestParseCreateTaskLeavesNonFirecrackerConfigAlong(t *testing.T) {
 	}
 	protoIn, err := typeurl.MarshalAny(in)
 	require.NoError(t, err, "unable to marshal proto message")
-	outFirecrackerConfig, outOpts, err := parseCreateTaskOpts(context.TODO(), protoIn)
+	outFirecrackerConfig, outOpts, err := parseCreateTaskOpts(protoIn)
 	require.NoError(t, err, "unable to parse firecracker config from proto message")
 	if outFirecrackerConfig != nil {
 		// assert.Equal is insufficient here as the nil comparison fails for