sweet/benchmarks/gvisor: use new diagnostics framework

Change-Id: I245edaf771c0d8698d4a28fcb064362d09911a4e
Cq-Include-Trybots: luci.golang.try:x_benchmarks-gotip-linux-amd64-longtest
Reviewed-on: https://go-review.googlesource.com/c/benchmarks/+/600067
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>
diff --git a/sweet/benchmarks/gvisor/common.go b/sweet/benchmarks/gvisor/common.go
index 7f18e16..65d1d27 100644
--- a/sweet/benchmarks/gvisor/common.go
+++ b/sweet/benchmarks/gvisor/common.go
@@ -9,6 +9,7 @@
 
 import (
 	"fmt"
+	"os"
 	"os/exec"
 	"path/filepath"
 	"syscall"
@@ -24,40 +25,44 @@
 	return filepath.Join(assetsDir, subBenchmark, "bin", platformDir, "workload")
 }
 
-func (c *config) profilePath(typ diagnostics.Type) string {
-	return filepath.Join(c.tmpDir, string(typ)+".prof")
-}
-
-func (cfg *config) runscCmd(arg ...string) *exec.Cmd {
+func (cfg *config) runscCmd(arg ...string) (*exec.Cmd, []func()) {
 	var cmd *exec.Cmd
+
+	cmdArgs := []string{cfg.runscPath}
+
 	goProfiling := false
-	for _, typ := range []diagnostics.Type{diagnostics.CPUProfile, diagnostics.MemProfile, diagnostics.Trace} {
-		if driver.DiagnosticEnabled(typ) {
+	var postExit []func()
+	addDiagnostic := func(typ diagnostics.Type, flag string) {
+		if df, err := cfg.diag.Create(typ); err != nil {
+			fmt.Fprintf(os.Stderr, "failed to create %s diagnostics: %s\n", typ, err)
+		} else if df != nil {
+			df.Close()
+			cmdArgs = append(cmdArgs, flag, df.Name())
 			goProfiling = true
-			break
+			postExit = append(postExit, df.Commit)
 		}
 	}
+	addDiagnostic(diagnostics.CPUProfile, "-profile-cpu")
+	addDiagnostic(diagnostics.MemProfile, "-profile-heap")
+	addDiagnostic(diagnostics.Trace, "-trace")
 	if goProfiling {
-		arg = append([]string{"-profile"}, arg...)
+		cmdArgs = append(cmdArgs, "-profile")
 	}
-	if driver.DiagnosticEnabled(diagnostics.CPUProfile) {
-		arg = append([]string{"-profile-cpu", cfg.profilePath(diagnostics.CPUProfile)}, arg...)
-	}
-	if driver.DiagnosticEnabled(diagnostics.MemProfile) {
-		arg = append([]string{"-profile-heap", cfg.profilePath(diagnostics.MemProfile)}, arg...)
-	}
-	if driver.DiagnosticEnabled(diagnostics.Trace) {
-		arg = append([]string{"-trace", cfg.profilePath(diagnostics.Trace)}, arg...)
-	}
-	if driver.DiagnosticEnabled(diagnostics.Perf) {
-		perfArgs := []string{"record", "-o", cfg.profilePath(diagnostics.Perf)}
+
+	if df, err := cfg.diag.Create(diagnostics.Perf); err != nil {
+		fmt.Fprintf(os.Stderr, "failed to create %s diagnostics: %s\n", diagnostics.Perf, err)
+	} else if df != nil {
+		df.Close()
+		postExit = append(postExit, df.Commit)
+
+		perfArgs := []string{"perf", "record", "-o", df.Name()}
 		perfArgs = append(perfArgs, driver.PerfFlags()...)
-		perfArgs = append(perfArgs, cfg.runscPath)
-		perfArgs = append(perfArgs, arg...)
-		cmd = exec.Command("perf", perfArgs...)
-	} else {
-		cmd = exec.Command(cfg.runscPath, arg...)
+		perfArgs = append(perfArgs, cmdArgs...)
+		cmdArgs = perfArgs
 	}
+
+	cmdArgs = append(cmdArgs, arg...)
+	cmd = exec.Command(cmdArgs[0], cmdArgs[1:]...)
 	cmd.SysProcAttr = &syscall.SysProcAttr{
 		// Try to bring down the sandbox if we unexpectedly exit.
 		Pdeathsig: syscall.SIGKILL,
@@ -66,5 +71,5 @@
 		// tree at once.
 		Setpgid: true,
 	}
-	return cmd
+	return cmd, postExit
 }
diff --git a/sweet/benchmarks/gvisor/http_server.go b/sweet/benchmarks/gvisor/http_server.go
index 1115237..70d46e2 100644
--- a/sweet/benchmarks/gvisor/http_server.go
+++ b/sweet/benchmarks/gvisor/http_server.go
@@ -101,7 +101,7 @@
 	defer runtime.GOMAXPROCS(procs)
 	clients := clientProcs
 
-	baseSrvCmd := cfg.runscCmd(
+	baseSrvCmd, postExit := cfg.runscCmd(
 		"-rootless", "do", "-ip", ip,
 		workloadsPath(cfg.assetsDir, "http"),
 		"-host", ip,
@@ -139,6 +139,10 @@
 			err = r
 			return
 		}
+
+		for _, fn := range postExit {
+			fn()
+		}
 	}()
 
 	err = driver.RunBenchmark(b.name()+"Startup", func(d *driver.B) error {
diff --git a/sweet/benchmarks/gvisor/main.go b/sweet/benchmarks/gvisor/main.go
index 988c336..0b1c3c2 100644
--- a/sweet/benchmarks/gvisor/main.go
+++ b/sweet/benchmarks/gvisor/main.go
@@ -17,7 +17,6 @@
 	"time"
 
 	"golang.org/x/benchmarks/sweet/benchmarks/internal/driver"
-	"golang.org/x/benchmarks/sweet/common/diagnostics"
 )
 
 type config struct {
@@ -25,6 +24,8 @@
 	assetsDir string
 	tmpDir    string
 	short     bool
+
+	diag *driver.Diagnostics
 }
 
 var cliCfg config
@@ -62,24 +63,20 @@
 
 	// Run each benchmark once.
 	for _, bench := range benchmarks {
+		cfg := cliCfg
+		cfg.diag = driver.NewDiagnostics(bench.name())
+
 		// Run the benchmark command under runsc.
 		var buf bytes.Buffer
-		if err := bench.run(&cliCfg, &buf); err != nil {
+		if err := bench.run(&cfg, &buf); err != nil {
 			if buf.Len() != 0 {
 				fmt.Fprintf(os.Stderr, "=== Benchmark %s stdout+stderr ===", bench.name())
 				fmt.Fprintf(os.Stderr, "%s\n", buf.String())
 			}
 			return err
 		}
-		for _, typ := range diagnostics.Types() {
-			if !driver.DiagnosticEnabled(typ) {
-				continue
-			}
-			// runscCmd ensures these are created if necessary.
-			if err := driver.CopyDiagnosticData(cliCfg.profilePath(typ), typ, bench.name()); err != nil {
-				return err
-			}
-		}
+
+		cfg.diag.Commit(nil)
 	}
 	return nil
 }
diff --git a/sweet/benchmarks/gvisor/startup.go b/sweet/benchmarks/gvisor/startup.go
index 0439b47..6fb2e64 100644
--- a/sweet/benchmarks/gvisor/startup.go
+++ b/sweet/benchmarks/gvisor/startup.go
@@ -21,10 +21,15 @@
 }
 
 func (b startup) run(cfg *config, out io.Writer) error {
-	cmd := cfg.runscCmd("-rootless", "-network=none", "run", "bench")
+	cmd, postExit := cfg.runscCmd("-rootless", "-network=none", "run", "bench")
 	cmd.Stdout = out
 	cmd.Stderr = out
 	cmd.Dir = filepath.Join(cfg.assetsDir, "startup")
+	defer func() {
+		for _, fn := range postExit {
+			fn()
+		}
+	}()
 	return driver.RunBenchmark(b.name(), func(d *driver.B) error {
 		return cmd.Run()
 	}, driver.DoTime(true))
diff --git a/sweet/benchmarks/gvisor/syscall.go b/sweet/benchmarks/gvisor/syscall.go
index e9b5a91..bfaf6da 100644
--- a/sweet/benchmarks/gvisor/syscall.go
+++ b/sweet/benchmarks/gvisor/syscall.go
@@ -23,13 +23,18 @@
 }
 
 func (b systemCall) run(cfg *config, out io.Writer) error {
-	baseCmd := cfg.runscCmd("-rootless", "do", workloadsPath(cfg.assetsDir, "syscall"))
+	baseCmd, postExit := cfg.runscCmd("-rootless", "do", workloadsPath(cfg.assetsDir, "syscall"))
 	baseCmd.Stdout = out
 	baseCmd.Stderr = out
 	cmd, err := cgroups.WrapCommand(baseCmd, "test-syscall.scope")
 	if err != nil {
 		return err
 	}
+	defer func() {
+		for _, fn := range postExit {
+			fn()
+		}
+	}()
 	return driver.RunBenchmark(b.name(), func(d *driver.B) error {
 		d.Ops(b.ops)
 		d.ResetTimer()