From 82d504ffbfd37c1d7d9adda574272e6837164947 Mon Sep 17 00:00:00 2001
From: Tim Ross <tim.ross@goteleport.com>
Date: Mon, 16 Dec 2024 15:32:48 -0500
Subject: [PATCH] Convert lib/benchmark to use slog

---
 examples/bench/example.go  |  5 +++--
 lib/benchmark/benchmark.go | 10 +++++-----
 tool/tsh/common/tsh.go     |  2 +-
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/examples/bench/example.go b/examples/bench/example.go
index 00b0c5b699cb9..64d2f397f8ce4 100644
--- a/examples/bench/example.go
+++ b/examples/bench/example.go
@@ -35,8 +35,9 @@ func main() {
 	}
 
 	// Run Linear generator
+	ctx := context.Background()
 	results, err := benchmark.Run(
-		context.TODO(),
+		ctx,
 		linear,
 		"host",
 		"username",
@@ -59,7 +60,7 @@ func main() {
 
 	// Export latency profile
 	responseHistogram := results[0].Histogram
-	_, err = benchmark.ExportLatencyProfile("profiles/", responseHistogram, 1, 1.0)
+	_, err = benchmark.ExportLatencyProfile(ctx, "profiles/", responseHistogram, 1, 1.0)
 	if err != nil {
 		fmt.Println(err)
 	}
diff --git a/lib/benchmark/benchmark.go b/lib/benchmark/benchmark.go
index dd919545f7b06..0fdc8f5f9679b 100644
--- a/lib/benchmark/benchmark.go
+++ b/lib/benchmark/benchmark.go
@@ -24,6 +24,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"log/slog"
 	"os"
 	"os/signal"
 	"path/filepath"
@@ -33,7 +34,6 @@ import (
 
 	"github.com/HdrHistogram/hdrhistogram-go"
 	"github.com/gravitational/trace"
-	"github.com/sirupsen/logrus"
 
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/observability/tracing"
@@ -102,7 +102,7 @@ func Run(ctx context.Context, lg *Linear, host, login, proxy string, suite Suite
 		signal.Notify(exitSignals, syscall.SIGTERM, syscall.SIGINT)
 		defer signal.Stop(exitSignals)
 		sig := <-exitSignals
-		logrus.Debugf("signal: %v", sig)
+		slog.DebugContext(ctx, "terminating benchmark due to signal", "signal", sig)
 		cancel()
 	}()
 	var results []Result
@@ -131,7 +131,7 @@ func Run(ctx context.Context, lg *Linear, host, login, proxy string, suite Suite
 }
 
 // ExportLatencyProfile exports the latency profile and returns the path as a string if no errors
-func ExportLatencyProfile(path string, h *hdrhistogram.Histogram, ticks int32, valueScale float64) (string, error) {
+func ExportLatencyProfile(ctx context.Context, path string, h *hdrhistogram.Histogram, ticks int32, valueScale float64) (string, error) {
 	timeStamp := time.Now().Format("2006-01-02_15:04:05")
 	suffix := fmt.Sprintf("latency_profile_%s.txt", timeStamp)
 	if path != "." {
@@ -147,7 +147,7 @@ func ExportLatencyProfile(path string, h *hdrhistogram.Histogram, ticks int32, v
 
 	if _, err := h.PercentilesPrint(fo, ticks, valueScale); err != nil {
 		if err := fo.Close(); err != nil {
-			logrus.WithError(err).Warningf("failed to close file")
+			slog.WarnContext(ctx, "failed to close latency profile file", "error", err)
 		}
 		return "", trace.Wrap(err)
 	}
@@ -259,7 +259,7 @@ func (c *Config) Benchmark(ctx context.Context, tc *client.TeleportClient, suite
 			result.Duration = time.Since(start)
 			return result, nil
 		case <-statusTicker.C:
-			logrus.Infof("working... current observation count: %d", result.RequestsOriginated)
+			slog.InfoContext(ctx, "working...", "current_observation_count", result.RequestsOriginated)
 		}
 	}
 }
diff --git a/tool/tsh/common/tsh.go b/tool/tsh/common/tsh.go
index aa6fb5cbf5666..582cce0fcd08d 100644
--- a/tool/tsh/common/tsh.go
+++ b/tool/tsh/common/tsh.go
@@ -3998,7 +3998,7 @@ func onBenchmark(cf *CLIConf, suite benchmark.Suite) error {
 	}
 	fmt.Fprintf(cf.Stdout(), "\n")
 	if cf.BenchExport {
-		path, err := benchmark.ExportLatencyProfile(cf.BenchExportPath, result.Histogram, cf.BenchTicks, cf.BenchValueScale)
+		path, err := benchmark.ExportLatencyProfile(cf.Context, cf.BenchExportPath, result.Histogram, cf.BenchTicks, cf.BenchValueScale)
 		if err != nil {
 			fmt.Fprintf(cf.Stderr(), "failed exporting latency profile: %s\n", utils.UserMessageFromError(err))
 		} else {