Fix mem leak in newHcsTask()

goroutine that waits on uvm exit in
newHcsTask() should also check for
container exit as containers can exit
before the uvm does. This will ensure
that we will not leak the goroutine.

Signed-off-by: Kirtana Ashok <[email protected]>

cmd/containerd-shim-runhcs-v1/task_hcs.go

@@ -252,7 +252,7 @@ func newHcsTask(
 	if parent != nil {
 		// We have a parent UVM. Listen for its exit and forcibly close this
 		// task. This is not expected but in the event of a UVM crash we need to
-		// handle this case.
+		// ensure the resources are cleaned up as expected.
 		go ht.waitForHostExit()
 	}
 
@@ -616,8 +616,10 @@ func (ht *hcsTask) waitInitExit() {
 	ht.close(ctx)
 }
 
-// waitForHostExit waits for the host virtual machine to exit. Once exited
-// forcibly exits all additional exec's in this task.
+// waitForHostExit waits for the host virtual machine to exit. Once exited,
+// it forcibly exits all additional execs in this task. Make sure to check
+// for container exit as well since the container could exit before
+// the UVM and leak this goroutine started for its task.
 //
 // This MUST be called via a goroutine to wait on a background thread.
 //
@@ -628,7 +630,7 @@ func (ht *hcsTask) waitForHostExit() {
 	defer span.End()
 	span.AddAttributes(trace.StringAttribute("tid", ht.id))
 
-	err := ht.host.WaitCtx(ctx)
+	err := ht.host.WaitForUvmOrContainerExit(ctx, ht.c)
 	if err != nil {
 		log.G(ctx).WithError(err).Error("failed to wait for host virtual machine exit")
 	} else {

internal/uvm/wait.go

@@ -9,9 +9,33 @@ import (
 
 	"github.com/sirupsen/logrus"
 
+	"github.com/Microsoft/hcsshim/internal/cow"
 	"github.com/Microsoft/hcsshim/internal/logfields"
 )
 
+// WaitForUvmOrContainerExit waits for the container `c` or its UVM
+// to exit. This is used to clean up hcs task and exec resources by
+// the caller.
+func (uvm *UtilityVM) WaitForUvmOrContainerExit(ctx context.Context, c cow.Container) (err error) {
+	select {
+	case <-c.WaitChannel():
+		return c.WaitError()
+	case <-uvm.hcsSystem.WaitChannel():
+		logrus.WithField(logfields.UVMID, uvm.id).Debug("uvm exited, waiting for output processing to complete")
+		var outputErr error
+		if uvm.outputProcessingDone != nil {
+			select {
+			case <-uvm.outputProcessingDone:
+			case <-ctx.Done():
+				outputErr = fmt.Errorf("failed to wait on uvm output processing: %w", ctx.Err())
+			}
+		}
+		return errors.Join(uvm.hcsSystem.WaitError(), outputErr)
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 // Wait waits synchronously for a utility VM to terminate.
 func (uvm *UtilityVM) Wait() error { return uvm.WaitCtx(context.Background()) }