Optimize mapr parsing and stabilize aggregate shutdown

2 files changed, 88 insertions, 49 deletions

diff --git a/internal/tools/profile/profile.go b/internal/tools/profile/profile.go
index d87662c..21508b2 100644
--- a/internal/tools/profile/profile.go
+++ b/internal/tools/profile/profile.go
@@ -64,29 +64,29 @@ func parseFlags() *Config {
 	flag.IntVar(&cfg.Runs, "runs", 1, "Number of profiling runs")
 	flag.BoolVar(&cfg.NoColor, "nocolor", false, "Disable colored output")
 	flag.DurationVar(&cfg.Timeout, "timeout", cfg.Timeout, "Timeout for profiling runs")
-	
+
 	// Custom command list
 	var cmdList string
 	flag.StringVar(&cmdList, "commands", "", "Comma-separated list of commands to profile")
-	
+
 	flag.Parse()
-	
+
 	if cmdList != "" {
 		cfg.Commands = strings.Split(cmdList, ",")
 	}
-	
+
 	return cfg
 }
 
 func runQuickProfile(cfg *Config) error {
 	common.PrintSection("DTail Quick Profiling")
-	
+
 	// Generate test data
 	gen := common.NewDataGenerator()
-	
+
 	logFile := filepath.Join(cfg.TestDataDir, "quick_test.log")
 	csvFile := filepath.Join(cfg.TestDataDir, "quick_test.csv")
-	
+
 	common.PrintInfo("Generating test data...\n")
 	if err := gen.GenerateFile(logFile, "10MB", common.FormatLog); err != nil {
 		return fmt.Errorf("failed to generate log file: %w", err)
@@ -94,31 +94,31 @@ func runQuickProfile(cfg *Config) error {
 	if err := gen.GenerateFile(csvFile, "10MB", common.FormatCSV); err != nil {
 		return fmt.Errorf("failed to generate CSV file: %w", err)
 	}
-	
+
 	// Build commands
 	common.PrintInfo("Building commands...\n")
 	if err := common.BuildCommands("dcat", "dgrep", "dmap"); err != nil {
 		return err
 	}
-	
+
 	// Profile each command
 	common.PrintSection("Running quick profiles...")
-	
+
 	// Profile dcat
 	if err := profileCommand("dcat", "dcat",
 		[]string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none", logFile},
 		cfg.Timeout); err != nil {
 		return err
 	}
-	
+
 	// Profile dgrep
 	if err := profileCommand("dgrep", "dgrep",
-		[]string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none", 
+		[]string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
 			"-regex", "user[0-9]+", logFile},
 		cfg.Timeout); err != nil {
 		return err
 	}
-	
+
 	// Profile dmap
 	query := `select count($line),avg($duration) group by $user logformat csv`
 	if err := profileCommand("dmap", "dmap",
@@ -127,24 +127,24 @@ func runQuickProfile(cfg *Config) error {
 		cfg.Timeout); err != nil {
 		return err
 	}
-	
+
 	// Analyze results
 	return analyzeLatestProfiles(cfg)
 }
 
 func runFullProfile(cfg *Config) error {
 	common.PrintSection("DTail Full Profiling")
-	
+
 	// Generate test data
 	gen := common.NewDataGenerator()
-	
+
 	testFiles := map[string]string{
 		"small.log":        "10MB",
 		"medium.log":       "100MB",
 		"test.csv":         "50MB",
 		"dtail_format.log": "100000", // lines
 	}
-	
+
 	common.PrintInfo("Generating test data...\n")
 	for filename, size := range testFiles {
 		fullPath := filepath.Join(cfg.TestDataDir, filename)
@@ -163,16 +163,16 @@ func runFullProfile(cfg *Config) error {
 			}
 		}
 	}
-	
+
 	// Build commands
 	common.PrintInfo("Building commands...\n")
 	if err := common.BuildCommands("dcat", "dgrep", "dmap"); err != nil {
 		return err
 	}
-	
+
 	// Run profiling
 	common.PrintSection("Running full profiling suite...")
-	
+
 	// Profile configurations
 	profiles := []struct {
 		cmd  string
@@ -184,24 +184,24 @@ func runFullProfile(cfg *Config) error {
 			filepath.Join(cfg.TestDataDir, "small.log")}},
 		{"dcat", "medium_file", []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
 			filepath.Join(cfg.TestDataDir, "medium.log")}},
-		
+
 		// dgrep profiles
 		{"dgrep", "simple_pattern", []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
 			"-regex", "ERROR", filepath.Join(cfg.TestDataDir, "medium.log")}},
 		{"dgrep", "complex_pattern", []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
 			"-regex", "(ERROR|WARN).*user[0-9]+", filepath.Join(cfg.TestDataDir, "medium.log")}},
