@@ -22,6 +22,7 @@ import (
22
22
"fmt"
23
23
"os"
24
24
"runtime"
25
+ "runtime/pprof"
25
26
"slices"
26
27
"sort"
27
28
"strings"
@@ -270,12 +271,34 @@ func main() {
270
271
Usage : "Address to which the gRPC tracing collector will send spans to." ,
271
272
Value : "127.0.0.1:4317" ,
272
273
},
274
+ & cli.StringFlag {
275
+ Name : "profile-cpu" ,
276
+ Usage : "Write a pprof CPU profile to the provided path." ,
277
+ },
278
+ & cli.StringFlag {
279
+ Name : "profile-mem" ,
280
+ Usage : "Write a pprof memory profile to the provided path." ,
281
+ },
273
282
}
274
283
275
284
app .Before = func (context * cli.Context ) (err error ) {
276
285
var config * common.ServerConfiguration
277
286
var exePath string
278
287
288
+ cpuProfilePath := context .String ("profile-cpu" )
289
+ if cpuProfilePath != "" {
290
+ logrus .Infof ("Creating CPU profile in: %s" , cpuProfilePath )
291
+
292
+ file , err := os .Create (cpuProfilePath )
293
+ if err != nil {
294
+ return fmt .Errorf ("could not create CPU profile: %w" , err )
295
+ }
296
+
297
+ if err := pprof .StartCPUProfile (file ); err != nil {
298
+ return fmt .Errorf ("could not start CPU profiling: %w" , err )
299
+ }
300
+ }
301
+
279
302
if exePath , err = os .Executable (); err != nil {
280
303
logrus .Fatal (err )
281
304
}
@@ -355,6 +378,29 @@ func main() {
355
378
356
379
return nil
357
380
}
381
+
382
+ app .After = func (ctx * cli.Context ) error {
383
+ memProfilePath := ctx .String ("profile-mem" )
384
+ if memProfilePath != "" {
385
+ logrus .Infof ("Creating memory profile in: %s" , memProfilePath )
386
+
387
+ file , err := os .Create (memProfilePath )
388
+ if err != nil {
389
+ return fmt .Errorf ("could not create memory profile: %w" , err )
390
+ }
391
+ defer file .Close ()
392
+
393
+ // Ensure up to date data
394
+ runtime .GC ()
395
+
396
+ if err := pprof .WriteHeapProfile (file ); err != nil {
397
+ return fmt .Errorf ("could not write memory profile: %w" , err )
398
+ }
399
+ }
400
+
401
+ return nil
402
+ }
403
+
358
404
// sort all flags
359
405
for _ , cmd := range app .Commands {
360
406
sort .Sort (cli .FlagsByName (cmd .Flags ))
@@ -373,4 +419,7 @@ func main() {
373
419
logrus .Errorf ("Unable to shutdown tracer provider: %v" , err )
374
420
}
375
421
}
422
+
423
+ // noop when profiling is not enabled
424
+ pprof .StopCPUProfile ()
376
425
}
0 commit comments