sweet/benchmarks/internal/server: collect at least one profile

For profiles that can't be truncated, make sure we collect at least
one 1-second profile from servers. This is important for PGO
benchmarks, which must first do a profile run. Etcd, in particular,
may not run long enough for a complete profile without this.

Change-Id: I4f26dd43d0a3cfd5fa7a4eda488e8b71bbfb65f0
Reviewed-on: https://go-review.googlesource.com/c/benchmarks/+/600756
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>

sweet/benchmarks/internal/server/server.go

diff --git a/sweet/benchmarks/internal/server/server.go b/sweet/benchmarks/internal/server/server.go
index 42681fe..c71b20c 100644
--- a/sweet/benchmarks/internal/server/server.go
+++ b/sweet/benchmarks/internal/server/server.go
@@ -13,6 +13,7 @@
 	"os"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/benchmarks/sweet/benchmarks/internal/driver"
 	"golang.org/x/benchmarks/sweet/common/diagnostics"
@@ -51,8 +52,18 @@
 	go func() {
 		defer wg.Done()
 
+		// If we can't truncate this diagnostic, make sure we collect it at
+		// least once. This is important for PGO, which first does a profiling run.
+		ctx1 := ctx
+		if typ.CanMerge() && !typ.CanTruncate() {
+			var cancel1 func()
+			ctx1, cancel1 = context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel1()
+		}
+
 		for {
-			err := collectTo(ctx, host, diag, typ, name)
+			err := collectTo(ctx1, host, diag, typ, name)
+			ctx1 = ctx
 			if err != nil {
 				if !errors.Is(err, context.Canceled) {
 					fmt.Fprintf(os.Stderr, "failed to read diagnostic %s: %v", typ, err)