summaryrefslogtreecommitdiff
path: root/internal/tools
diff options
context:
space:
mode:
Diffstat (limited to 'internal/tools')
-rw-r--r--internal/tools/benchmark/benchmark.go385
-rw-r--r--internal/tools/common/data_generator.go268
-rw-r--r--internal/tools/common/utils.go213
-rw-r--r--internal/tools/pgo/pgo.go1219
-rw-r--r--internal/tools/pgo/pgo_test.go132
-rw-r--r--internal/tools/profile/analyze.go221
-rw-r--r--internal/tools/profile/profile.go367
-rw-r--r--internal/tools/profile/profile_test.go30
8 files changed, 2835 insertions, 0 deletions
diff --git a/internal/tools/benchmark/benchmark.go b/internal/tools/benchmark/benchmark.go
new file mode 100644
index 0000000..b728329
--- /dev/null
+++ b/internal/tools/benchmark/benchmark.go
@@ -0,0 +1,385 @@
+package benchmark
+
+import (
+ "bufio"
+ "flag"
+ "fmt"
+ "io"
+ "os"
+ "os/exec"
+ "path/filepath"
+ "sort"
+ "strings"
+ "time"
+
+ "github.com/mimecast/dtail/internal/tools/common"
+)
+
+// Config holds benchmark configuration
+type Config struct {
+ Mode string
+ BaselineDir string
+ Tag string
+ Quick bool
+ Memory bool
+ OutputFile string
+ Verbose bool
+ Iterations string
+ BaselinePath string
+}
+
+// Run executes the benchmark command
+func Run() error {
+ cfg := parseFlags()
+
+ // Create baseline directory if needed
+ if err := common.EnsureDirectory(cfg.BaselineDir); err != nil {
+ return fmt.Errorf("failed to create baseline directory: %w", err)
+ }
+
+ switch cfg.Mode {
+ case "run":
+ return runBenchmarks(cfg)
+ case "baseline":
+ return createBaseline(cfg)
+ case "compare":
+ return compareWithBaseline(cfg)
+ case "list":
+ return listBaselines(cfg)
+ case "clean":
+ return cleanBaselines(cfg)
+ default:
+ return fmt.Errorf("unknown benchmark mode: %s", cfg.Mode)
+ }
+}
+
+func parseFlags() *Config {
+ cfg := &Config{
+ BaselineDir: "benchmarks/baselines",
+ Iterations: "1x",
+ }
+
+ flag.StringVar(&cfg.Mode, "mode", "run", "Benchmark mode: run, baseline, compare, list, clean")
+ flag.StringVar(&cfg.BaselineDir, "dir", cfg.BaselineDir, "Baseline directory")
+ flag.StringVar(&cfg.Tag, "tag", "", "Tag for baseline (e.g., 'before-optimization')")
+ flag.BoolVar(&cfg.Quick, "quick", false, "Run only quick benchmarks")
+ flag.BoolVar(&cfg.Memory, "memory", false, "Include memory profiling")
+ flag.StringVar(&cfg.OutputFile, "output", "", "Output file for results")
+ flag.BoolVar(&cfg.Verbose, "verbose", false, "Verbose output")
+ flag.StringVar(&cfg.Iterations, "iterations", cfg.Iterations, "Benchmark iterations (e.g., 3x)")
+ flag.StringVar(&cfg.BaselinePath, "baseline", "", "Baseline file for comparison")
+
+ flag.Parse()
+
+ // Handle positional arguments for compare mode
+ if cfg.Mode == "compare" && cfg.BaselinePath == "" {
+ args := flag.Args()
+ if len(args) > 0 {
+ cfg.BaselinePath = args[0]
+ }
+ }
+
+ return cfg
+}
+
+func runBenchmarks(cfg *Config) error {
+ common.PrintSection("Running DTail Benchmarks")
+
+ // Build binaries
+ common.PrintInfo("Building binaries...\n")
+ if err := common.BuildCommands("dcat", "dgrep", "dmap", "dtail", "dserver"); err != nil {
+ return fmt.Errorf("failed to build binaries: %w", err)
+ }
+
+ // Prepare benchmark command
+ args := []string{"test", "-bench=."}
+ if cfg.Quick {
+ args = append(args, "-bench=BenchmarkQuick")
+ }
+ if cfg.Memory {
+ args = append(args, "-benchmem")
+ }
+ if cfg.Iterations != "1x" {
+ args = append(args, fmt.Sprintf("-benchtime=%s", cfg.Iterations))
+ }
+ if cfg.Verbose {
+ args = append(args, "-v")
+ }
+ args = append(args, "./benchmarks")
+
+ // Run benchmarks
+ cmd := exec.Command("go", args...)
+
+ var output []byte
+ var err error
+
+ if cfg.OutputFile != "" {
+ // Capture output for file
+ output, err = cmd.CombinedOutput()
+ if err != nil {
+ return fmt.Errorf("benchmark failed: %w\n%s", err, string(output))
+ }
+
+ // Write to file
+ if err := os.WriteFile(cfg.OutputFile, output, 0644); err != nil {
+ return fmt.Errorf("failed to write output file: %w", err)
+ }
+
+ // Also print to stdout
+ fmt.Print(string(output))
+ common.PrintSuccess("\nResults saved to: %s\n", cfg.OutputFile)
+ } else {
+ // Direct output to stdout
+ cmd.Stdout = os.Stdout
+ cmd.Stderr = os.Stderr
+ if err := cmd.Run(); err != nil {
+ return fmt.Errorf("benchmark failed: %w", err)
+ }
+ }
+
+ return nil
+}
+
+func createBaseline(cfg *Config) error {
+ if cfg.Tag == "" {
+ return fmt.Errorf("baseline tag is required (use -tag)")
+ }
+
+ common.PrintSection("Creating Benchmark Baseline")
+
+ // Generate filename
+ timestamp := time.Now().Format("20060102_150405")
+ safeTag := strings.ReplaceAll(cfg.Tag, " ", "_")
+ safeTag = strings.Map(func(r rune) rune {
+ if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') ||
+ (r >= '0' && r <= '9') || r == '.' || r == '_' || r == '-' {
+ return r
+ }
+ return '_'
+ }, safeTag)
+
+ filename := filepath.Join(cfg.BaselineDir,
+ fmt.Sprintf("baseline_%s_%s.txt", timestamp, safeTag))
+
+ // Create baseline file with metadata
+ file, err := os.Create(filename)
+ if err != nil {
+ return fmt.Errorf("failed to create baseline file: %w", err)
+ }
+ defer file.Close()
+
+ // Write metadata
+ fmt.Fprintf(file, "Git commit: %s\n", common.GetGitCommit())
+ fmt.Fprintf(file, "Date: %s\n", time.Now().Format(time.RFC3339))
+ fmt.Fprintf(file, "Tag: %s\n", cfg.Tag)
+ fmt.Fprintf(file, "----------------------------------------\n")
+
+ // Run benchmarks and capture output
+ args := []string{"test", "-bench=.", "-benchmem"}
+ if cfg.Quick {
+ args = append(args, "-bench=BenchmarkQuick")
+ }
+ if cfg.Iterations != "1x" && cfg.Iterations != "" {
+ args = append(args, fmt.Sprintf("-benchtime=%s", cfg.Iterations))
+ }
+ args = append(args, "./benchmarks")
+
+ cmd := exec.Command("go", args...)
+ cmd.Stdout = io.MultiWriter(file, os.Stdout)
+ cmd.Stderr = os.Stderr
+
+ common.PrintInfo("Running benchmarks for baseline...\n")
+ if err := cmd.Run(); err != nil {
+ return fmt.Errorf("benchmark failed: %w", err)
+ }
+
+ common.PrintSuccess("\nBaseline saved to: %s\n", filename)
+ return nil
+}
+
+func compareWithBaseline(cfg *Config) error {
+ if cfg.BaselinePath == "" {
+ return fmt.Errorf("baseline file required (use -baseline or specify as argument)")
+ }
+
+ if !common.FileExists(cfg.BaselinePath) {
+ return fmt.Errorf("baseline file not found: %s", cfg.BaselinePath)
+ }
+
+ common.PrintSection("Comparing with Baseline")
+ fmt.Printf("Baseline: %s\n\n", cfg.BaselinePath)
+
+ // Run current benchmarks
+ currentFile := filepath.Join(cfg.BaselineDir, "current.txt")
+ args := []string{"test", "-bench=.", "-benchmem"}
+
+ // Check if baseline is quick mode
+ baselineContent, err := os.ReadFile(cfg.BaselinePath)
+ if err != nil {
+ return fmt.Errorf("failed to read baseline: %w", err)
+ }
+ if strings.Contains(string(baselineContent), "BenchmarkQuick") {
+ args = append(args, "-bench=BenchmarkQuick")
+ }
+
+ args = append(args, "./benchmarks")
+
+ cmd := exec.Command("go", args...)
+ output, err := cmd.CombinedOutput()
+ if err != nil {
+ return fmt.Errorf("benchmark failed: %w\n%s", err, string(output))
+ }
+
+ // Save current results
+ if err := os.WriteFile(currentFile, output, 0644); err != nil {
+ return fmt.Errorf("failed to write current results: %w", err)
+ }
+
+ // Print current results
+ fmt.Println("Current benchmark results:")
+ fmt.Println(string(output))
+
+ common.PrintSection("Comparison Report")
+
+ // Try benchstat first
+ if err := runBenchstat(cfg.BaselinePath, currentFile); err != nil {
+ // Fall back to simple diff
+ common.PrintInfo("benchstat not found, showing simple diff:\n\n")
+ if err := showSimpleDiff(cfg.BaselinePath, currentFile); err != nil {
+ return fmt.Errorf("failed to show diff: %w", err)
+ }
+ }
+
+ // Save comparison report
+ reportFile := filepath.Join(cfg.BaselineDir,
+ fmt.Sprintf("comparison_%s.txt", time.Now().Format("20060102_150405")))
+
+ report := fmt.Sprintf("Comparison Report\n"+
+ "Generated: %s\n"+
+ "Baseline: %s\n"+
+ "Current: %s\n"+
+ "================================================================================\n\n",
+ time.Now().Format(time.RFC3339),
+ cfg.BaselinePath,
+ currentFile)
+
+ if err := os.WriteFile(reportFile, []byte(report), 0644); err != nil {
+ common.PrintError("Failed to save comparison report: %v\n", err)
+ } else {
+ common.PrintInfo("\nComparison report saved to: %s\n", reportFile)
+ }
+
+ return nil
+}
+
+func listBaselines(cfg *Config) error {
+ common.PrintSection("Available Baselines")
+
+ pattern := filepath.Join(cfg.BaselineDir, "baseline_*.txt")
+ files, err := filepath.Glob(pattern)
+ if err != nil {
+ return fmt.Errorf("failed to list baselines: %w", err)
+ }
+
+ if len(files) == 0 {
+ fmt.Printf("No baselines found in %s\n", cfg.BaselineDir)
+ return nil
+ }
+
+ // Sort by modification time (newest first)
+ sort.Slice(files, func(i, j int) bool {
+ fi, _ := os.Stat(files[i])
+ fj, _ := os.Stat(files[j])
+ return fi.ModTime().After(fj.ModTime())
+ })
+
+ // Display baselines
+ for _, file := range files {
+ info, err := os.Stat(file)
+ if err != nil {
+ continue
+ }
+
+ // Try to extract tag from file
+ tag := extractTagFromBaseline(file)
+
+ fmt.Printf(" %s %8s %-40s %s\n",
+ info.ModTime().Format("2006-01-02 15:04:05"),
+ common.FormatSize(info.Size()),
+ filepath.Base(file),
+ tag)
+ }
+
+ fmt.Printf("\nTotal: %d baselines\n", len(files))
+ fmt.Printf("\nUsage: dtail-tools benchmark -mode compare <baseline_file>\n")
+
+ return nil
+}
+
+func cleanBaselines(cfg *Config) error {
+ common.PrintSection("Cleaning Old Baselines")
+
+ pattern := filepath.Join(cfg.BaselineDir, "baseline_*.txt")
+ files, err := filepath.Glob(pattern)
+ if err != nil {
+ return fmt.Errorf("failed to list baselines: %w", err)
+ }
+
+ if len(files) <= 10 {
+ fmt.Println("No old baselines to clean (keeping last 10)")
+ return nil
+ }
+
+ // Sort by modification time (oldest first)
+ sort.Slice(files, func(i, j int) bool {
+ fi, _ := os.Stat(files[i])
+ fj, _ := os.Stat(files[j])
+ return fi.ModTime().Before(fj.ModTime())
+ })
+
+ // Remove old files
+ toRemove := files[:len(files)-10]
+ for _, file := range toRemove {
+ fmt.Printf("Removing: %s\n", filepath.Base(file))
+ if err := os.Remove(file); err != nil {
+ common.PrintError("Failed to remove %s: %v\n", file, err)
+ }
+ }
+
+ common.PrintSuccess("\nRemoved %d old baselines\n", len(toRemove))
+ return nil
+}
+
+func extractTagFromBaseline(filename string) string {
+ file, err := os.Open(filename)
+ if err != nil {
+ return ""
+ }
+ defer file.Close()
+
+ scanner := bufio.NewScanner(file)
+ for scanner.Scan() {
+ line := scanner.Text()
+ if strings.HasPrefix(line, "Tag: ") {
+ return strings.TrimPrefix(line, "Tag: ")
+ }
+ if strings.HasPrefix(line, "----") {
+ break
+ }
+ }
+ return ""
+}
+
+func runBenchstat(baseline, current string) error {
+ cmd := exec.Command("benchstat", baseline, current)
+ cmd.Stdout = os.Stdout
+ cmd.Stderr = os.Stderr
+ return cmd.Run()
+}
+
+func showSimpleDiff(baseline, current string) error {
+ cmd := exec.Command("diff", "-u", baseline, current)
+ output, _ := cmd.CombinedOutput()
+ fmt.Print(string(output))
+ return nil
+} \ No newline at end of file
diff --git a/internal/tools/common/data_generator.go b/internal/tools/common/data_generator.go
new file mode 100644
index 0000000..d3d4225
--- /dev/null
+++ b/internal/tools/common/data_generator.go
@@ -0,0 +1,268 @@
+package common
+
+import (
+ "bufio"
+ "fmt"
+ "math/rand"
+ "os"
+ "path/filepath"
+ "time"
+)
+
+// DataFormat represents the format of generated data
+type DataFormat string
+
+// Supported data generator output formats.
+const (
+ // FormatLog generates generic log lines.
+ FormatLog DataFormat = "log"
+ FormatCSV DataFormat = "csv"
+ FormatDTail DataFormat = "dtail"
+ FormatMapReduce DataFormat = "mapreduce"
+)
+
+// DataGenerator generates test data for profiling and benchmarking
+type DataGenerator struct {
+ rand *rand.Rand
+}
+
+// NewDataGenerator creates a new data generator
+func NewDataGenerator() *DataGenerator {
+ return &DataGenerator{
+ rand: rand.New(rand.NewSource(time.Now().UnixNano())),
+ }
+}
+
+// GenerateFile generates a test data file of the specified size and format
+func (g *DataGenerator) GenerateFile(filename string, sizeStr string, format DataFormat) error {
+ size, err := ParseSize(sizeStr)
+ if err != nil {
+ return fmt.Errorf("invalid size: %w", err)
+ }
+
+ // Create directory if needed
+ dir := filepath.Dir(filename)
+ if err := os.MkdirAll(dir, 0755); err != nil {
+ return fmt.Errorf("failed to create directory: %w", err)
+ }
+
+ // Check if file already exists
+ if _, err := os.Stat(filename); err == nil {
+ return nil // File exists, skip generation
+ }
+
+ switch format {
+ case FormatLog:
+ return g.generateLogFile(filename, size)
+ case FormatCSV:
+ return g.generateCSVFile(filename, size)
+ case FormatDTail, FormatMapReduce:
+ return g.generateDTailFormatFile(filename, size)
+ default:
+ return fmt.Errorf("unsupported format: %s", format)
+ }
+}
+
+// GenerateLogFileWithLines generates a log file with specific number of lines
+func (g *DataGenerator) GenerateLogFileWithLines(filename string, lines int, format DataFormat) error {
+ // Create directory if needed
+ dir := filepath.Dir(filename)
+ if err := os.MkdirAll(dir, 0755); err != nil {
+ return fmt.Errorf("failed to create directory: %w", err)
+ }
+
+ // Check if file already exists
+ if _, err := os.Stat(filename); err == nil {
+ return nil // File exists, skip generation
+ }
+
+ switch format {
+ case FormatDTail, FormatMapReduce:
+ return g.generateDTailFormatFileWithLines(filename, lines)
+ default:
+ return fmt.Errorf("line-based generation only supported for dtail/mapreduce format")
+ }
+}
+
+func (g *DataGenerator) generateLogFile(filename string, targetSize int64) error {
+ file, err := os.Create(filename)
+ if err != nil {
+ return err
+ }
+ defer file.Close()
+
+ writer := bufio.NewWriter(file)
+ defer writer.Flush()
+
+ var currentSize int64
+ lineNum := 0
+ levels := []string{"INFO", "DEBUG", "WARN", "ERROR"}
+ users := []string{"user1", "user2", "user3", "user4", "user5", "admin", "guest", "service", "monitor", "test"}
+ actions := []string{"login", "logout", "query", "update", "delete", "create", "read", "write", "sync", "backup"}
+
+ for currentSize < targetSize {
+ lineNum++
+ timestamp := time.Now().Add(time.Duration(-lineNum) * time.Second).Format("2006-01-02 15:04:05")
+ level := levels[g.rand.Intn(len(levels))]
+ user := users[g.rand.Intn(len(users))]
+ action := actions[g.rand.Intn(len(actions))]
+ duration := g.rand.Intn(5000) + 100
+ status := "success"
+ if g.rand.Float32() < 0.1 {
+ status = "failure"
+ }
+
+ line := fmt.Sprintf("[%s] %s - User %s performed %s action (duration: %dms, status: %s)\n",
+ timestamp, level, user, action, duration, status)
+
+ n, err := writer.WriteString(line)
+ if err != nil {
+ return err
+ }
+ currentSize += int64(n)
+ }
+
+ return nil
+}
+
+func (g *DataGenerator) generateCSVFile(filename string, targetSize int64) error {
+ file, err := os.Create(filename)
+ if err != nil {
+ return err
+ }
+ defer file.Close()
+
+ writer := bufio.NewWriter(file)
+ defer writer.Flush()
+
+ // Write header
+ header := "timestamp,user,action,duration,status\n"
+ n, err := writer.WriteString(header)
+ if err != nil {
+ return err
+ }
+ currentSize := int64(n)
+
+ lineNum := 0
+ users := []string{"user1", "user2", "user3", "user4", "user5", "admin", "guest", "service", "monitor", "test"}
+ actions := []string{"login", "logout", "query", "update", "delete", "create", "read", "write", "sync", "backup"}
+
+ for currentSize < targetSize {
+ lineNum++
+ timestamp := time.Now().Add(time.Duration(-lineNum) * time.Second).Format("2006-01-02 15:04:05")
+ user := users[g.rand.Intn(len(users))]
+ action := actions[g.rand.Intn(len(actions))]
+ duration := g.rand.Intn(5000) + 100
+ status := "success"
+ if g.rand.Float32() < 0.1 {
+ status = "failure"
+ }
+
+ line := fmt.Sprintf("%s,%s,%s,%d,%s\n", timestamp, user, action, duration, status)
+
+ n, err := writer.WriteString(line)
+ if err != nil {
+ return err
+ }
+ currentSize += int64(n)
+ }
+
+ return nil
+}
+
+func (g *DataGenerator) generateDTailFormatFile(filename string, targetSize int64) error {
+ file, err := os.Create(filename)
+ if err != nil {
+ return err
+ }
+ defer file.Close()
+
+ writer := bufio.NewWriter(file)
+ defer writer.Flush()
+
+ var currentSize int64
+ lineNum := 0
+ hostnames := []string{"server01", "server02", "server03", "server04", "server05",
+ "server06", "server07", "server08", "server09", "server10"}
+
+ for currentSize < targetSize {
+ lineNum++
+ hostname := hostnames[lineNum%len(hostnames)]
+ timestamp := fmt.Sprintf("%02d%02d-%02d%02d%02d",
+ 10+(lineNum/86400)%12, (lineNum/3600)%30+1,
+ (lineNum/3600)%24, (lineNum/60)%60, lineNum%60)
+ goroutines := 10 + (lineNum % 50)
+ cgocalls := lineNum % 100
+ cpus := 1 + (lineNum % 8)
+ loadavg := float64(lineNum%100) / 100.0
+ uptime := fmt.Sprintf("%dh%dm%ds", lineNum/3600, (lineNum/60)%60, lineNum%60)
+ currentConnections := lineNum % 20
+ lifetimeConnections := 1000 + lineNum
+
+ line := fmt.Sprintf("INFO|%s|1|stats.go:56|%d|%d|%d|%.2f|%s|MAPREDUCE:STATS|hostname=%s|currentConnections=%d|lifetimeConnections=%d\n",
+ timestamp, cpus, goroutines, cgocalls, loadavg, uptime, hostname, currentConnections, lifetimeConnections)
+
+ n, err := writer.WriteString(line)
+ if err != nil {
+ return err
+ }
+ currentSize += int64(n)
+ }
+
+ return nil
+}
+
+func (g *DataGenerator) generateDTailFormatFileWithLines(filename string, lines int) error {
+ file, err := os.Create(filename)
+ if err != nil {
+ return err
+ }
+ defer file.Close()
+
+ writer := bufio.NewWriter(file)
+ defer writer.Flush()
+
+ hostnames := []string{"server01", "server02", "server03", "server04", "server05",
+ "server06", "server07", "server08", "server09", "server10"}
+
+ for i := 1; i <= lines; i++ {
+ hostname := hostnames[i%len(hostnames)]
+ timestamp := fmt.Sprintf("%02d%02d-%02d%02d%02d",
+ 10+(i/86400)%12, (i/3600)%30+1,
+ (i/3600)%24, (i/60)%60, i%60)
+ goroutines := 10 + (i % 50)
+ cgocalls := i % 100
+ cpus := 1 + (i % 8)
+ loadavg := float64(i%100) / 100.0
+ uptime := fmt.Sprintf("%dh%dm%ds", i/3600, (i/60)%60, i%60)
+ currentConnections := i % 20
+ lifetimeConnections := 1000 + i
+
+ line := fmt.Sprintf("INFO|%s|1|stats.go:56|%d|%d|%d|%.2f|%s|MAPREDUCE:STATS|hostname=%s|currentConnections=%d|lifetimeConnections=%d\n",
+ timestamp, cpus, goroutines, cgocalls, loadavg, uptime, hostname, currentConnections, lifetimeConnections)
+
+ if _, err := writer.WriteString(line); err != nil {
+ return err
+ }
+ }
+
+ return nil
+}
+
+// GenerateLogFile generates a log file with specified number of lines
+// This is a convenience function for PGO module
+func GenerateLogFile(filename string, lines int) error {
+ g := NewDataGenerator()
+ // Estimate size based on average line length (about 100 bytes per line)
+ estimatedSize := int64(lines * 100)
+ return g.generateLogFile(filename, estimatedSize)
+}
+
+// GenerateCSVFile generates a CSV file with specified number of lines
+// This is a convenience function for PGO module
+func GenerateCSVFile(filename string, lines int) error {
+ g := NewDataGenerator()
+ // Estimate size based on average line length (about 50 bytes per line)
+ estimatedSize := int64(lines * 50)
+ return g.generateCSVFile(filename, estimatedSize)
+}
diff --git a/internal/tools/common/utils.go b/internal/tools/common/utils.go
new file mode 100644
index 0000000..37f115a
--- /dev/null
+++ b/internal/tools/common/utils.go
@@ -0,0 +1,213 @@
+package common
+
+import (
+ "fmt"
+ "os"
+ "os/exec"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "time"
+)
+
+// ParseSize parses a size string like "10MB", "1GB" into bytes
+func ParseSize(sizeStr string) (int64, error) {
+ originalStr := sizeStr
+ sizeStr = strings.ToUpper(strings.TrimSpace(sizeStr))
+
+ // Handle single-letter suffixes (K, M, G, T) by adding B
+ if len(sizeStr) > 1 {
+ lastChar := sizeStr[len(sizeStr)-1]
+ secondLastChar := byte('0')
+ if len(sizeStr) > 1 {
+ secondLastChar = sizeStr[len(sizeStr)-2]
+ }
+
+ // If ends with K, M, G, or T and the character before it is a digit, add B
+ if (lastChar == 'K' || lastChar == 'M' || lastChar == 'G' || lastChar == 'T') &&
+ (secondLastChar >= '0' && secondLastChar <= '9') {
+ sizeStr = sizeStr + "B"
+ }
+ }
+
+ // Order matters - check longer suffixes first
+ suffixes := []struct {
+ suffix string
+ multiplier int64
+ }{
+ {"TB", 1024 * 1024 * 1024 * 1024},
+ {"GB", 1024 * 1024 * 1024},
+ {"MB", 1024 * 1024},
+ {"KB", 1024},
+ {"B", 1},
+ }
+
+ for _, s := range suffixes {
+ if strings.HasSuffix(sizeStr, s.suffix) {
+ numStr := strings.TrimSuffix(sizeStr, s.suffix)
+ numStr = strings.TrimSpace(numStr)
+ if numStr == "" {
+ return 0, fmt.Errorf("no number before size suffix")
+ }
+ num, err := strconv.ParseFloat(numStr, 64)
+ if err != nil {
+ return 0, fmt.Errorf("invalid size number: %s (original: %s, processed: %s)", numStr, originalStr, sizeStr)
+ }
+ return int64(num * float64(s.multiplier)), nil
+ }
+ }
+
+ // Try parsing as plain number (assume bytes)
+ num, err := strconv.ParseInt(sizeStr, 10, 64)
+ if err != nil {
+ return 0, fmt.Errorf("invalid size format: %s", sizeStr)
+ }
+ return num, nil
+}
+
+// FormatSize formats bytes into human-readable size
+func FormatSize(bytes int64) string {
+ const unit = 1024
+ if bytes < unit {
+ return fmt.Sprintf("%d B", bytes)
+ }
+ div, exp := int64(unit), 0
+ for n := bytes / unit; n >= unit; n /= unit {
+ div *= unit
+ exp++
+ }
+ return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
+}
+
+// BuildCommand builds a dtail command if it doesn't exist
+func BuildCommand(cmd string) error {
+ // Check if binary exists
+ if _, err := os.Stat(cmd); err == nil {
+ return nil // Already exists
+ }
+
+ // Build the command
+ cmdName := filepath.Base(cmd)
+ buildCmd := exec.Command("go", "build", "-o", cmd, fmt.Sprintf("./cmd/%s/main.go", cmdName))
+ buildCmd.Stdout = os.Stdout
+ buildCmd.Stderr = os.Stderr
+
+ fmt.Printf("Building %s...\n", cmdName)
+ return buildCmd.Run()
+}
+
+// BuildCommands builds multiple dtail commands
+func BuildCommands(commands ...string) error {
+ for _, cmd := range commands {
+ if err := BuildCommand(cmd); err != nil {
+ return fmt.Errorf("failed to build %s: %w", cmd, err)
+ }
+ }
+ return nil
+}
+
+// EnsureDirectory creates a directory if it doesn't exist
+func EnsureDirectory(dir string) error {
+ return os.MkdirAll(dir, 0755)
+}
+
+// FileExists checks if a file exists
+func FileExists(path string) bool {
+ _, err := os.Stat(path)
+ return err == nil
+}
+
+// GetTimestamp returns a timestamp string for file naming
+func GetTimestamp() string {
+ return time.Now().Format("20060102_150405")
+}
+
+// GetGitCommit returns the current git commit hash (short form)
+func GetGitCommit() string {
+ cmd := exec.Command("git", "rev-parse", "--short", "HEAD")
+ output, err := cmd.Output()
+ if err != nil {
+ return "unknown"
+ }
+ return strings.TrimSpace(string(output))
+}
+
+// RunCommandWithTimeout runs a command with a timeout
+func RunCommandWithTimeout(timeout time.Duration, name string, args ...string) error {
+ cmd := exec.Command(name, args...)
+ cmd.Stdout = os.Stdout
+ cmd.Stderr = os.Stderr
+
+ if err := cmd.Start(); err != nil {
+ return err
+ }
+
+ done := make(chan error, 1)
+ go func() {
+ done <- cmd.Wait()
+ }()
+
+ select {
+ case <-time.After(timeout):
+ if err := cmd.Process.Kill(); err != nil {
+ return fmt.Errorf("failed to kill process: %w", err)
+ }
+ return fmt.Errorf("command timed out after %v", timeout)
+ case err := <-done:
+ return err
+ }
+}
+
+// CleanupFiles removes temporary files matching patterns
+func CleanupFiles(patterns ...string) error {
+ for _, pattern := range patterns {
+ matches, err := filepath.Glob(pattern)
+ if err != nil {
+ return fmt.Errorf("invalid pattern %s: %w", pattern, err)
+ }
+ for _, match := range matches {
+ if err := os.Remove(match); err != nil && !os.IsNotExist(err) {
+ return fmt.Errorf("failed to remove %s: %w", match, err)
+ }
+ }
+ }
+ return nil
+}
+
+// Colors for terminal output
+const (
+ ColorReset = "\033[0m"
+ ColorRed = "\033[0;31m"
+ ColorGreen = "\033[0;32m"
+ ColorYellow = "\033[1;33m"
+ ColorBlue = "\033[0;34m"
+ ColorPurple = "\033[0;35m"
+ ColorCyan = "\033[0;36m"
+ ColorWhite = "\033[0;37m"
+)
+
+// PrintColored prints colored text to stdout
+func PrintColored(color, format string, args ...interface{}) {
+ fmt.Printf(color+format+ColorReset, args...)
+}
+
+// PrintSection prints a section header
+func PrintSection(title string) {
+ PrintColored(ColorGreen, "%s\n", title)
+ fmt.Println(strings.Repeat("=", len(title)))
+}
+
+// PrintInfo prints an info message
+func PrintInfo(format string, args ...interface{}) {
+ PrintColored(ColorYellow, format, args...)
+}
+
+// PrintError prints an error message
+func PrintError(format string, args ...interface{}) {
+ PrintColored(ColorRed, format, args...)
+}
+
+// PrintSuccess prints a success message
+func PrintSuccess(format string, args ...interface{}) {
+ PrintColored(ColorGreen, format, args...)
+} \ No newline at end of file
diff --git a/internal/tools/pgo/pgo.go b/internal/tools/pgo/pgo.go
new file mode 100644
index 0000000..3cabf23
--- /dev/null
+++ b/internal/tools/pgo/pgo.go
@@ -0,0 +1,1219 @@
+package pgo
+
+import (
+ "bytes"
+ "crypto/rand"
+ "crypto/rsa"
+ "crypto/x509"
+ "encoding/pem"
+ "flag"
+ "fmt"
+ "io"
+ "net"
+ "net/http"
+ "os"
+ "os/exec"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/mimecast/dtail/internal/tools/common"
+
+ "golang.org/x/crypto/ssh"
+)
+
+// Config holds PGO configuration
+type Config struct {
+ Command string // Command to build with PGO (dtail, dcat, etc.)
+ ProfileDir string // Directory containing profile data
+ OutputDir string // Directory for PGO-optimized binaries
+ TestDataSize int // Size of test data for profile generation
+ TestIterations int // Number of iterations for profile generation
+ Verbose bool // Verbose output
+ Commands []string // Specific commands to optimize (empty = all)
+ ProfileOnly bool // Only generate profiles, don't build optimized binaries
+}
+
+// Run executes the PGO workflow
+func Run() error {
+ var cfg Config
+
+ // Define flags
+ flag.StringVar(&cfg.ProfileDir, "profiledir", "pgo-profiles", "Directory for profile data")
+ flag.StringVar(&cfg.OutputDir, "outdir", "pgo-build", "Directory for PGO-optimized binaries")
+ flag.IntVar(&cfg.TestDataSize, "datasize", 1000000, "Lines of test data for profile generation")
+ flag.IntVar(&cfg.TestIterations, "iterations", 3, "Number of profile generation iterations")
+ flag.BoolVar(&cfg.Verbose, "verbose", false, "Verbose output")
+ flag.BoolVar(&cfg.Verbose, "v", false, "Verbose output (short)")
+ flag.BoolVar(&cfg.ProfileOnly, "profileonly", false, "Only generate profiles, don't build optimized binaries")
+
+ // Custom usage
+ flag.Usage = func() {
+ fmt.Fprintf(os.Stderr, "Usage: dtail-tools pgo [options] [commands...]\n\n")
+ fmt.Fprintf(os.Stderr, "Profile-Guided Optimization (PGO) for DTail commands\n\n")
+ fmt.Fprintf(os.Stderr, "Options:\n")
+ flag.PrintDefaults()
+ fmt.Fprintf(os.Stderr, "\nCommands:\n")
+ fmt.Fprintf(os.Stderr, " If no commands specified, all dtail commands will be optimized\n")
+ fmt.Fprintf(os.Stderr, " Available: dtail, dcat, dgrep, dmap, dserver\n\n")
+ fmt.Fprintf(os.Stderr, "Example:\n")
+ fmt.Fprintf(os.Stderr, " dtail-tools pgo # Optimize all commands\n")
+ fmt.Fprintf(os.Stderr, " dtail-tools pgo dcat dgrep # Optimize specific commands\n")
+ fmt.Fprintf(os.Stderr, " dtail-tools pgo -v -iterations 5 # Verbose with 5 iterations\n")
+ }
+
+ flag.Parse()
+
+ // Get commands from remaining args
+ cfg.Commands = flag.Args()
+ if len(cfg.Commands) == 0 {
+ // All commands can now be profiled to completion. dtail was previously
+ // excluded because its follow client never returned from client.Start on
+ // -shutdownAfter/SIGINT/SIGTERM, so it never flushed a CPU profile (that
+ // is how the committed dtail.pprof came to be 0 bytes). That follow
+ // shutdown is now honoured (task 1v0): runDtailWorkload profiles a real
+ // follow session bounded by -shutdownAfter, which returns and flushes.
+ // dtail (and dserver) run after the always-serverless dcat/dgrep/dmap so
+ // that a setup problem in the SSH-based workloads cannot abort the run
+ // before the serverless profiles are captured.
+ cfg.Commands = []string{"dcat", "dgrep", "dmap", "dtail", "dserver"}
+ }
+
+ return runPGO(&cfg)
+}
+
+func runPGO(cfg *Config) error {
+ // Create directories
+ if err := os.MkdirAll(cfg.ProfileDir, 0755); err != nil {
+ return fmt.Errorf("creating profile directory: %w", err)
+ }
+ if err := os.MkdirAll(cfg.OutputDir, 0755); err != nil {
+ return fmt.Errorf("creating output directory: %w", err)
+ }
+
+ fmt.Println("DTail Profile-Guided Optimization")
+ fmt.Println("=================================")
+ fmt.Printf("Commands: %s\n", strings.Join(cfg.Commands, ", "))
+ fmt.Printf("Profile directory: %s\n", cfg.ProfileDir)
+ fmt.Printf("Output directory: %s\n", cfg.OutputDir)
+ fmt.Printf("Test data size: %d lines\n", cfg.TestDataSize)
+ fmt.Printf("Iterations: %d\n\n", cfg.TestIterations)
+
+ // Step 1: Build baseline binaries
+ fmt.Println("Step 1: Building baseline binaries...")
+ if err := buildBaseline(cfg); err != nil {
+ return fmt.Errorf("building baseline: %w", err)
+ }
+
+ // Step 2: Generate profiles
+ fmt.Println("\nStep 2: Generating profiles...")
+ if err := generateProfiles(cfg); err != nil {
+ return fmt.Errorf("generating profiles: %w", err)
+ }
+
+ // If profile-only mode, stop here
+ if cfg.ProfileOnly {
+ fmt.Println("\nProfile generation complete!")
+ fmt.Printf("Profiles saved in: %s\n", cfg.ProfileDir)
+ return nil
+ }
+
+ // Step 3: Build PGO-optimized binaries
+ fmt.Println("\nStep 3: Building PGO-optimized binaries...")
+ if err := buildWithPGO(cfg); err != nil {
+ return fmt.Errorf("building with PGO: %w", err)
+ }
+
+ // Step 4: Compare performance
+ fmt.Println("\nStep 4: Comparing performance...")
+ if err := comparePerformance(cfg); err != nil {
+ return fmt.Errorf("comparing performance: %w", err)
+ }
+
+ fmt.Println("\nPGO optimization complete!")
+ fmt.Printf("Optimized binaries are in: %s\n", cfg.OutputDir)
+
+ return nil
+}
+
+func buildBaseline(cfg *Config) error {
+ for _, cmd := range cfg.Commands {
+ if cfg.Verbose {
+ fmt.Printf("Building %s...\n", cmd)
+ }
+
+ // Build command
+ buildCmd := exec.Command("go", "build",
+ "-o", filepath.Join(cfg.OutputDir, cmd+"-baseline"),
+ fmt.Sprintf("./cmd/%s", cmd))
+
+ if cfg.Verbose {
+ buildCmd.Stdout = os.Stdout
+ buildCmd.Stderr = os.Stderr
+ }
+
+ if err := buildCmd.Run(); err != nil {
+ return fmt.Errorf("building %s: %w", cmd, err)
+ }
+ }
+
+ return nil
+}
+
+func generateProfiles(cfg *Config) error {
+ // Generate test data
+ testFiles, err := generateTestData(cfg)
+ if err != nil {
+ return fmt.Errorf("generating test data: %w", err)
+ }
+ defer cleanupTestData(testFiles)
+
+ // Run each command to generate profiles
+ for _, cmd := range cfg.Commands {
+ fmt.Printf("\nGenerating profile for %s...\n", cmd)
+
+ profilePath := filepath.Join(cfg.ProfileDir, fmt.Sprintf("%s.pprof", cmd))
+
+ // Run iterations to collect profile data
+ if err := runProfileWorkload(cfg, cmd, testFiles, profilePath); err != nil {
+ return fmt.Errorf("running workload for %s: %w", cmd, err)
+ }
+
+ // Sanity-check the freshly captured profile. A zero-sample or empty
+ // profile is worthless for PGO and must never be silently accepted.
+ if err := verifyProfileNonEmpty(cmd, profilePath); err != nil {
+ return fmt.Errorf("verifying profile for %s: %w", cmd, err)
+ }
+ }
+
+ return nil
+}
+
+// countRawSamples counts the sample rows in the textual output of
+// "go tool pprof -raw". That output lists each captured sample on its own
+// indented line between the "Samples:" header (which is followed by a single
+// units line, e.g. "samples/count cpu/nanoseconds") and the "Locations"
+// section:
+//
+// Samples:
+// samples/count cpu/nanoseconds
+// 1 10000000: 1 2 3 4 5 6 7 8
+// 3 30000000: 9 10 11 5 6 7 8
+// Locations
+//
+// A profile captured from an idle process has the header but no data rows.
+// Parsing the textual form keeps this dependency-free (go.mod is intentionally
+// lean) and mirrors the existing use of "go tool pprof" for merging.
+func countRawSamples(raw string) int {
+ inSamples := false
+ sawUnits := false
+ count := 0
+ for _, line := range strings.Split(raw, "\n") {
+ trimmed := strings.TrimSpace(line)
+ if !inSamples {
+ if trimmed == "Samples:" {
+ inSamples = true
+ }
+ continue
+ }
+ // The "Locations" line terminates the sample table.
+ if strings.HasPrefix(trimmed, "Locations") {
+ break