1 files changed, 19 insertions, 15 deletions
diff --git a/test/benchmarks/harness/util.go b/test/benchmarks/harness/util.go
index 36abe1069..f7e569751 100644
--- a/test/benchmarks/harness/util.go
+++ b/test/benchmarks/harness/util.go
@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/docker/docker/api/types/mount"
+	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/test/dockerutil"
 	"gvisor.dev/gvisor/pkg/test/testutil"
 )
@@ -58,52 +59,55 @@ func DebugLog(b *testing.B, msg string, args ...interface{}) {
 	}
 }
 
+// FileSystemType represents a type container mount.
+type FileSystemType string
+
 const (
 	// BindFS indicates a bind mount should be created.
-	BindFS = "bindfs"
+	BindFS FileSystemType = "bindfs"
 	// TmpFS indicates a tmpfs mount should be created.
-	TmpFS = "tmpfs"
+	TmpFS FileSystemType = "tmpfs"
 	// RootFS indicates no mount should be created and the root mount should be used.
-	RootFS = "rootfs"
+	RootFS FileSystemType = "rootfs"
 )
 
 // MakeMount makes a mount and cleanup based on the requested type. Bind
 // and volume mounts are backed by a temp directory made with mktemp.
 // tmpfs mounts require no such backing and are just made.
 // rootfs mounts do not make a mount, but instead return a target direectory at root.
-// It is up to the caller to call the returned cleanup.
-func MakeMount(machine Machine, fsType string) ([]mount.Mount, string, func(), error) {
+// It is up to the caller to call Clean on the passed *cleanup.Cleanup
+func MakeMount(machine Machine, fsType FileSystemType, cu *cleanup.Cleanup) ([]mount.Mount, string, error) {
 	mounts := make([]mount.Mount, 0, 1)
+	target := "/data"
 	switch fsType {
 	case BindFS:
 		dir, err := machine.RunCommand("mktemp", "-d")
 		if err != nil {
-			return mounts, "", func() {}, fmt.Errorf("failed to create tempdir: %v", err)
+			return mounts, "", fmt.Errorf("failed to create tempdir: %v", err)
 		}
 		dir = strings.TrimSuffix(dir, "\n")
-
+		cu.Add(func() {
+			machine.RunCommand("rm", "-rf", dir)
+		})
 		out, err := machine.RunCommand("chmod", "777", dir)
 		if err != nil {
-			machine.RunCommand("rm", "-rf", dir)
-			return mounts, "", func() {}, fmt.Errorf("failed modify directory: %v %s", err, out)
+			return mounts, "", fmt.Errorf("failed modify directory: %v %s", err, out)
 		}
-		target := "/data"
 		mounts = append(mounts, mount.Mount{
 			Target: target,
 			Source: dir,
 			Type:   mount.TypeBind,
 		})
-		return mounts, target, func() { machine.RunCommand("rm", "-rf", dir) }, nil
+		return mounts, target, nil
 	case RootFS:
-		return mounts, "/", func() {}, nil
+		return mounts, target, nil
 	case TmpFS:
-		target := "/data"
 		mounts = append(mounts, mount.Mount{
 			Target: target,
 			Type:   mount.TypeTmpfs,
 		})
-		return mounts, target, func() {}, nil
+		return mounts, target, nil
 	default:
-		return mounts, "", func() {}, fmt.Errorf("illegal mount type not supported: %v", fsType)
+		return mounts, "", fmt.Errorf("illegal mount type not supported: %v", fsType)
 	}
 }