@@ -18,6 +18,8 @@ import (
1818 "fmt"
1919 "log/slog"
2020 "os"
21+ "runtime/pprof"
22+ "strings"
2123
2224 "github.com/blinklabs-io/dingo/internal/config"
2325 "github.com/blinklabs-io/dingo/internal/version"
@@ -37,7 +39,9 @@ func slogPrintf(format string, v ...any) {
3739
3840var (
3941 globalFlags = struct {
40- debug bool
42+ cpuprofile string
43+ memprofile string
44+ debug bool
4145 }{}
4246 configFile string
4347)
@@ -72,6 +76,50 @@ func commonRun() *slog.Logger {
7276}
7377
7478func main () {
79+ // Parse profiling flags before cobra setup (handle both --flag=value and --flag value syntax)
80+ cpuprofile := ""
81+ memprofile := ""
82+ args := os .Args
83+ if len (args ) > 0 {
84+ args = args [1 :] // Skip program name
85+ } else {
86+ args = []string {}
87+ }
88+ for i := 0 ; i < len (args ); i ++ {
89+ arg := args [i ]
90+ switch {
91+ case strings .HasPrefix (arg , "--cpuprofile=" ):
92+ cpuprofile = strings .TrimPrefix (arg , "--cpuprofile=" )
93+ case arg == "--cpuprofile" && i + 1 < len (args ):
94+ cpuprofile = args [i + 1 ]
95+ i ++ // Skip next arg
96+ case strings .HasPrefix (arg , "--memprofile=" ):
97+ memprofile = strings .TrimPrefix (arg , "--memprofile=" )
98+ case arg == "--memprofile" && i + 1 < len (args ):
99+ memprofile = args [i + 1 ]
100+ i ++ // Skip next arg
101+ }
102+ }
103+
104+ // Initialize CPU profiling (starts immediately, stops on exit)
105+ if cpuprofile != "" {
106+ fmt .Fprintf (os .Stderr , "Starting CPU profiling to %s\n " , cpuprofile )
107+ f , err := os .Create (cpuprofile )
108+ if err != nil {
109+ fmt .Fprintf (os .Stderr , "could not create CPU profile: %v\n " , err )
110+ os .Exit (1 )
111+ }
112+ defer f .Close ()
113+ if err := pprof .StartCPUProfile (f ); err != nil {
114+ fmt .Fprintf (os .Stderr , "could not start CPU profile: %v\n " , err )
115+ os .Exit (1 )
116+ }
117+ defer func () {
118+ pprof .StopCPUProfile ()
119+ fmt .Fprintf (os .Stderr , "CPU profiling stopped\n " )
120+ }()
121+ }
122+
75123 rootCmd := & cobra.Command {
76124 Use : programName ,
77125 Run : func (cmd * cobra.Command , args []string ) {
@@ -89,6 +137,10 @@ func main() {
89137 BoolVarP (& globalFlags .debug , "debug" , "D" , false , "enable debug logging" )
90138 rootCmd .PersistentFlags ().
91139 StringVar (& configFile , "config" , "" , "path to config file" )
140+ rootCmd .PersistentFlags ().
141+ StringVar (& globalFlags .cpuprofile , "cpuprofile" , "" , "write cpu profile to file" )
142+ rootCmd .PersistentFlags ().
143+ StringVar (& globalFlags .memprofile , "memprofile" , "" , "write memory profile to file" )
92144
93145 rootCmd .PersistentPreRunE = func (cmd * cobra.Command , args []string ) error {
94146 cfg , err := config .LoadConfig (configFile )
@@ -109,4 +161,19 @@ func main() {
109161 // NOTE: we purposely don't display the error, since cobra will have already displayed it
110162 os .Exit (1 )
111163 }
164+
165+ // Finalize memory profiling (captures heap state at program end)
166+ if memprofile != "" {
167+ f , err := os .Create (memprofile )
168+ if err != nil {
169+ fmt .Fprintf (os .Stderr , "could not create memory profile: %v\n " , err )
170+ os .Exit (1 )
171+ }
172+ defer f .Close ()
173+ if err := pprof .WriteHeapProfile (f ); err != nil {
174+ fmt .Fprintf (os .Stderr , "could not write memory profile: %v\n " , err )
175+ os .Exit (1 )
176+ }
177+ fmt .Fprintf (os .Stderr , "Memory profiling complete\n " )
178+ }
112179}
0 commit comments