Fix mem leak in newHcsTask()

newHcsTask() starts a goroutine for hyperv
containers to enable hcsTask and hcsExec
resource cleanup as expected. But when
containers exit before the uvm in the
normal case, this goroutine is leaked
as it is waiting on a block wait on the UVM.
This commit ensures that we wait for container
exit as well so that we do not leak the routine.

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()) }