-		
+
 		// dmap profiles
 		{"dmap", "simple_count", []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
 			"-query", "from STATS select count(*)", "-files", filepath.Join(cfg.TestDataDir, "dtail_format.log")}},
 		{"dmap", "aggregations", []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
-			"-query", "from STATS select sum($goroutines),avg($cgocalls),max(lifetimeConnections)", 
+			"-query", "from STATS select sum($goroutines),avg($cgocalls),max(lifetimeConnections)",
 			"-files", filepath.Join(cfg.TestDataDir, "dtail_format.log")}},
 		{"dmap", "csv_query", []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
 			"-query", `select count($line),count($user),count($action) group by $user,$action where $status eq "success" logformat csv`,
 			"-files", filepath.Join(cfg.TestDataDir, "test.csv")}},
 	}
-	
+
 	for _, p := range profiles {
 		common.PrintInfo("\nProfiling %s - %s\n", p.cmd, p.name)
 		for i := 1; i <= cfg.Runs; i++ {
@@ -216,19 +216,19 @@ func runFullProfile(cfg *Config) error {
 			}
 		}
 	}
-	
+
 	return analyzeLatestProfiles(cfg)
 }
 
 func runDMapProfile(cfg *Config) error {
 	common.PrintSection("DTail dmap Profiling")
-	
+
 	// Generate MapReduce test data
 	gen := common.NewDataGenerator()
-	
+
 	smallFile := filepath.Join(cfg.TestDataDir, "stats_small.log")
 	mediumFile := filepath.Join(cfg.TestDataDir, "stats_medium.log")
-	
+
 	common.PrintInfo("Preparing MapReduce test data...\n")
 	if err := gen.GenerateLogFileWithLines(smallFile, 1000, common.FormatDTail); err != nil {
 		return fmt.Errorf("failed to generate small file: %w", err)
@@ -236,16 +236,16 @@ func runDMapProfile(cfg *Config) error {
 	if err := gen.GenerateLogFileWithLines(mediumFile, 1000000, common.FormatDTail); err != nil {
 		return fmt.Errorf("failed to generate medium file: %w", err)
 	}
-	
+
 	// Build dmap
 	common.PrintInfo("Building dmap...\n")
 	if err := common.BuildCommand("dmap"); err != nil {
 		return err
 	}
-	
+
 	// Profile different queries
 	common.PrintSection("Profiling dmap queries...")
-	
+
 	queries := []struct {
 		name  string
 		query string
@@ -256,7 +256,7 @@ func runDMapProfile(cfg *Config) error {
 		{"Min and max", "from STATS select min(currentConnections),max(lifetimeConnections) group by hostname", smallFile},
 		{"Large file processing", "from STATS select count($line),avg($goroutines) group by hostname", mediumFile},
 	}
-	
+
 	for _, q := range queries {
 		common.PrintInfo("\nQuery: %s\n", q.name)
 		args := []string{"-profile", "-profiledir", cfg.ProfileDir, "-plain", "-cfg", "none",
@@ -265,26 +265,26 @@ func runDMapProfile(cfg *Config) error {
 			return fmt.Errorf("failed to profile query %s: %w", q.name, err)
 		}
 	}
-	
+
 	return analyzeLatestProfiles(cfg)
 }
 
 func profileCommand(name, cmd string, args []string, timeout time.Duration) error {
 	fmt.Printf("Command: %s %s\n", cmd, strings.Join(args, " "))
-	
+
 	command := exec.Command("./"+cmd, args...)
 	command.Stdout = nil // Suppress output during profiling
 	command.Stderr = os.Stderr
-	
+
 	if err := command.Start(); err != nil {
 		return err
 	}
-	
+
 	done := make(chan error, 1)
 	go func() {
 		done <- command.Wait()
 	}()
-	
+
 	select {
 	case <-time.After(timeout):
 		command.Process.Kill()
@@ -294,9 +294,9 @@ func profileCommand(name, cmd string, args []string, timeout time.Duration) erro
 			return err
 		}
 	}
-	
+
 	// Find generated profile
-	pattern := filepath.Join("profiles", fmt.Sprintf("%s_cpu_*.prof", name))
+	pattern := filepath.Join(profileDirFromArgs(args), fmt.Sprintf("%s_cpu_*.prof", name))
 	matches, _ := filepath.Glob(pattern)
 	if len(matches) > 0 {
 		// Sort by modification time and get the latest
@@ -307,52 +307,61 @@ func profileCommand(name, cmd string, args []string, timeout time.Duration) erro
 		})
 		fmt.Printf("  Generated: %s\n", filepath.Base(matches[0]))
 	}
-	
+
 	return nil
 }
 
+func profileDirFromArgs(args []string) string {
+	for i := 0; i < len(args)-1; i++ {
+		if args[i] == "-profiledir" {
+			return args[i+1]
+		}
+	}
+	return "profiles"
+}
+
 func analyzeLatestProfiles(cfg *Config) error {
 	common.PrintSection("Profile Analysis")
-	
+
 	// Find latest profiles for each command
 	for _, cmd := range cfg.Commands {
 		cpuPattern := filepath.Join(cfg.ProfileDir, fmt.Sprintf("%s_cpu_*.prof", cmd))
 		memPattern := filepath.Join(cfg.ProfileDir, fmt.Sprintf("%s_mem_*.prof", cmd))
-		
+
 		cpuProfiles, _ := filepath.Glob(cpuPattern)
 		memProfiles, _ := filepath.Glob(memPattern)
-		
+
 		if len(cpuProfiles) > 0 {
 			sort.Slice(cpuProfiles, func(i, j int) bool {
 				fi, _ := os.Stat(cpuProfiles[i])
 				fj, _ := os.Stat(cpuProfiles[j])
 				return fi.ModTime().After(fj.ModTime())
 			})
-			
+
 			fmt.Printf("\n%s CPU Profile: %s\n", cmd, filepath.Base(cpuProfiles[0]))
 			if err := showTopFunctions(cpuProfiles[0], 5, false); err != nil {
 				fmt.Printf("  Analysis failed: %v\n", err)
 			}
 		}
-		
+
 		if len(memProfiles) > 0 {
 			sort.Slice(memProfiles, func(i, j int) bool {
 				fi, _ := os.Stat(memProfiles[i])
 				fj, _ := os.Stat(memProfiles[j])
 				return fi.ModTime().After(fj.ModTime())
 			})
-			
+
 			fmt.Printf("\n%s Memory Profile: %s\n", cmd, filepath.Base(memProfiles[0]))
 			if err := showTopFunctions(memProfiles[0], 5, true); err != nil {
 				fmt.Printf("  Analysis failed: %v\n", err)
 			}
 		}
 	}
-	
+
 	common.PrintSuccess("\nProfiling complete!\n")
 	fmt.Println("\nTo analyze profiles in detail:")
 	fmt.Printf("  go tool pprof %s/<profile_file>\n", cfg.ProfileDir)
 	fmt.Printf("  dtail-tools profile -mode analyze <profile_file>\n")
-	
+
 	return nil
-}
\ No newline at end of file
+}
diff --git a/internal/tools/profile/profile_test.go b/internal/tools/profile/profile_test.go
new file mode 100644
index 0000000..1a11fdd
--- /dev/null
+++ b/internal/tools/profile/profile_test.go
@@ -0,0 +1,30 @@
+package profile
+
+import "testing"
+
+func TestProfileDirFromArgs(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+		want string
+	}{
+		{
+			name: "explicit profile dir",
+			args: []string{"-profile", "-profiledir", "custom-profiles", "-plain"},
+			want: "custom-profiles",
+		},
+		{
+			name: "missing profile dir falls back to default",
+			args: []string{"-profile", "-plain"},
+			want: "profiles",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := profileDirFromArgs(tt.args); got != tt.want {
+				t.Fatalf("profileDirFromArgs(%v) = %q, want %q", tt.args, got, tt.want)
+			}
+		})
+	}
+}