From b1b59a6d6490e4aa3a99fb253eccfeff9e658ee5 Mon Sep 17 00:00:00 2001
From: Tim Ross <tim.ross@goteleport.com>
Date: Wed, 3 May 2023 15:33:07 -0400
Subject: [PATCH] Reuse resources for `tsh bench ssh --random`

Retrieves the list of resources a single time when the benchmark
is built and reuses the list within each workload function.
---
 lib/benchmark/ssh.go | 45 +++++++++++++++++++-------------------------
 1 file changed, 19 insertions(+), 26 deletions(-)

diff --git a/lib/benchmark/ssh.go b/lib/benchmark/ssh.go
index 15605cfdee2d8..3184a3c055d79 100644
--- a/lib/benchmark/ssh.go
+++ b/lib/benchmark/ssh.go
@@ -37,42 +37,35 @@ type SSHBenchmark struct {
 
 // BenchBuilder returns a WorkloadFunc for the given benchmark suite.
 func (s SSHBenchmark) BenchBuilder(ctx context.Context, tc *client.TeleportClient) (WorkloadFunc, error) {
+	var resources []types.Server
 	if s.Random {
 		if tc.Host != "all" {
 			return nil, trace.BadParameter("random ssh bench commands must use the format <user>@all <command>")
 		}
 
-		fn, err := s.random(ctx, tc)
-		return fn, trace.Wrap(err)
-	}
-
-	return func(ctx context.Context) error {
-		return tc.SSH(ctx, s.Command, false)
-	}, nil
-}
-
-// random creates a [WorkloadFunc] that executes the provided command on
-// a random host that the user has access to. If hosts disappear or go
-// offline during the benchmark they will not be removed from the list and
-// can have a negative impact on the results.
-func (s SSHBenchmark) random(ctx context.Context, tc *client.TeleportClient) (WorkloadFunc, error) {
-	clt, err := tc.ConnectToCluster(ctx)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-	defer clt.Close()
+		clt, err := tc.ConnectToCluster(ctx)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		defer clt.Close()
 
-	resources, err := apiclient.GetAllResources[types.Server](ctx, clt.AuthClient, tc.ResourceFilter(types.KindNode))
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
+		resources, err = apiclient.GetAllResources[types.Server](ctx, clt.AuthClient, tc.ResourceFilter(types.KindNode))
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
 
-	if len(resources) == 0 {
-		return nil, trace.BadParameter("no target hosts available")
+		if len(resources) == 0 {
+			return nil, trace.BadParameter("no target hosts available")
+		}
 	}
 
 	return func(ctx context.Context) error {
-		return tc.SSH(ctx, s.Command, false, client.WithHostAddress(chooseRandomHost(resources)))
+		var opts []func(*client.SSHOptions)
+		if len(resources) > 0 {
+			opts = append(opts, client.WithHostAddress(chooseRandomHost(resources)))
+		}
+
+		return tc.SSH(ctx, s.Command, false, opts...)
 	}, nil
 }