sweet/benchmarks/cockroachdb: use new diagnostics framework

Change-Id: I6e155eee914428ebb665a02264f648ea17bcd73c
Cq-Include-Trybots: luci.golang.try:x_benchmarks-gotip-linux-amd64-longtest
Reviewed-on: https://go-review.googlesource.com/c/benchmarks/+/600063
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Michael Pratt <[email protected]>

sweet/benchmarks/cockroachdb/main.go

@@ -20,7 +20,6 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
-	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -606,52 +605,24 @@ func run(cfg *config) (err error) {
 	return driver.RunBenchmark(cfg.bench.reportName, func(d *driver.B) error {
 		// Set up diagnostics.
 		var stopAll par.Funcs
-		if driver.DiagnosticEnabled(diagnostics.CPUProfile) {
-			for _, inst := range instances {
-				stop := server.PollDiagnostic(
-					inst.httpAddr(),
-					cfg.tmpDir,
-					cfg.bench.reportName,
-					diagnostics.CPUProfile,
-				)
-				stopAll.Add(stop)
-			}
-		}
-		if driver.DiagnosticEnabled(diagnostics.Trace) {
-			var sum atomic.Uint64
-			for _, inst := range instances {
-				stopTrace := server.PollDiagnostic(
-					inst.httpAddr(),
-					cfg.tmpDir,
-					cfg.bench.reportName,
-					diagnostics.Trace,
-				)
-				stopAll.Add(func() {
-					n := stopTrace()
-					sum.Add(n)
-				})
+		diag := driver.NewDiagnostics(cfg.bench.reportName)
+		for _, typ := range diagnostics.Types() {
+			if typ.HTTPEndpoint() == "" {
+				continue
 			}
-			defer func() {
-				d.Report("trace-bytes", sum.Load())
-			}()
-		}
-		if driver.DiagnosticEnabled(diagnostics.MemProfile) {
-			for _, inst := range instances {
-				inst := inst
-				stopAll.Add(func() {
-					_, err := server.CollectDiagnostic(
-						inst.httpAddr(),
-						cfg.tmpDir,
-						cfg.bench.reportName,
-						diagnostics.MemProfile,
-					)
-					if err != nil {
-						fmt.Fprintf(os.Stderr, "failed to read memprofile: %v", err)
-					}
-				})
+			for i, inst := range instances {
+				name := ""
+				if !typ.CanMerge() {
+					// Create a separate file for each instance.
+					name = fmt.Sprintf("inst%d", i)
+				}
+				stop := server.FetchDiagnostic(inst.httpAddr(), diag, typ, name)
+				stopAll.Add(stop)
 			}
 		}
+		defer diag.Commit(d)
 		defer stopAll.Run()
+
 		// Actually run the benchmark.
 		log.Println("running benchmark")
 		return runBenchmark(d, cfg, instances)