From 139998538e898a2f2fd6fbee0e463cf1ba3f8c2f Mon Sep 17 00:00:00 2001
From: Sambhav Jain <sambhav.jain@cockroachlabs.com>
Date: Mon, 8 Jul 2024 16:00:01 +0530
Subject: [PATCH] roachprod-microbench: add clean command

Epic: none

Release note: None

Previously cleaning benchmark output was only internally exposed in roachprod-microbench, but we now need it as part of a CI job, hence a new utility command "clean" has been added to roachprod-microbench
See: https://github.com/cockroachdb/cockroach/issues/106661
---
 pkg/cmd/roachprod-microbench/BUILD.bazel |  1 +
 pkg/cmd/roachprod-microbench/clean.go    | 85 ++++++++++++++++++++++++
 pkg/cmd/roachprod-microbench/main.go     | 25 +++++++
 3 files changed, 111 insertions(+)
 create mode 100644 pkg/cmd/roachprod-microbench/clean.go

diff --git a/pkg/cmd/roachprod-microbench/BUILD.bazel b/pkg/cmd/roachprod-microbench/BUILD.bazel
index c5eab7be98f6..559d21b209fa 100644
--- a/pkg/cmd/roachprod-microbench/BUILD.bazel
+++ b/pkg/cmd/roachprod-microbench/BUILD.bazel
@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library", "go_test")
 go_library(
     name = "roachprod-microbench_lib",
     srcs = [
+        "clean.go",
         "compare.go",
         "executor.go",
         "export.go",
diff --git a/pkg/cmd/roachprod-microbench/clean.go b/pkg/cmd/roachprod-microbench/clean.go
new file mode 100644
index 000000000000..1f05566e547a
--- /dev/null
+++ b/pkg/cmd/roachprod-microbench/clean.go
@@ -0,0 +1,85 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/cockroachdb/errors"
+)
+
+type cleanConfig struct {
+	inputFilePath  string
+	outputFilePath string
+}
+
+type clean struct {
+	cleanConfig
+	inputFile *os.File
+}
+
+func defaultCleanConfig() cleanConfig {
+	return cleanConfig{
+		inputFilePath:  "",
+		outputFilePath: "",
+	}
+}
+
+func newClean(config cleanConfig) (*clean, error) {
+	file, err := os.Open(config.inputFilePath)
+	if err != nil {
+		return nil, err
+	}
+
+	return &clean{cleanConfig: config, inputFile: file}, nil
+}
+
+func (c *clean) writeCleanOutputToFile(cleanedBenchmarkOutputLog benchmarkExtractionResult) error {
+
+	if err := os.MkdirAll(filepath.Dir(c.outputFilePath), os.ModePerm); err != nil {
+		return err
+	}
+
+	outputFile, err := os.Create(c.outputFilePath)
+	if err != nil {
+		return err
+	}
+	defer outputFile.Close()
+
+	for _, benchmarkResult := range cleanedBenchmarkOutputLog.results {
+		if _, writeErr := outputFile.WriteString(
+			fmt.Sprintf("%s\n", strings.Join(benchmarkResult, " "))); writeErr != nil {
+			return errors.Wrap(writeErr, "failed to write benchmark result to file")
+		}
+	}
+
+	return nil
+}
+
+func (c *clean) cleanBenchmarkOutputLog() error {
+	defer c.inputFile.Close()
+
+	rawBenchmarkLogs, err := io.ReadAll(c.inputFile)
+	if err != nil {
+		return err
+	}
+
+	cleanedBenchmarkOutputLog := extractBenchmarkResults(string(rawBenchmarkLogs))
+	if err = c.writeCleanOutputToFile(cleanedBenchmarkOutputLog); err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/pkg/cmd/roachprod-microbench/main.go b/pkg/cmd/roachprod-microbench/main.go
index 17bb9537928f..1c050c7b5c6d 100644
--- a/pkg/cmd/roachprod-microbench/main.go
+++ b/pkg/cmd/roachprod-microbench/main.go
@@ -48,10 +48,35 @@ Typical usage:
 	command.AddCommand(makeCompareCommand())
 	command.AddCommand(makeRunCommand())
 	command.AddCommand(makeExportCommand())
+	command.AddCommand(makeCleanCommand())
 
 	return command
 }
 
+func makeCleanCommand() *cobra.Command {
+	config := defaultCleanConfig()
+	runCmdFunc := func(cmd *cobra.Command, commandLine []string) error {
+		args, _ := splitArgsAtDash(cmd, commandLine)
+
+		config.inputFilePath = args[0]
+		config.outputFilePath = args[1]
+		c, err := newClean(config)
+		if err != nil {
+			return err
+		}
+
+		return c.cleanBenchmarkOutputLog()
+	}
+	command := &cobra.Command{
+		Use:   "clean <inputFilePath> <outputFilePath>",
+		Short: "remove noisy logs from the benchmark output and dump it to a file",
+		Long:  `remove noisy logs from the benchmark output and dump it to a file`,
+		Args:  cobra.ExactArgs(2),
+		RunE:  runCmdFunc,
+	}
+	return command
+}
+
 func makeRunCommand() *cobra.Command {
 	config := defaultExecutorConfig()
 	runCmdFunc := func(cmd *cobra.Command, commandLine []string) error {